A P ROBABILISTIC P ERSPECTIVE ON R EINFORCEMENT L EARNING VIA S UPERVISED L EARNING
A BSTRACT
Reinforcement Learning via Supervised Learning (RvS) only uses supervised techniques to learn desirable behaviors from large datasets.
RvS has attracted much attention lately due to its simplicity and ability to leverage diverse trajectories.
We introduce Density to Decision (D2D), a new framework, to unify a myriad of RvS algorithms.
The Density to Decision framework formulates RvS as a two-step process: i) density estimation via supervised learning and ii) decision making via exponential tilting of the density
Using our framework , we categorise popular RvS algorithms and show how they are different by the design choices in their implementation
We then introduce a novel algorithm Implicit
RvS leveraging powerful density estimation techniques that can easily be tilted to produce desirable behaviors
We compare the performance of a suite of RvS algorithms on the D4RL benchmark.
Finally, we highlight the limitations of current
RvS algorithms in comparison with traditional RL ones.
I NTRODUCTION
Leveraging large amounts of off-policy data is essential for using Reinforcement Learning (RL) in real-world settings (Levine et al., 2020).
Recently, a straightforward and efficient approach to use large datasets has been to treat RL as a Supervised Learning problem, namely RL via Supervised
Learning (RvS) (Emmons et al., 2021).
The simplicity of RvS comes from bypassing the difficulties of Temporal Difference (TD) learning (Sutton, 1988), to instead use the observed return to learn good policies.
Furthermore, the efficiency of RvS comes from the ability to learn from sub-optimal trajectories present in the dataset.
The RvS algorithms usually convert the RL problem into behavior cloning where the samples are weighted by, filtered by, or conditioned on the observed Monte Carlo (MC) return.
It is however unclear what objective these RvS algorithms are maximising.
We introduce the Density Estimation to Decision Making (D2D) framework to unify the RvS algorithms.
This novel framework allows us to formulate RvS as a two steps process.
The first step consists of performing density estimation over the dataset, which can be done via supervised learning.
The second step consists of sampling from the exponential tilt of the learned density to obtain samples that are both likely and have high observed MC returns.
Using the D2D framework, we can derive recent algorithms such as Return Conditioned Policy (Kumar et al., 2019), Upside-Down RL (Srivastava et al., 2019), Trajectory Transformer (Janner et al., 2021), Decision Transformer (Chen et al., 2021), and others.
We also introduce implicit RL via
Supervised learning, a new implicit algorithm leveraging advances in density modeling.
We benchmark these different algorithms using the D4RL framework (Fu et al., 2020).
Finally, we highlight the limitations of current RvS algorithms in comparison with traditional RL algorithms. asasasa t s t Grrrr t ......
A P ROBABILISTIC I NFERENCE P ERSPECTIVE ON RL VIA S UPERVISED L EARNING
Reinforcement Learning via Supervised Learning algorithms are trained on large datasets collected by agents of diverse expertise.
The promise of RvS algorithms is in leveraging large amounts of sub-optimal and near-optimal data to achieve performance beyond what would be attainable by only using less but near-optimal data.
Current RvS algorithms use the dataset to learn a good policy by weighting the transitions as a function of the return, by conditioning the action distribution on the return, or by learning a forward model of the environment to estimate the return of an action.
These different approaches do not look related at first glance, but consolidating them into an unified framework would improve our understanding of the connections between these methods and allows us to design better algorithms.
We introduce the Density to Decision (D2D) framework to unify the RvS methods listed above.
This framework consists of two steps: i) learning a density model using the dataset, and ii) using the learned density for decision making.
Using this framework, we demonstrate that these RvS methods approximate the same density but make different design choices in their approximation.
Finally, we propose an algorithm using a different design choice better suited for density estimation and efficient decision making.
2. F ROM D ENSITY E STIMATION TO D ECISION M AKING
Density Estimation.
The dataset is composed of the triplets: states, action and MC return ( s t , a t , G t ) .
Their graphical structure and dependencies are depicted in Figure 1.
The MC return (Sutton & Barto, 2018) can be understood as providing hindsight information about the outcome of the trajectory (Furuta et al., 2021).
The first step consists of learning a density over the triplets which can easily be estimated via supervised learning methods.
Specifically, the objective can be defined as the negative log-likelihood of the data:
[Equation]
Decision Making.
We are interested in making decisions better or as good as the best ones observed in the dataset.
Thus, we cannot simply sample from the learned density as this would result in decisions present in the dataset ignoring the quality of these decisions.
To improve decision making we have to direct the sampling towards the best decisions in the dataset.
A natural approach is to use the Exponential Tilt (Asmussen & Glynn, 2007; O’Donoghue et al., 2020) of the MC return, which is defined as:
[Equation]
where κ ( η ) is the cumulant generating function log E [ e η − 1 G ] and s t is observed.
As we decrease η → 0 , the density becomes concentrated around regions that are potentially less likely in the dataset but where the MC return is large.
We recover the dataset density by increasing η → ∞ Density Estimation ( E [Equation] Decision Making
[Equation] [Equation] independently of the MC return. Making η an important hyper-parameter as we want samples that are both likely in the dataset and have high MC return.
Alternatively, we can examine the log of the exponential tilt to gain intuition. Specifically, taking the log on both sides we obtain and note that log E [ e η − 1 G ] is a constant, maximizing p θ ( a , G | s t ; η ) is equivalent to maximizing the following: [Equation] (3) resulting in action MC return pairs that are likely under the dataset and where the MC return is large.
Connection to Control-as-Inference.
We note the similarity between D2D decision-making objective and the Control-as-Inference target density.
In both cases, the exponential function is used to turn a trajectory return into a positive number which can then be treated as a probability density.
Thus, trajectories with larger return have higher density making them more likely to be sampled.
The main difference with Control-as-Inference is in the training process.
Control-as-Inference introduces observed “optimality variables“ (Levine, 2018).
These optimality variables are used to define a backward message.
The trajectory density is decomposed by time step using the backward message which must be approximated using (soft) TD-learning.
Thus making training more challenging than D2D which can be performed by simple supervised learning.
Algorithm 1 Density to Decision
Input: Dataset D = { s t , a t , G t } Tt =
Input: Learning rate α
Input: Temperature η # Density Estimation while not converged do
Sample minibatch B ∼ D
# Compute loss
L ( θ ) ← (cid:80) ( s t , a t ,G t ) ∈B − log p θ ( a t , G t | s t )
Update parameters θ ← θ + α ∇ θ L ( θ ) end while # Decision Making s 0 ∼ µ ( s 0 ) t ←while episode not done do a [Equation] s [Equation]
∼ p
[Equation]
s [Equation] [Equation] ) t ← t +end while
2. A P ROBABILISTIC I NFERENCE P ERSPECTIVE ON R V S A LGORITHMS
Using our D2D framework, we can categorize multiple RvS algorithms as different design choices to approximate the same density.
An overview of the different methods an design choices are provided in Table 1.
Return Conditioned Policy (RCP).
One can model the return conditional policy p ( a | s , G ) and a state conditional distribution over return p ( G | s ) using supervised learning.
Conditioning on the return G allows the model to leverage trajectories with a sub-optimal observed return, potentially resulting in more efficient learning.
An additional advantage of this approach is that a simple Gaussian approximation to p θ ( a | s , G ) can capture a multimodal action distribution p ( a | s ) , which can be obtained by marginalising the return G , i.e., p ( a | s ) =
E p θ ( G | s )
[ p θ ( a | s , G )] .
To recover the richness of the training distribution, one would need to learn a multimodal distribution over MC return p θ ( G | s ) , but in practice only a subset of the training behaviors is desired: the best behavior.
The learning is done over the following factorisation:
[Equation]
where θ denotes the density that is learned, and the decision making distribution is obtained:
[Equation]
Kumar et al.
(2019) and Srivastava et al.
(2019) estimate the maximum return and then condition on it at the beginning of the episode.
For simplicity, it is also possible to ignore η and simply condition the policy on the maximum observed return in the dataset: p θ ( a | s t , G max ) as is done in Decision
Transformer (DT) (Chen et al., 2021).
Reweighted Behavior Cloning (RBC).
One can model a state-conditional action distribution p θ ( a | s ) and a state-action conditional return distribution p θ ( G | s , a ) .
Contrary to RCP which can capture a multimodal action distribution p ( a | s ) by combining simpler return conditional distributions, in turn RBC must directly learn a multimodal action distribution p θ ( a | s ) .
The decision making distribution can be approximated via reweighting the action samples.
Density estimation and decision making can be respectively done as:
[Equation]
where C is a normalising constant G ∼ p ( G | a , s t ) and a ∼ p ( a | s t ) .
For example, Trajectory Transformer (TT) (Janner et al., 2021) performs decision making by sampling multiple actions from p θ ( a | s t ) and evaluates p θ ( G | s t , a ) for each action using model roll-outs sampled from p
θ ( s t +1 , r t | s t , a t ) .
Again, η can be ignored in this case by taking the action with the highest observed return.
The ancestral sampling (Bishop & Nasrabadi, 2006) procedure used to select action in RBC and TT does not scale well for large action spaces and long planning horizons.
Return Weighted Regression (RWR).
It is possible to approximate the target density by bypassing the estimation of the return altogether.
A policy can be learned using importance weights to bias the policy toward actions with a higher observed return.
This has the advantage of being simple and computationally efficient.
Specifically, one can directly amortize the distribution using importance weights proportional to the exponential of the observed MC return: [Equation]
[Equation]
where c does not depend on θ .
Action selection can be performed efficiently using a forward pass p θ ( a | s ) .
We note the similarity with reward weighted regression (Peters & Schaal, 2007).
Alternatively, the weights could also be an indicator variable I { G>g } which is 1 when the observed MC return is higher than a return threshold g and 0 otherwise.
I MPLICIT RL VIA S UPERVISED L EARNING
Using the D2D framework to investigate current RvS algorithms highlights the design choices that are currently made for the density estimation and decision making phases.
We notice that no method directly model the joint distribution over the action and the return.
Modeling the joint presents certain difficulties, but could potentially provides certain advantages.
We now introduce a novel implicit method to model the joint distribution and that can be use to perform decision making in an efficient manner.
Using the D2D framework, we introduce a novel RvS algorithm.
As opposed to the previous methods that factorize the joint distribution into different conditional distributions, it is possible to model the joint distribution directly.
Modeling the joint has the advantages of using both the action distribution and the future return distribution signals to train a better model.
It also does not require ancestral sampling to produce and evaluate an action as sampling can be done at the same time.
3. D ENSITY E STIMATION P ROCEDURE
In order to model the joint distribution over action and return we leverages advances in Energy
Based Models (EBMs) (LeCun et al., 2006).
EBMs are a promising class of generative models that can easily model discontinuity and their effectiveness for behavior cloning has been demonstrated in Florence et al.
Specifically, the energy function is defined as:
[Equation]
where E θ ( s , a , G ) is the parametrized energy function and Z θ = (cid:82) a (cid:82) G exp( E θ ( s , a , G )) dGd a is the normalising constant.
Training the density model can be performed using the InfoNCE loss function defined as:
[Equation]
where { ˜ a ji , ˜ G ji } N neg.
j =1 is a set of negative counter-examples.
Negative sampling can be performed efficiently by minimizing ˜ a , ˜ G = arg min a ,G E θ ( a , G | s t ) using stochastic gradient Langevin dynamics (Welling & Teh, 2011) (SGLD), see Algorithm 2 for details. 3. D ECISION -M AKING P ROCEDURE
We now have access to a differentiable joint density p θ ( a , G | s ) from which we can sample action and MC return pairs using SGLD.
Doing so will result in action and MC return likely under the training dataset irrespective of the observed MC return, which is equivalent to implicit behavior cloning as done by Florence et al.
Instead, we want to bias sampling towards high region of observed MC return. To do so, we will use the exponential tilt in G to bias the sampling towards action, MC return pairs with high observed MC returns.
Specifically, we use the exponential tilt in G is defined as: [Equation] where κ ( η ) is the cumulant generating function log E [ e η − 1 G ] .
We note that the density is still differentiable and can be efficiently sampled via SGLD.
As in the D2D framework, η controls the trade-off between the likeliness of a ( a , G ) pair under the dataset and the MC return.
3. R ELATED W ORK
We note that other works have used EBM in RL.
Interestingly, most of the previous works treat the
Q -value as the energy and does not model how likely the Q function is (or the MC return in our case).
Sallans & Hinton (2004) parametrize the Restricted Boltzman Machine (Freund & Haussler, 1991;
Welling et al., 2004) negative energy function as the Q value.
And more recently, Haarnoja et al.
(2017) learn a Q function via the (soft-) Bellman equation, then treat the Q function as the negative energy function which is sampled via a stochastic neural network.
2 action 0.
M C R e t u r n ( G ) (a) p ( a , G | s 0 )    2 action 0.
M C R e t u r n ( G ) (b) p ( a , G | s 0 )
exp( G 0 .
1 )    2 action 0.
M C R e t u r n ( G ) (c) p ( a , G | s 0 ) exp( G 0 . 05 )  E XPERIMENTS 4. D IDACTIC E XAMPLE We present a simple didactic experiment to communicate better the advantages of modeling the joint distribution over the state, action, and MC return p ( a , G | s ) .
The agent starts the episode at the point (0 , 0) with a randomly selected target and must choose an angle α ∈ [ − π/ 2 , π/ 2] (the half-circle) as an action.
Reaching one of the targets results in a positive reward while the other one results in a negative reward.
The angle is then converted to a 2d action as a = (∆ x , ∆ y )
= (cos( α ) , sin( α )) .
The agent deterministically choose the angle to minimize the distance with the selected target.
A small amount of Gaussian noise is then added to the action a to make the policy stochastic.
There is no learning in this experiment.
Graphical depiction of the environment and of random trajectories can be seen in Figure 3 in Appendix A.1.
In Figure 2a, we examine the joint distribution of the angle α (x-axis) and the MC return (y-axis) conditioned on the initial state s 0 = (0 , 0) .
The distribution colors represent the (unobserved) target on which the behavior policy was conditioned.
We can observe that the distribution is bimodal in both the MC return and the angle distribution.
The stochasticity of the discounted MC return comes only from the random noise added to the policy as some policies take longer to reach the target.
For each target conditioned policy, we can observe a mode for the action at around π/ 4 ( − π/ 4 ) for the blue (orange) policy representing an angle of 45 ◦ ( − 45 ◦ )
Given the observed angle and return for s 0 in the dataset in Figure 2a, we are interested in generating a distribution over the best action.
Since the environment dynamics are fixed and the stochasticity comes from the policy alone, we can use the exponentially tilted density p ( a , s , G ) exp( η − 1 G ) with respect to G to generate a distribution over best actions.
In Figures 2b and 2c, we can observe that as we decrease the temperature, the density over the maximum return increases and so does the density over the angle α = π/ 4 .
We observe that even without access to the target on which the policy was conditioned, modeling the joint density including the MC return G allows us to recover the distribution over the desired policy.
Similarly, in offline RL, we do not observe the expertise level of the policies that collected the data.
We can nevertheless use the observed return as guidance to learn desirable behaviors.
4. C ONTINUOUS C ONTROL D ATASET
To better understand the importance of the different design choices in RvS, we compare the algorithms on the D4RL benchmark (Fu et al., 2020).
The D4RL benchmark provides a suite of tasks and datasets to evaluate offline RL agents (Levine et al., 2020).
We compare the different algorithms on three continuous locomotion tasks with diverse datasets. [Table]
For each algorithm, we fine tuned the hyper-parameters for 3 seeds on the half-cheetah mediumexpert, walker2d medium-expert and hopper medium-expert datasets.
We note that for RCP the best return to condition on is tricky to tweak and wildly differs across environments.
We tune the target return .
We note that our RCP implementation performs slightly better than Emmons et al. (2021) implementation on certain tasks such as walker2d medium-expert. While significantly worst on the hopper suite as it requires wildly different hyper-parameters than the half-cheetah and walker2d tasks. Algorithm 1 outlines the training and evaluation procedure.
We train the density estimation algorithms for 100k gradient steps and then evaluate them on 100 episodes (of 1000 steps each).
We repeat this procedure for three seeds.
As in Emmons et al.
we define G as the average future reward.
Additionally, we normalize each state coordinate to have mean 0 and standard deviation 1, and G to lie in [ − 1 , 1] .
We focus on the Gym Locomotion v2 tasks, consisting of the HalfCheetah, Hopper, and Walker2D datasets from the D4RL offline RL benchmark (Fu et al., 2020).
These datasets are collected by agents with medium and expert levels of expertise in solving the respective tasks.
The difficulty lies in leveraging sub-optimal and near-optimal trajectories to learn a policy.
In ?? we show the baseline comparisons for the medium-expert dataset for the previously mentioned tasks.
The values reported are the average return and standard deviation for 3 training initialisation seeds over 100 episodes each.
We can see that IRvS performs on par with the best among the other methods for the Hopper and Walker tasks.
D ISCUSSION
In this section, we highlight the differences between D2D and traditional Offline RL algorithms and discuss the advantages of both approaches.
We first compare the supervised density estimation procedure used by D2D algorithms as opposed to the TD learning approach used by Offline RL algorithms to evaluate actions.
We then compare the Exponential Tilt approach of D2D algorithms as opposed to the maximization of a learned Q value by Offline RL algorithms to perform action selection.
The first difference to note is the action evaluation procedure.
RvS uses the observed Monte Carlo return rather than a parametrized Q value to learn a policy.
On one hand, Q values can be challenging to learn when function approximation, off-policy, and bootstrapping are used as it can diverge (Van Hasselt et al., 2018), while the observed MC return provides an easy signal to use.
On the other hand, Q -values learned via TD learning (Sutton, 1988) perform trajectory stitching thus a good action in a bad trajectory can be evaluated independently of the trajectory outcome and be recognised as a good action.
While in D2D, a good action in a bad trajectory will only be evaluated based on the trajectory outcome.
A potential middle ground would be to use TD learning to perform on-policy evaluation (no policy optimization) as it is stable and provides good performances (Goo & Niekum, 2022; Brandfonbrener et al., 2021).
The second main difference between D2D and Offline RL algorithms is the action selection strategy.
For a given state, D2D algorithms limit action selection candidates to the actions that have been observed for that state in the dataset.
The agent might be optimistic for state-action pair that have not been observed in the dataset, and selecting unobserved actions would potentially result in worst performance than limiting action selection to observed actions.
Yet, limiting action selection could potentially limit the asymptotic performance of the agent.
D2D algorithms can explicitly control the risk-reward tradeoff via the temperature η .
For small η , the agent might select action with high, but less likely, MC return.
As we increase the temperature, the agent will seek less risk.
A similar strategy is used by popular offline RL algorithms.
For example, by using explicit density estimation to constraint the action selection to the actions observed in the dataset (Wu et al., 2022), training the
Q -value network to be conservative for unobserved actions (Kumar et al., 2020) or by constraining the policy using behavior cloning regularization (Fujimoto & Gu, 2021).
C ONCLUSION
In this work, we introduce the Density Estimation to Decision Making framework which unifies commonly used RvS Algorithms.
The 2 steps process brings clarity on the training and action selection of these algorithms, and also provides guidance on designing novel algorithms.
Using our framework we introduce IRvS which directly targets the joint density distribution over the action and return instead of the commonly used arbitrary density factorization.
We show that IRvS is competitive with the state-of-the-art RvS algorithms on the D4RL Mujoco locomotion tasks.
Finally, we discuss the trade-off of using RvS over traditional Offline RL algorithms.
Reinforcement Learning via Supervised Learning algorithms have been shown to be competitive with Offline RL algorithms for trajectory optimization.
An interesting next step would be to compare their performance, generalisation capabilities and robustness on more difficult tasks
A. S TOCHASTIC G RADIENT L ANGEVIN D YNAMICS
Input: Initial state s
Input: Initial action a ∼ U ( − 1 , 1)
Input: Initial MC return G ∼ U ( − 1 , 1)
Input: Learning rate α # Density Estimation for i ∈ { 1 , . . .
, 100 } do g ← ∇ ( a ,G ) E
θ ( s , a , G ) u ← clip ( 12 g + ϵ, − 0 . 5 , 0 . 5) , ϵ ∼ N (0 , 1) ( a , G ) ← ( a , G ) − αu end for
Return: ( s , a , G )