Policy Gradients Incorporating the Future

Reasoning about the future – understanding how decisions in the present time affect outcomes in the future – is one of the central challenges for reinforcement learning ( RL ) , especially in highly-stochastic or partially observable environments . While predicting the future directly is hard , in this work we introduce a method that allows an agent to “ look into the future ” without explicitly predicting it . Namely , we propose to allow an agent , during its training on past experience , to observe what actually happened in the future at that time , while enforcing an information bottleneck to avoid the agent overly relying on this privileged information . Coupled with recent advances in variational inference and a latent-variable autoregressive model , this gives our agent the ability to utilize rich and useful information about the future trajectory dynamics in addition to the present . Our method , Policy Gradients Incorporating the Future ( PGIF ) , is easy to implement and versatile , being applicable to virtually any policy gradient algorithm . We apply our proposed method to a number of off-the-shelf RL algorithms and show that PGIF is able to achieve higher reward faster in a variety of online and offline RL domains , as well as sparse-reward and partially observable environments . 1 INTRODUCTION . Fundamentally , reinforcement learning ( RL ) is composed of gathering useful information ( exploration ) and assigning credit to that information ( credit assignment ) . Both of these problems present their own unique learning challenges . In this work , we focus on credit assignment , which refers to the challenge of matching observed outcomes in the future to decisions made in the past . Humans appear to do this in a sample efficient manner ( Johnson-Laird , 2010 ) , and so it is natural to expect our own RL agents to do so as well . One of the most popular approaches to credit assignment , known as model-free RL , is to learn a value function to approximate the future return given a starting state and action . The value function is learned using experience of the agent acting in the environment via temporal difference ( TD ) methods ( Sutton , 1988 ) , which regress the value function to a target based on a combination of groundtruth returns achieved in the environment and the approximate value function itself . The need to bootstrap learning of the value function on its own estimates is known to lead to difficulties in practice , where one must achieve a careful balance between bias and variance ( Harutyunyan et al. , 2019 ; Weaver & Tao , 2001 ; Schulman et al. , 2016 ; Mnih et al. ) . If a slight imbalance arises , the consequences can be disastrous for learning ( Tsitsiklis & Van Roy , 1996 ; van Hasselt et al. , 2018 ; Sutton & Barto , 2018 ) . For example , in offline RL this issue is so pronounced that algorithms must apply strong regularizations on both learned policy and value function to achieve stable performance ( Wu et al. , 2020 ; Kumar et al. , 2019 ; Zhang et al. , 2021 ; Nachum et al. , 2019 ) . The model-free approach plays dual to the model-based approach , where an agent learns a dynamics and reward model of the environment , and then learns an agent to optimize behavior in this model . Thus , credit assignment boils down to utilizing an appropriate planning algorithm that can perform multiple rollouts in the model , effectively allowing the agent to “ look into the future ” ( Racanière et al. , 2017 ; Pascanu et al. , 2017 ) to determine cause-and-effect ( Sutton , 1991 ; Peng et al. , 2018 ; Abbas et al. ) . While model-based RL may appear more straightforward , learning an accurate model is a challenge in practice , presenting its own sample-efficiency problems ( Wang et al. , 2020 ) as well as memory and computational issues ( Łukasz Kaiser et al. , 2020 ) . Model-based approaches are thus most beneficial when the environment exhibits some level of regularity ( Fra , 2019 ) . Beyond these issues , credit assignment in both model-free and model-based RL is further exacerbated by partially observable environments , in which the full environment state is not known to the learning agent . Thus it is infeasible to predict future events accurately . When applied to such non-Markovian domains , model-free algorithms relying on bootstrapping and value function approximation tend to be biased ( Singh , 1994 ) . On the other hand for model-based approaches , learning an accurate dynamics model in such domains is a difficult , potentially ill-defined problem ( Suematsu & Hayashi , 1999 ; Bush & Pineau , 2009 ) . In this work , we aim to circumvent these challenges . We propose a simple modification to model-free RL that allows the learned policy and value function to “ look into the future ” but without the need to learn an accurate model . Namely , we propose to modify the policy and value function to not only condition on the presently observed state and action but also on the subsequent trajectory ( sampled by the agent as it was interacting with the environment ) following this state and action . This way , our method mitigates potential approximation or feasibility issues in accurately modeling the future . To ensure that the learned policy and value function remains relevant during inference ( i.e. , data collection ) when the future trajectory is unavailable , we place an information bottleneck ( Saxe et al. , 2018 ; Tishby et al. , 1999 ) on the additional inputs , encouraging the learned functions to minimize their reliance on this privileged information . One may thus view our method as an instance of teacher forcing or Z-forcing ( Goyal et al. , 2017 ; Lamb et al. , 2016 ) where our student is the learned policy and value function and the teacher is some function of the information in the future trajectory . It is well known that combining a strong autoregressive decoder with latent variables while ensuring that they carry useful information is difficult ( Bowman et al. , 2016 ; Chen et al. , 2017 ) . Z-forcing enforces the incorporation of relevant information by learning latent variables that predict the future steps . Practically , our method , Policy Gradients Incorporating the Future ( PGIF ) , is versatile and easy to implement . We use either a backwards RNN or a transformer to inject downstream information from the observed trajectories by way of latent variables , with a KL divergence regularization on these latents . We apply PGIF on top of a variety of off-the-shelf RL algorithms , including RNNbased PPO ( Schulman et al. , 2017 ) , SAC ( Haarnoja et al. , 2018 ) , and BRAC ( Wu et al. , 2020 ) , and evaluate these algorithms on online and offline RL as well as sparse-reward and partially observable environments . In all of these domains , we demonstrate the ability of PGIF to achieve higher returns faster compared to these existing RL algorithms on their own , thus showing that our proposed method is both versatile and beneficial in practice . 2 BACKGROUND AND NOTATION . We begin by providing a brief overview of the notation and preliminary concepts that we will use in our later derivations . Markov Decision Processes ( MDPs ) MDPs are defined by a tuple 〈S , A , P , R , ρ0 , γ〉 where S is a set of states , A is a set of actions , P is a transition kernel giving a probability P ( s′|s , a ) over next states given the current state and action , R : S × A → [ Rmin , Rmax ] is a reward function , ρ0 is an initial state distribution , and γ ∈ [ 0 , 1 ) is a discount factor . An agent in this MDP is a stationary policy π giving a probability π ( a|s ) over actions at any state s ∈ S. A policy π interacts with the MDP by starting at s0 ∼ ρ0 and then at time t ≥ 0 sampling an action at ∼ π ( st ) at which point the MDP provides an immediate reward R ( st , at ) and transitions to a next state st+1 ∼ P ( st , at ) . The interaction ends when the agent encounters some terminal state sT . The value function V π : S → R of a policy is defined as V π ( s ) = Eπ [ ∑T−1 t=0 γ trt|s0 = s ] , where Eπ denotes the expectation of following π in the MDP and T is a random variable denoting when a terminal state is reached . Similarly , the state-action value function Qπ : S × A → R is defined as Qπ ( s , a ) = Eπ [ ∑T−1 t=0 γ trt|s0 = s , a0 = a ] . The advantage Aπ is then given by Aπ ( s , a ) = Qπ ( s , a ) − V π ( s ) . We denote ρπ as the distribution over trajectories τ = ( s0 , a0 , r0 , . . . , sT ) sampled by π when interacting with the MDP . During learning , π is typically parameterized ( e.g. , by a neural network ) , and in this case , we use πθ to denote this parameterized policy with learning parameters given by θ . The policy gradient theorem ( Sutton et al. , 1999 ) states that , in order to optimize the RL objective Es0∼ρ0 [ V π ( s0 ) ] , a parameterized policy should be updated with respect to the gradient of the following loss ( ignoring any gradients through ρπθ , where we denote stopping the gradient with θ̄ ) , JPG ( πθ ) = Eτ∼ρπθ̄ [ ∑T−1 t=0 γ t·Q̂t log πθ ( at|st ) ] , ( 1 ) where Q̂t is an unbiased estimate of Qπ ( st , at ) . In the simplest case , Q̂t is the empirically observed future discounted return following st , at . In other cases , an approximate Q-value or advantage function is incorporated to trade-off between the bias and variance in the policy gradients . When the Q or V value function is parameterized , we will use ψ to denote its parameters . For example , the policy gradient loss with a parameterized Qψ is given by , JPG ( πθ , Qψ ) = Eτ∼ρπθ [ T−1∑ t=0 γt ·Qψ ( st , at ) log πθ ( at|st ) ] . ( 2 ) The value function Qψ is typically learned via some regression-based temporal differencing method . For example , JTD ( Qψ ) = Eτ∼ρπθ [ T−1∑ t=0 ( Q̂t −Qψ ( st , at ) ) 2 ] . ( 3 ) Stochastic Latent Variable Models In our derivations , we will utilize parameterized policies and value functions conditioned on auxiliary inputs given by stochastic latent variables . That is , we consider a latent space Z , typically a real-valued vector space . We define a parameterized policy that is conditioned on this latent variable as πθ ( a|s , z ) for a ∈ A , s ∈ S , z ∈ Z ; i.e. , πθ takes in states and latent variables and produces a distribution over actions . In this way , one can consider the latent variable z as modulating the behavior of πθ in the MDP . During interactions with the MDP or during training , the latent variables themselves are generated by some stochastic process , thus determining the behavior of πθ . For example , in the simplest case z may be sampled from a latent prior pυ ( Z ) ( z|s ) , parameterized by υ ( Z ) . Thus , during interactions with the MDP actions at are sampled as at ∼ πθ ( st , zt ) , zt ∼ pυ ( Z ) ( st ) . We treat parameterized latent variable value functions analogously . Specifically , we consider a latent space U and a parameterized Q-value function as Qψ ( s , a , u ) . A prior distribution over these latent variables may be denoted by pυ ( U ) ( u|s ) .

