\documentclass{article} \usepackage{amsmath} \usepackage{algorithm} \usepackage{algpseudocode} \usepackage{fancyhdr} \usepackage{lastpage} \usepackage{extramarks} \usepackage[usenames,dvipsnames]{color} \usepackage{graphicx} % Required to insert images \usepackage{listings} % Required for insertion of code \usepackage{courier} % Required for the courier font \usepackage{enumerate} \usepackage{enumitem} \usepackage{booktabs} \usepackage{amsmath, amsthm, amssymb} \usepackage{caption} \usepackage{subcaption} \captionsetup[table]{skip=4pt} \usepackage{framed} \usepackage{bm} \usepackage[most]{tcolorbox} \usepackage{physics} \usepackage{xcolor} \graphicspath{{img/}} % set of paths to search for images \usepackage{amsmath} \usepackage{amsfonts} \usepackage{amssymb} \usepackage{hyperref} \usepackage{float} \usepackage{xcolor} \hypersetup{ colorlinks = true, linkcolor = blue, anchorcolor = blue, citecolor = red, filecolor = blue, urlcolor = blue } \newenvironment{myitemize} { \begin{itemize} \setlength{\itemsep}{0pt} \setlength{\parskip}{0pt} \setlength{\parsep}{0pt} } { \end{itemize} } \usepackage{biblatex} % bibliography \addbibresource{papers.bib} \usepackage{tikz} \usetikzlibrary{positioning,patterns,fit} \newcommand{\ifans}[1]{\ifanswers \color{red} \textbf{Solution: } #1 \color{black}} \newcommand{\E}{\mathbb{E}} \newcommand{\given}{\,|\,} \newcommand{\bE}{\mathbb{E}} \newcommand{\bH}{\mathbb{H}} \newcommand{\bI}{\mathbb{I}} \newcommand{\bN}{\mathbb{N}} \newcommand{\bR}{\mathbb{R}} \newcommand{\mc}[1]{\mathcal{#1}} \newcommand{\ra}{\rightarrow} \DeclareMathOperator*{\argmin}{arg\,min} \DeclareMathOperator*{\argmax}{arg\,max} \newcommand{\safereward}{r_{\text{safe}}} \newcommand{\lowreward}{\underline{r}_{\text{risk}}} \newcommand{\highreward}{\overline{r}_{\text{risk}}} \newcommand{\consreward}{r_{\text{cons}}} % Margins \topmargin=-0.45in \evensidemargin=0in \oddsidemargin=0in \textwidth=6.5in \textheight=9.0in \headsep=0.25in \linespread{1.1} % Line spacing % Set up the header and footer \pagestyle{fancy} \rhead{\hmwkAuthorName} % Top left header \lhead{\hmwkClass: \hmwkTitle} % Top center head \lfoot{\lastxmark} % Bottom left footer \cfoot{} % Bottom center footer \rfoot{Page\ \thepage\ of\ \protect\pageref{LastPage}} % Bottom right footer \renewcommand\headrulewidth{0.4pt} % Size of the header rule \renewcommand\footrulewidth{0.4pt} % Size of the footer rule \setlength\parindent{0pt} % Removes all indentation from paragraphs %---------------------------------------------------------------------------------------- % CODE INCLUSION CONFIGURATION %---------------------------------------------------------------------------------------- \definecolor{MyDarkGreen}{rgb}{0.0,0.4,0.0} % This is the color used for comments \lstloadlanguages{Python} \lstset{language=Python, frame=single, % Single frame around code basicstyle=\footnotesize\ttfamily, % Use small true type font keywordstyle=[1]\color{Blue}\bf, keywordstyle=[2]\color{Purple}, keywordstyle=[3]\color{Blue}\underbar, % Custom functions underlined and blue identifierstyle=, % Nothing special about identifiers commentstyle=\usefont{T1}{pcr}{m}{sl}\color{MyDarkGreen}\small, % Comments small dark green courier font stringstyle=\color{Purple}, % Strings are purple showstringspaces=false, % Don't put marks in string spaces tabsize=5, % 5 spaces per tab morekeywords={rand}, morekeywords=[2]{on, off, interp}, morekeywords=[3]{test}, morecomment=[l][\color{Blue}]{...}, % Line continuation (...) like blue comment numbers=left, % Line numbers on left firstnumber=1, % Line numbers start with line 1 numberstyle=\tiny\color{Blue}, % Line numbers are blue and small stepnumber=5 % Line numbers go in steps of 5 } \newcommand{\perlscript}[2]{ \begin{itemize} \item[]\lstinputlisting[caption=#2,label=#1]{#1.pl} \end{itemize} } %---------------------------------------------------------------------------------------- % NAME AND CLASS SECTION %---------------------------------------------------------------------------------------- \newcommand{\hmwkTitle}{Assignment \#3} % Assignment title \newcommand{\hmwkClass}{CS\ 234} % Course/class \newcommand{\hmwkAuthorName}{} % Your name %---------------------------------------------------------------------------------------- % TITLE PAGE %---------------------------------------------------------------------------------------- \title{ \vspace{-1in} \textmd{\textbf{\hmwkClass:\ \hmwkTitle}}} \author{} \date{} % Insert date here if you want it to appear below your name \begin{document} \maketitle \vspace{-.5in} \begin{framed} {\bf Due date: May 17, 2024 at 6:00 PM (18:00) PST} \\[1em] These questions require thought but do not require long answers. Please be as concise as possible. \\[1em] We encourage students to discuss in groups for assignments. \textbf{However, each student must finish the problem set and programming assignment individually, and must turn in her/his assignment.} We ask that you abide by the university Honor Code and that of the Computer Science department, and make sure that all of your submitted work is done by yourself. If you have discussed the problems with others, please include a statement saying who you discussed problems with. Failure to follow these instructions will be reported to the Office of Community Standards. We reserve the right to run a fraud-detection software on your code. \\[1em] Please review any additional instructions posted on the assignment page at http://web.stanford.edu/class/cs234/assignments.html. When you are ready to submit, please follow the instructions on the course website. \\[1em] \end{framed} \section*{An introduction to reinforcement learning from human preferences} The goal of this assignment is to give you some hands-on experience with reinforcement learning from human preferences (RLHF). You will implement, compare, and contrast several different approaches for a robotics task in MuJoCo. \section{Reward engineering (6 pts writeup + 7 pts coding)} In Assignment 2 you applied PPO to solve an environment with a provided reward function. The process of deriving a reward function for a specific task is called reward engineering. \subsection{Written Questions (6 pts)} \begin{enumerate} \item[(a) (2 pts)] (\textbf{written}) Why is reward engineering usually hard? What are potential risks that come with specifying an incorrect reward function? Provide an example of a problem and a reward function that appears to be adequate but may have unintended consequences. \item[(b) (2 pts)] (\textbf{written}) Read the description of the \href{https://gymnasium.farama.org/environments/mujoco/hopper/}{Hopper environment}. Using your own words, describe the goal of the environment, and how each term of the reward functions contributes to encourage the agent to achieve it. (\textbf{Optional:} Do you agree with this reward function? Would you change it? How?) \item[(c) (2 pts)] (\textbf{written}) By default, the episode terminates when the agent leaves the set of ``healthy'' states. What do these ``healthy'' states mean? Name one advantage and one disadvantage of this early termination. \end{enumerate} \subsection{Coding Questions (7 pts)} \begin{enumerate} \item[(d) (2 pts)] (\textbf{coding}) Use the provided starter code to train a policy using PPO to solve the Hopper environment for 3 different seeds. Do this with and without early termination. \begin{tcolorbox} \begin{verbatim} python ppo_hopper.py [--early-termination] --seed SEED \end{verbatim} \end{tcolorbox} \item[(e) (2 pts)] (\textbf{written}) Attach here the plot of the episodic returns along training, with and without early termination. You can generate the plot by running \begin{tcolorbox} \begin{verbatim} python plot.py --directory results --seeds SEEDS \end{verbatim} \end{tcolorbox} where \texttt{SEEDS} is a comma-separated list of the seeds you used. Comment on the performance in terms of training epochs and wall time. Is the standard error in the average returns high or low? How could you obtain a better estimate of the average return on Hopper achieved by a policy optimized with PPO? \item[(f) (2 pts)] (\textbf{written}) Pick one of the trained policies and render a video of an evaluation rollout. \begin{tcolorbox} \begin{verbatim} python render.py --checkpoint [PATH TO MODEL CHECKPOINT] \end{verbatim} \end{tcolorbox} Does the agent successfully complete the assigned task? Does it complete it in the way you would expect it to, or are you surprised by the agent behavior? \item[(g) (1 pts)] (\textbf{written}) Render another video for another policy. How do the two rollouts compare? Do you prefer one over the other? \end{enumerate} \section{Learning from preferences (5 pts writeup + 25 pts coding)} In the previous part you trained multiple policies from scratch and compared them at the end of training. In this section, we will see how we can use human preferences on two roll-outs to learn a reward function. We will follow the framework proposed by \cite{NIPS2017_d5e2c0ad}. A reward function $r: \mathcal{O} \times \mathcal{A} \rightarrow \mathbb{R}$ defines a preference relation $\succ$ if for all trajectories $\sigma^i = (o^i_t,a^i_t)_{t=0,...,T}$ we have that $$ \left(\left(o_0^1, a_0^1\right), \ldots,\left(o_{T}^1, a_{T}^1\right)\right) \succ\left(\left(o_0^2, a_0^2\right), \ldots,\left(o_{T}^2, a_{T}^2\right)\right) $$ whenever $$ r\left(o_0^1, a_0^1\right)+\cdots+r\left(o_{T}^1, a_{T}^1\right)>r\left(o_0^2, a_0^2\right)+\cdots+r\left(o_{T}^2, a_{T}^2\right) . $$ Following the Bradley-Terry preference model \cite{19ff28b9-64f9-3656-ba40-08326a05748e}, we can calculate the probability of one trajectory $\sigma^1$ being preferred over $\sigma^2$ as follows: $$ \hat{P}\left[\sigma^1 \succ \sigma^2\right]=\frac{\exp \sum \hat{r}\left(o_t^1, a_t^1\right)}{\exp \sum \hat{r}\left(o_t^1, a_t^1\right)+\exp \sum \hat{r}\left(o_t^2, a_t^2\right)}, $$ where $\hat{r}$ is an estimate of the reward for a state-action pair. This is similar to a classification problem, and we can fit a function approximator to $\hat{r}$ by minimizing the cross-entropy loss between the values predicted with the above formula and ground truth human preference labels $\mu(1)$ and $\mu(2)$: $$ \operatorname{loss}(\hat{r})=-\sum_{\left(\sigma^1, \sigma^2, \mu\right) \in \mathcal{D}} \mu(1) \log \hat{P}\left[\sigma^1 \succ \sigma^2\right]+\mu(2) \log \hat{P}\left[\sigma^2 \succ \sigma^1\right] . $$ % Once we have learned the reward function\footnote{Recent work on RLHF for reinforcement learning suggests that the pairwise feedback provided by humans on partial trajectories may be more consistent with regret, and that the learned reward function may be better viewed as an advantage function. See Knox et al. AAAI 2024 "Learning optimal advantage from preferences and mistaking it for reward." \url{https://openreview.net/forum?id=euZXhbTmQ7}}, we can apply any policy optimization algorithm (such as PPO) to maximize the returns of a model under it. \subsection{Written questions (5 pts)} \begin{enumerate} \item[(a) (5 pt)] Let $\hat{r}(o,\ a) = \phi_w(o,\ a)$. Calculate $\nabla_w \operatorname{loss}(\hat{r}(o,\ a))$. \end{enumerate} \subsection{Coding questions (25 pts)} In this problem we are trying to solve the same task as in the previous part, but this time we will learn a reward function from a dataset of preferences rather than manually specifying a reward function. \begin{enumerate} \item[(b) (5 pt)] Load one of the samples from the preference dataset we provide you, and render a video of the two trajectories using the following command \begin{tcolorbox} \begin{verbatim} python render.py --dataset [PATH TO PREFERENCE DATASET FILE] --idx IDX \end{verbatim} \end{tcolorbox} where \texttt{IDX} is an index into the preference dataset (if ommitted a sequence will be chosen at random). Bear in mind that each sequence in the dataset has 25 timesteps, which means that the resulting videos will have 0.2 seconds. Take note of which sequence was labeled as preferred (this information will appear in the name of the generated videos, but for the coming parts it is helpful to know that $0$ means the first sequence was preferred, $1$ means the second one, and $0.5$ means neither is preferred over the other). Do you agree with the label (that is, if shown the two trajectories, would you have ranked them the same way they appear in the dataset, knowing that we are trying to solve the Hopper environment)? \item[(c) (3 pt)] Repeat the previous question for 4 more samples, keeping track of whether you personally agree with the dataset preference label. Use this to estimate how much you agree with whoever ranked the trajectories. How much did you get? Based on this agreement estimate, would you trust a reward function learned on this data? \item[(d) (8 pt)] Implement the functions in the \texttt{RewardModel} class (\texttt{run\_rlhf.py}), which is responsible for learning a reward function from preference data. \item[(e) (5 pt)] Train a model using PPO and the learned reward function with 3 different random seeds. Plot the average returns for both the original reward function and the learned reward function and include it in your response. \begin{tcolorbox} \begin{verbatim} python plot.py --rlhf-directory results_rlhf --output \ results_rlhf/hopper_rlhf.png \end{verbatim} \end{tcolorbox} Do the two correlate? \item[(f) (1 pt)] Is the learned reward function identifiable? \item[(g) (3 pt)] Pick one of the policies and render a video of the agent behavior at the end of training. \begin{tcolorbox} \begin{verbatim} python render.py --checkpoint [PATH TO MODEL CHECKPOINT] \end{verbatim} \end{tcolorbox} How does it compare to the behavior of the agent generated by policies trained from scratch? How does it compare to the demonstrations you've seen from the dataset? \end{enumerate} \section{Direct preference optimization (25 pts coding)} In the previous question we saw how we could train a model based on preference data. However, suppose you are given a pre-trained model and the corresponding preference data. Ideally, you would like to optimize the model directly on the preference data, instead of having to learn a reward function, and then run PPO on it. That is the idea behind direct preference optimization (DPO) \cite{NEURIPS2023_a85b405e}. The algorithm proposed in the original paper allows us to skip the reward learning and reinforcement learning steps, and optimize the model directly from preference data by optimizing the following loss: $$\mathcal{L}_{\mathrm{DPO}}\left(\pi_\theta ; \pi_{\mathrm{ref}}\right)=-\mathbb{E}_{\left(x, y_w, y_l\right) \sim \mathcal{D}}\left[\log \sigma\left(\beta \log \frac{\pi_\theta\left(y_w \mid x\right)}{\pi_{\mathrm{ref}}\left(y_w \mid x\right)}-\beta \log \frac{\pi_\theta\left(y_l \mid x\right)}{\pi_{\mathrm{ref}}\left(y_l \mid x\right)}\right)\right],$$ where $\pi_{\mathrm{ref}}$ is the policy from which we sampled $y_w$ and $y_l$ given $x$, $\pi_\theta$ is the policy we are optimizing, and $\sigma$ is the sigmoid function. We will use DPO in this question as well. To provide some context, let us consider the general approach for RLHF for text generation: \begin{enumerate} \item Train a large language model (LLM) to do next token prediction given a context (the tokens that came previously). \item Given a fixed context $x$, generate possible next token sequence predictions $y_1$ and $y_2$, and store the triple $(x, y_1, y_2)$. \item Ask human supervisors to rank $y_1$ and $y_2$ given $x$ according to individual preference. \item Update the LLM to maximize the probability of giving the preferred answers using reinforcement learning. \end{enumerate} In a similar way, given an observation $x$ we could have two ranked sequences of actions $a^1_{1:T}$ and $a^2_{1:T}$, train the model to generate the preferred sequence of actions, and then execute them all\footnote{To understand why we are considering sequences of actions rather than a single action for the next time, recall that $25$ actions corresponded to $0.2$ seconds of video. If you found it difficult to rank a sequence of $25$ actions based on such a short video, imagine ranking the effect of a single action!}. If the length of the generated action sequence is equal to the environment time horizon, this is called open-loop control. However, this approach lacks robustness, since the plan of actions will not change in response to disturbances or compounding errors. Instead, we are going to adapt this into a more robust scheme by training our policy to predict a sequence of actions for the next $T$ time steps (where $T$ is the length of the trajectories in our preference dataset), but only take the first action in the plan generated by the policy. In this way, we re-plan our actions at every time step, ensuring the ability to respond to disturbances. \subsection{Coding questions (25 pts)} \begin{enumerate} \item[(a) (9 pts)] Implement the \texttt{ActionSequenceModel} class instance methods. When called, the model should return a probability distribution for the actions over the number of next time steps specified at initialization. Use a multivariate normal distribution for each action, with mean and standard deviation predicted by a neural network (see the starter code for more details).\footnote{We have prepared a \href{https://colab.research.google.com/drive/1sw8FJIR5865laTJiI0fqHG3KcGVT23ED}{notebook} to illustrate the behavior of \href{https://pytorch.org/docs/stable/distributions.html\#independent}{\texttt{torch.distributions.Independent}}.} \item[(b) (3 pts)] Implement the \texttt{update} method of the \texttt{SFT} class. This class will be used to pre-train a policy on the preference data by maximizing the log probabilities of the preferred actions given the observations in the dataset. \item[(c) (5 pts)] Implement the \texttt{update} method of the \texttt{DPO} class. This should minimize the DPO loss described above. \item[(d) (5 pts)] Run DPO for 3 different random seeds (you may want to tweak the number of DPO steps to get better results), and plot the evolution of returns over time. \begin{tcolorbox} \begin{verbatim} python plot.py --dpo-directory results_dpo --output \ results_dpo/hopper_dpo.png \end{verbatim} \end{tcolorbox} Include that plot in your response. How does it compare to the returns achieved using RLHF? Comment on the pros and cons of each method applied to this specific example. \item[(e) (3 pts)] Render a video of an episode generated by the pre-trained policy, and a video of an episode generated by the policy tuned by DPO. \begin{tcolorbox} \begin{verbatim} python render.py --dpo --checkpoint [PATH TO MODEL CHECKPOINT] \end{verbatim} \end{tcolorbox} How do they compare? \end{enumerate} \section{Best Arm Identification in Multi-armed Bandit (25pts)} In many experimental settings we are interested in quickly identifying the ``best" of a set of potential interventions, such as finding the best of a set of experimental drugs at treating cancer, or the website design that maximizes user subscriptions. Here we may be interested in efficient pure exploration, seeking to quickly identify the best arm for future use. In this problem, we bound how many samples may be needed to find the best or near-optimal intervention. We frame this as a multi-armed bandit with rewards bounded in $[0,1]$. Recall a bandit problem can be considered as a finite-horizon MDP with just one state ($|\mathcal{S}| = 1$) and horizon $1$: each episode consists of taking a single action and observing a reward. In the bandit setting -- unlike in standard RL -- the action (or ``arm") taken does not affect the distribution of future states. We assume a simple multi-armed bandit, meaning that $1 < |\mathcal{A}| < \infty$. Since there is only one state, a policy is simply a distribution over actions. There are exactly $|\mathcal{A}|$ different deterministic policies. Your goal is to design a simple algorithm to identify a near-optimal arm with high probability. We recall Hoeffding's inequality: if $X_1,\dots,X_n$ are i.i.d. random variables satisfying $0 \le X_i \le 1$ with probability $1$ for all $i$, $\overline X = \E[X_1] = \dots = \E[X_n]$ is the expected value of the random variables, and $\widehat X = \frac{1}{n} \sum_{i=1}^n X_i$ is the sample mean, then for any $\delta > 0$ we have \begin{align} \Pr\Bigg(|\widehat X - \overline X | > \sqrt{\frac{\log(2/\delta)}{2n}} \Bigg) < \delta. \end{align} Assuming that the rewards are bounded in $[0,1]$, we propose this simple strategy: pull each arm $n_e$ times, and return the action with the highest average payout $\widehat r_a$. The purpose of this exercise is to study the number of samples required to output an arm that is at least $\epsilon$-optimal with high probability. Intuitively, as $n_e$ increases the empirical average of the payout $\widehat r_a$ converges to its expected value $\overline r_a$ for every action $a$, and so choosing the arm with the highest empirical payout $\widehat r_a$ corresponds to approximately choosing the arm with the highest expected payout $\overline r_a$. \begin{enumerate} \item[(a) (10 pts)] We start by bounding the probability of the ``bad event'' in which the empirical mean of some arm differs significantly from its expected return. Starting from Hoeffding's inequality with $n_e$ samples allocated to every action, show that: \begin{align} \Pr\Bigg(\exists a \in \mathcal{A} \quad \text{s.t.} \quad |\widehat r_a - \overline r_a | > \sqrt{\frac{\log(2/\delta)}{2n_e}} \Bigg) < |\mathcal{A}|\delta. \end{align} Note that, depending on your derivation, you may come up with a tighter upper bound than $|\mathcal{A}|\delta$. This is also acceptable (as long as you argue that your bound is tighter), but showing the inequality above is sufficient. \item[(b) (15 pts)] After pulling each arm (action) $n_e$ times our algorithm returns the arm with the highest empirical mean: \begin{equation} a^\dagger = \arg\max_{a} \widehat r_a \end{equation} Notice that $a^\dagger$ is a random variable. Let ${a^\star} = \arg\max_a \overline r_{a}$ be the true optimal arm. Suppose that we want our algorithm to return at least an $\epsilon$-optimal arm with probability at least $1-\delta'$, as follows: \begin{equation} \label{eqn:maxa} \Pr \Bigg(\overline r_{a^\dagger} \geq \overline r_{a^\star} - \epsilon \Bigg) \geq 1-\delta'. \end{equation} How accurately do we need to estimate each arm in order to pick an arm that is $\epsilon$-optimal? Then derive how many total samples we need total (across all arms) to return an $\epsilon$-optimal arm with prob at least 1- $\delta'$ (that satisfies Equation~\ref{eqn:maxa}). Express your result as a function of the number of actions, the required precision $\epsilon$ and the failure probability $\delta'$. \item[(c) (0 pts)] (Optional challenge, will not be graded) The above derivation only assumed the outcomes were bounded between 0 and 1. In practice people often assume outcomes are drawn from a parametric distribution, and under mild assumptions, one can use the central limit theorem to assume the average outcomes for an arm will follow a normal distribution. Repeat the above analysis under this assumption, for a multi-armed bandit with two arms. Is the resulting number of samples significantly smaller under these assumptions? In real settings it is often very expensive to run experiments. Do you think the method and bound derived in (a-b) would be preferable to making a normal assumption and why or why not? \end{enumerate} \section{Challenges of RLHF (5 pts)} RLHF and DPO leverage humans providing pairwise comparisons. \begin{enumerate} \item[(a) (2 pts)] Who are human raters likely to be that are employed to provide pairwise preferences for generic LLM RLHF/DPO (in terms of socioeconomic profiles)? What ethical issues may this have, and what impact may this have on the resulting trained models? \item[(b) (1 pt)] Would RLHF be a good way to train a system to make better medical diagnoses? Why or why not? \item[(c) (2 pts)]There are a number of limitations of RLHF. Please describe one here (different than parts a and b). You may find it interesting to read \href{https://arxiv.org/pdf/2307.15217}{``Open Problems and Fundamental Limitations of Reinforcement Learning from Human Feedback''} and you are welcome to describe challenges listed there, or make up your own. \end{enumerate} \end{document}