Behavior Cloning to Policy Gradients

Behavior cloning to score-based policy gradient

This note makes the explicit connection between supervised learning (i.e., behavior cloning) and reinforcement learning (i.e., score-based policy gradients.

We start with defining behavior cloning, which assumes access to expert demonstrations, and calculate the gradient for maximizing its likelihood. Then, we introduce a dataset with bad demonstrations, which we try to (additionally) minimize likelihood on. Then, we introduce reward/advantage weighted regression, which replaces the explicit good/bad labels with returns / advantages. This then ends up looking like the standard score-based policy gradient.

Supervised learning - Behavior cloning with maximum log likelihood

We assume there is:

1. An expert policy ${\pi }_{+}(a|s)$ that defines a distribution over actions conditioned on states,
2. A dataset of n state-action pairs $D_{+}=\{(s_i,a_i){\}}_{i=1}^n$ that is generated by some (arbitrary) state distribution $d_{+}(s)$ and querying the expert for actions in those states, so that state-action samples in the dataset $D_{+}$ are jointly distributed according to $P{r}_{+}(s,a)=d_{+}(s){\pi }_{+}(a|s)$.

The maximum log-likelihood objective is to find a policy ${\pi }_{\theta }$ that optimizes expected log likelihood of state-action pairs from the state-action distribution $P{r}_{+}(s,a)$:

$$\begin{array}{rl}L({\pi }_{\theta })& ={\mathbb{E}}_{(s,a)\sim P{r}_{+}(\ast )}[\log ({\pi }_{\theta }(a|s))]\\ & \approx \frac{1}{n}\sum _{i=1}^n\log ({\pi }_{\theta }(a_i|s_i))\\ {\nabla }_{\theta }L({\pi }_{\theta })& \approx \frac{1}{n}\sum _{i=1}^n{\nabla }_{\theta }\log ({\pi }_{\theta }(a_i|s_i))\end{array}$$

• write code for this

Ultimately, we are finding a policy that maximizes the likelihood of actions that the expert generates, weighted according to an arbitrary state distribution $d_{+}(s)$.

Positive and negative examples, and a single dataset to include them all

Let’s assume we also have a dataset of $m$ “bad” state-action pairs $D_{-}=\{(s_j,a_j){\}}_{j=1}^m$, where the state-action pairs are sampled from a joint distribution that has (a potentially different) starting state distribution $d_{-}(s)$ and actions generated by a “bad” policy ${\pi }_{-}$, resulting a joint distribution that factors in a similar way to the good dataset: $P{r}_{-}(s,a)=d_{-}(s)\ast {\pi }_{-}(a|s)$.

We can now define an objective where in addition to maximizing likelihood of expert actions, we additional minimize the likelihood of bad examples:

$$\begin{array}{rl}L({\pi }_{\theta })& ={\mathbb{E}}_{(s,a)\sim P{r}_{+}(\ast )}[\log ({\pi }_{\theta }(a\mid s))]-{\mathbb{E}}_{(s,a)\sim P{r}_{-}(\ast )}[\log ({\pi }_{\theta }(a\mid s))]\\ & \approx \frac{1}{n}\sum _{i=1}^n\log ({\pi }_{\theta }(a_i|s_i))-\frac{1}{m}\sum _{j=1}^m\log ({\pi }_{\theta }(a_j|s_j))\\ {\nabla }_{\theta }L({\pi }_{\theta })& \approx \frac{1}{n}\sum _{i=1}^n{\nabla }_{\theta }\log ({\pi }_{\theta }(a_i|s_i))-\frac{1}{m}\sum _{j=1}^m{\nabla }_{\theta }\log ({\pi }_{\theta }(a_j|s_j))\end{array}$$

Since the only difference between maximizing and minimizing is whether the we’re multiplying the gradient of the loglikelihood by either positive/negative one, we can rewrite this objective over a single dataset $D$ that is a combination of the positive and negative datasets $D_{+}$ and $D_{-}$ (therefore containing $q=n+m$ state-action pairs) and add a label (-1/1) to each transition to denote whether it’s from the good/bad dataset:

$$\begin{array}{rl}D& =\{(s_k,a_k,{\Psi }_k){\}}_{k=1}^q\\ {\Psi }_k& =\{\begin{array}{llc}1,& & \text{(}k\le n,\text{good state-action pair)}\\ -1,& & \text{(}k>n,\text{bad state-action pair)}\end{array}\end{array}$$

Since this dataset $D$ is just a union of the good and bad datasets $D_{+}$ and $D_{-}$ (with an additional label for which dataset it came from) and has size $q=n+m$, we can calculate the probability of a state-action-label datapoint in the new dataset $D$ by taking into account what proportion of the dataset is good vs. bad data:

$$\begin{array}{rl}Pr(s_k,a_k,{\Psi }_k)& =Pr({\Psi }_k)d(s_k\mid {\Psi }_k)\pi (a_k\mid s_k,{\Psi }_k)\\ Pr({\Psi }_k)& =\{\begin{array}{llc}\frac{n}{n+m}& & \text{(fraction of data that is good)}\\ \frac{m}{n+m}& & \text{(fraction of data that is bad)}\end{array}\\ d(s_k|{\Psi }_k)& =\{\begin{array}{llc}{d}_{+}(s_k),& & k\le n\\ d_{-}(s_k),& & k>n\end{array}\\ \pi (a_k\mid s_k,{\Psi }_k)& =\{\begin{array}{llc}{\pi }_{+}(a_k\mid s_k),& & k\le n\\ {\pi }_{-}(a_k\mid s_k),& & k>n\end{array}\end{array}$$

Now we can write an objective using the single dataset, and using the label $\Psi$ to handle the sign by multiplying it by the log-likelihood so we can write the loss as one expectation/sum instead of two:

$$\begin{array}{rl}L({\pi }_{\theta })& ={\mathbb{E}}_{(s,a,\Psi )\sim P{r}_{(}\ast )}[\Psi \log ({\pi }_{\theta }(a\mid s))]\\ & \approx \frac{1}{q}\sum _{k=1}^q{\Psi }_k\log ({\pi }_{\theta }(a_k|s_k))\\ {\nabla }_{\theta }L({\pi }_{\theta })& \approx \frac{1}{q}\sum _{k=1}^q{\Psi }_k{\nabla }_{\theta }\log ({\pi }_{\theta }(a_k|s_k))\end{array}$$

• prove that this is very similar to the loss above that has the 2 separate datasets, but with a reweighting on the data.

Reward-weighted regression (RWR)

References

• original paper (Relative Entropy Policy Search)
• advantage-weighted regression paper: more interpretable / up-to-date imo

At each iteration, reward-weighted regression (RWR) solves the following regression problem:

$$\begin{array}{r}{\pi }_{k+1}=\arg \underset{\pi }{\max }{\mathbb{E}}_{s\sim d_{{\pi }_k}(s)}{\mathbb{E}}_{a\sim {\pi }_k(a|s)}[\exp (\frac{1}{\beta }{R}_{s,a})\log \pi (a|s)]\end{array}$$

Where:

• ${\pi }_k$ represents the policy at the kth iteration of the algorithm
• $R_{s,a}={\sum }_{t=0}^{\infty }{\gamma }^t{r}_t$ is the return (and can be interpreted as an empirical sample from $Q^{{\pi }_k}(s,a)$)
• $d_{{\pi }_k}(s)$ is the unnormalized discounted state distribution induced by the policy ${\pi }_k$

“The RWR update can be interpreted as solving a maximum likelihood problem that fits a new policy ${\pi }_{k+1}$ to samples collected under the current policy ${\pi }_k$, where the likelihood of each action is weighted by the exponentiated return received for that action, with a temperature paramtere $\beta >0$”

To clarify how similar this is to the Positive and negative examples, and a single dataset to include them all loss and policy gradient, we will

1. Set $R_{s,a}=\Psi (s,a)$, so that the returns are being used like a “label” to give more weight to maximizing the likelihood of actions based on their value.
2. write the optimization problem above as a loss $L_k({\pi }_{\theta })$ at each iteration $k$ and give the gradient w.r.t policy parameters:

$$\begin{array}{rl}{L}_k({\pi }_{\theta })& ={\mathbb{E}}_{s\sim d_{{\pi }_k}(s),a\sim {\pi }_k(a|s)}[\Psi (s,a)\log {\pi }_{\theta }(a|s)]\\ & \approx \frac{1}{n}\sum _{i=1}^n\Psi (s_i,a_i)\log {\pi }_{\theta }(a_i|s_i)\\ {\nabla }_{\theta }{L}_k({\pi }_{\theta })& \approx \frac{1}{n}\sum _{i=1}^n\Psi (s_i,a_i){\nabla }_{\theta }\log {\pi }_{\theta }(a_i|s_i)\end{array}$$

The biggest differences between the previous objective and this one are:

• The distribution of the expectation (RWR is on-policy, and the previous approach used an arbitrary state-distribution and assume actions came from the expert policy). Every time we update the policy, RWR collects a new dataset, whereas we used the same dataset each time we updated the policy in the supervised learning setting.

• In our newest objective, the label $\Psi$ is state-action conditioned and is derived from interaction samples with the environment produced by the policy, whereas before it was just a provided label from an expert. Also, the original label was just $-1$ and $1$, whereas the new label $\Psi (s,a)$ can be an arbitrary scalar.

• todo: provide justifcation from original paper for this objective (maximizing returns while staying close to original policy, KL constrained policy improvement)

• mention that this is on-policy, but we can do this is in an off-policy fashion by incorporating importance sampling (mentioned in AWR paper, Section 3.2).

Advantage-weighted regression (AWR)

Advantage-weighted regression extends the RWR objective above in two key ways:

1. Rather than operate on on-policy data, it instead learns from off-policy data (different state-action distribution for the expectation). More specifically, it starts with an empty replay buffer $D$ and adds data from each iteration of the policy ${\pi }_k$ (Let $D_k$ represent the replay buffer at iteration $k$).
2. It replaces the returns $R_{s,a}$ with an estimate of the advantage $A_k^D=R_{s,a}-V_k^D$, where $R_{s,a}$ is a trajectory from the dataset containing the mixture of rollouts from past policies, and $V_k^D$ is a value function that is fitted with empirical returns from $D_k$.

$$\begin{array}{rl}{\pi }_{k+1}& =\arg \underset{\pi }{\max }{\mathbb{E}}_{s,a\sim D_k}[\log \pi (a|s)\exp (\frac{1}{\beta }(R_{s,a}-V_k^D(s)))]\\ & =\arg \underset{\pi }{\max }{\mathbb{E}}_{s,a\sim D_k}[\log \pi (a|s)\exp (\frac{1}{\beta }{A}_k^D(s,a))]\end{array}$$

• todo: give details on how this objective is motivated from policy improvement objective instead of policy maximization.

To make the relationship to Reward-weighted regression (RWR) we can set $A_k^{D}s,a=\Psi (s,a)$ and define the loss/gradient for each iteration:

$$\begin{array}{rl}{L}_k({\pi }_{\theta })& ={\mathbb{E}}_{s,a\sim D_k}[\Psi (s,a)\log {\pi }_{\theta }(a|s)]\\ & \approx \frac{1}{n}\sum _{i=1}^n\Psi (s_i,a_i)\log {\pi }_{\theta }(a_i|s_i)\\ {\nabla }_{\theta }{L}_k({\pi }_{\theta })& \approx \frac{1}{n}\sum _{i=1}^n\Psi (s_i,a_i){\nabla }_{\theta }\log {\pi }_{\theta }(a_i|s_i)\end{array}$$

References

• original DDPG paper (connects DPPG to stochastic policy gradient)
• REPPO
  • Establishes pathwise policy gradient
• Proof that the optima policy to extract from a value function with entropy penalty is soft max.
• • From policy gradient to actor-critic methods
    • pretty sparse, does go over BC (regression and max-likelihood) to RWR and policy gradients. But isn’t formal about the distributions and how to make them look exactly the same / not much math notation.
• Offline Reinforcement Learning: Tutorial, Review, and Perspectives on Open Problems
  • good information on AWR (ctrl+f AWR for info)

TODOS:

• provide code for each loss and gradient.
• provide the empirical estimates because they match our code implementation.