Summary: The work proposes to run two recurrent LSTM neural networks backwards from the end of an episode, giving the rewards $r$ and states $s$ as input to the networks, and producing two separate latent state ($z$, $u$) distributions for each time step. Such procedure can be done during training after having collected an episode. Once, the $z$ and $u$ are computed, these are given as an input to the Q-function $Q(s,a,u)$, and the policy network $\pi(a|s,z)$ (note that the action was already sampled during the episode, but here the probability is recomputed using the new $z$). And the policy gradient/critic updates are compute as usual together with this augmented state. One remaining question is how to pick the $z$ during the episode when the future states are not known yet. This is done by training a prior network $p(z|s)$ by minimizing the KL divergence to the latent state distribution obtained from the backwards network (then one can just sample a $z$ from this prior network during the episode). To aid the backwards network in learning meaningful latent representations, they use z-forcing, which is a technique where the variable $z$ is used directly to predict some quantity to ensure that it incorporates useful information. They consider two options: 1) predict $b$ the hidden state of the backward RNN, 2) predict the values, rewards, discounts. (Note that this z-forcing is done for both $z$ and $u$).  They perform experiments on Bsuite umbrella-length, Gym-minigrid, custom partially observable MuJoCo environments, Offline RL on MuJoCo D4RL tasks. They compare with PPO, SAC, and other baselines. The experimental results improved in all cases, but their method also requires more computational time due to having to train the LSTMs.