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: This paper presents Policy Gradients Incorporating the Future (PGIF), a novel approach to incorporate future information during training to improve performance of model-free RL agents that must overcome the challenges of planning in environments with sparse rewards. The approach involves relying on recent work in "Z-forcing" in which training incorporates information about the future, encoded as a learned latent state that depends on the full not-yet-executed trajectory and state information. The future information helps the agent to learn effectively and overcome challenges associated with credit assignment notorious in partially observed environments; additional losses are added to ensure that the latent sate contains sufficient task-relevant information and that the planner does not rely too heavily on future information. The authors show solid results in a number of challenging environments, outperforming competitive approaches (including PPO) in nearly all of these, showing the effectiveness of the technique in both Online and Offline RL experiments.