Year: 2021

Abstract:
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.

text: 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.

heading: 1 INTRODUCTION
text: 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.

heading: 2 BACKGROUND AND NOTATION
text: 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).

heading: 3 POLICY GRADIENTS INCORPORATING THE FUTURE
text: Our method aims to allow a policy during training to leverage future information for learning control. We propose to utilize stochastic latent variable models to this end. Namely, we propose to train πθ and Qψ with latent variables (z,u) = {(zt, ut)}T−1t=0 sampled from a learned function qφ(τ) which has access to the full trajectory. For example, the PGIF form of the policy gradient objective in (2) may be expressed as
JPGIF, PG(πθ, Qψ, qφ) = Eτ∼ρπ,(z,u)∼qφ(τ) [ T−1∑ t=0 γt ·Qψ(st, at, ut) log πθ(at|st, zt) ] . (4)
The PGIF form of the temporal difference objective in (3) may be expressed analogously as,
JPGIF,TD(Qψ, qφ) = Eτ∼ρπθ ,u∼qφ(τ) [ T−1∑ t=0 (Q̂t −Qψ(st, at, ut))2 ] . (5)
It is clear that any RL objective which trains policies and/or value functions on trajectories can be adapted to a PGIF form in a straightforward manner. For example, in our experiments we will apply PGIF to an LSTM-based PPO (Schulman et al., 2017), SAC (Haarnoja et al., 2018), and BRAC (Wu et al., 2020).
While the PGIF-style objectives above adequately achieve our aim of allowing a policy to leverage future trajectory information during training, they also present a challenge during inference. When performing online interactions with the environment, one cannot evaluate qφ(τ), since the full trajectory τ is not yet observed. Therefore, while we want to give πθ and Qψ the ability to look at the full τ during training, we do not want their predictions to overly rely on this privileged information. To this end, we introduce a regularization on qφ in terms of a KL divergence from a prior distribution pυ (τ) := {pυ (zt, ut|st)}T−1t=0 which conditions (zt, ut) only on st. Thus, in the case of policy gradient, the full loss is,
JPGIF-KL(πθ, Qψ, qφ) = JPGIF, PG(πθ, Qψ, qφ) + βEτ∼ρπ [DKL(qφ(τ)‖pυ(τ))] , (6)
where β is the weight of the divergence term. The introduction of this prior thus solves two problems: (1) it encourages the learned policies and value functions to not overly rely on information beyond the immediate state; (2) it provides a mechanism for inference, namely using latent samples from pυ (s) when interacting with the environment.
Parameterization of qφ In our implementation, we parameterize qφ(τ) as an RNN operating in reverse order on τ . Specifically, we use an LSTM network to process the states in τ backwards, to yield LSTM hidden states b = {bt}T−1t=0 . The function qφ is then given by Gaussian distributions with mean and variance at time t derived from the backwards state bt. In practice, to avoid potential interfering gradients from the objectives of πθ and Qψ, we use separate RNNs with independent parameters qφ(Z) , qφ(U) for zt, ut, respectively. See Figure 1 for a graphical diagram of the training and inference procedure for PGIF. In our empirical studies shown in Appendix G, we will also show that a transformer (Vaswani et al., 2017) can be used in place of an RNN with minimal decrease in performance, providing more computational efficiency and potentially allowing for better propagation of information over time.

heading: 3.1 VARIATIONAL INFORMATION BOTTLENECK INTERPRETATION
text: The KL regularization we employ above may be interpreted as a variational information bottleneck, constraining the mutual information between the latent variable distribution and the trajectory τ . Here we provide a brief derivation establishing this equivalence.
For simplicity, we consider a specific timestep t ∈ N and a starting state st = s. Let T≥t denote the random variable for all information contained after and including timestep t in trajectory τ induced by π. Let Ut be the random variable for latents ut induced by qφ(τ≥t) conditioned on all steps in the trajectory after and including t. Consider a constrained objective minimizing JPGIF,TD while enforcing an upper bound Imax on the mutual information between the distribution of trajectory steps and the distribution of latent variables I(T≥t,Ut|st = s). This objective is given by,
min ψ,φ JPGIF,TD(Qψ, qφ|st = s) := Eτ∼ρπθ (·|st=s),ut∼qφ(τ≥t)[(Q̂t −Qψ(st, at, ut)) 2]
s.t. I(T≥t,Ut|st = s) ≤ Imax. (7)
Recall the definition of mutual information: I(T≥t,Ut|st = s) = ∫ p(τ≥t, ut|st = s) log
p(τ≥t, ut|st = s) p(τ≥t|st = s)p(ut|st = s) dutdτ≥t
= ∫ ρπθ (τ≥t|st = s)qφ(ut|τ≥t) log
qφ(ut|τ≥t) p(ut|st = s) dutdτ≥t,
(8)
where p(ut|st = s) is the marginal distribution over the latent variable p(ut|st = s) =∫ qφ(ut|τ≥t)ρπ(τ≥t|st = s)dτ≥t. This marginal is intractable to compute directly, and so to approximate this marginal we introduce a variational distribution h(ut|st = s). By definition we know that DKL[p(ut|st = s)‖h(ut|st = s)] ≥ 0. We can then see that ∫ p(ut|st = s) log p(ut|st = s)dut ≥∫
p(ut|st = s) log h(ut|st = s)dut. We therefore derive the upper bound for use in equation 7 as, I(T≥t,Ut|st = s) ≤ ∫ ρπ(τ≥t|st = s)qφ(ut|τ≥t) log
qφ(ut|τ≥t) h(ut|st = s) dutdτ≥t
≤ Eτ∼ρπθ (·|st=s) [ DKL(qφ(u|τ≥t)‖h(ut|st = s)) ] .
(9)
We can subsume the constraint into the objective as,
min ψ,φ JPGIF,TD(Qψ, qφ|st = s) + β ( Eτ∼ρπθ (·|st=s) [DKL (qφ(τ≥t)‖h(ut|st = s))]− Imax ) .
By taking h to be our learned prior pυ , we recover the single step (st = s), TD analogue of the PGIF objective in (6), offset by a constant β · Imax, as desired.

heading: 3.2 Z-FORCING WITH AUXILIARY LOSSES
text: While our proposed training architecture enables the policy and value function to look at the full trajectory τ , in practice it may be difficult for the trajectory information to propagate, especially in settings with highly sparse learning signals. In fact, it is known that such latent variable models may ignore the latent variables due to optimization issues, completely negating any potential benefit (Bengio et al., 2015). To circumvent these issues, we make use of the idea of Z-forcing (Goyal et al., 2021), which employs auxiliary losses and models to force the latent variables to encode information about the future. We denote this loss as JAx(ζ) where ζ is the set of parameters in any auxiliary models, and elaborate on the main forms of this loss which we consider below. We do not combine different types of auxiliary losses in our total loss function.
We emphasize that these auxiliary losses are applied exclusively to the backwards encoder qφ(τ) and have no direct impact on the learned policy or value functions. In fact, we conduct ablations in Appendix F showing that these auxiliary losses applied directly to policy and value functions without PGIF perform worse.
State based forcing (Force) A simple way to force state information to be encoded is to derive conditional generative models pζ(bt|zt) over the backwards states given the inferred latent variables zt ∼ qφ(Z)(zt|bt), and similarly for the latents ut. We can write this auxiliary objective as a maximum log-likelihood loss JAx(ζ) = −Eq
φ(Z) (zt|bt)[log pζ(bt|zt)]. This way, we enforce the noisy mapping
bt → zt defined by qφ(Z) to not be too noisy so as to completely remove any information from bt.
Value prediction networks (VPN) A more sophisticated approach to force information to be propagated is to use an autoencoder-like, model-based auxiliary loss. To this end, we take inspiration from VPNs (Oh et al., 2017), and apply an auxiliary loss that uses bt to predict future rewards, values, and discounts. Note that, in principle, bt already has access to this information, by virtue of the backwards RNN or transformer conditioned on the future trajectory. Thus, this auxiliary loss only serves to enforce that the RNN or transformer dutifully propagates this information from its inputs. We also note that, in contrast to the state based forcing described above, this approach only enforces bt to contain the relevant information, and it is up to the RL loss whether this information should be propagated to the latents zt, ut. We give a detailed explanation of VPNs in Appendix B.

heading: 3.3 FULL ALGORITHM
text: The full learning objective for PGIF is thus composed of three components: First, a latent-variable augmented RL objective, e.g., policy gradient as shown in (4). Second, a KL regularizer, e.g., as shown in (6). Finally, an auxiliary loss, given by either state based forcing (Force) or value prediction networks (VPN). We present an example pseudocode of a PGIF-style policy gradient with learned value function in Algorithm 1. See Appendix C for further details, including how to adaptively tune the coefficients on the KL and auxiliary loss components as well as more specific pseudocode algorithms for advantage policy gradient and soft actor-critic.
Algorithm 1 PGIF Algorithm with State-Action Value Function Estimation Require: Initial parameters: θ, υ(U), υ(Z), φ(U), φ(Z), ψ, ζPG, ζTD, Weights: αPG, αTD, βPG, βTD 1: for policy-step k = 0, 1, 2, . . . , N do 2: Collect set of Trajectories D = {τi, . . . } : 3: repeat 4: zt ∼ pυ(Z)(zt) 5: Execute: πθ(at|st, zt) and observe rt, st+1 from environment. 6: until episode termination 7: for Trajectory: τi ∈ D do 8: bZ = BackwardsLSTMZ(τi) (Operates over the entire trajectory) 9: bU = BackwardsLSTMU(τi) 10: τi = τi ∪ {bZ,bU} 11: ∀{s,a, r,bZ,bU} ∈ D: 12: Derive JAx-PG(ζPG), JAx-TD(ζTD) according to any auxiliary loss. 13: DTD = DKL(qφ(U)(u|bU)‖pυ(U)(u|s)) 14: JTD = Eτ∈D [ JPGIF,TD(Qψ, qφ(U)) + αTDJAx-TD(ζTD) + βTDDTD
] 15: DPG = DKL(qφ(Z)(z|bZ)‖pυ(Z)(z|s)) 16: JPG = Eτ∈D [ JPGIF, PG(πθ, Qψ, qφ(Z)) + αPGJAx-PG(ζPG) + βPGDPG
] 17: Update all parameters w.r.t: JPG and JTD

heading: 4 RELATED WORK
text: We review relevant works in the literature in this section, with additional discussions in Appendix B.
Incorporating the future Recent works in model-based RL have considered incorporating the future by way of dynamically leveraging rollouts of various horizon lengths and then using them for policy improvement (Buckman et al., 2018) . Z-forcing and stochastic dynamics models have been applied directly to learning environmental models and for behavioral cloning while incorporating the future but not for online or offline continuous control (Ke et al., 2019). Our present work is unique for incorporating Z-forcing and conditioning on the future in the model-free RL setting. A few other methods explore the future in less direct ways. For example, RL Upside down (Schmidhuber, 2020) uses both reward (or desired return) and state to predict actions, turning RL into a supervised learning problem.
Hindsight Hindsight credit assignment introduces the notion of incorporating the future of a trajectory by assigning credit based on the likelihood of an action leading to an outcome in the future (Harutyunyan et al., 2019). These methods were extended using a framework similar to ours, leveraging a backwards RNN to incorporate information in hindsight (Mesnard et al., 2021). Still, there are a number of differences compared to our own work. (1) Only the value function (rather than both the value and policy functions) is provided access to the future trajectory, whereas we show that allowing the actor access has benefits in some tasks (Appendix F). (2) There is no KL information bottleneck; rather information is constrained via an action prediction objective. (3) These previous works do not employ any Z-forcing, while it is well-known that learning useful latent variables in the presence of an autoregressive decoder is difficult without Z-forcing (Bayer & Osendorfer, 2015); in fact, in our own preliminary experiments we found our algorithm performs significantly worse without any auxiliary losses. Value driven hindsight modelling (HiMo) (Guez et al., 2020) proposes a hindsight value function, which is conditioned on future information in the trajectory. In this work, they primarily use the hindsight value function (separate from the agent’s value function) to learn a low-dimensional representation, which is distilled to a non-hindsight representation that is then used by the actual actor and critic. Thus there are a few key differences from our work: (1) There are no gradients passing from the RL loss to the representation loss (only the separate value prediction loss is used to train the representation), thus this method is arguably less end-to-end than PGIF; (2) The only mechanism for controlling the amount of information in the representation is through its dimension size and look-ahead, while using a KL penalty is more flexible. Our method is more versatile than previous works, being applicable to off-policy and offline RL settings rather than purely on-policy RL, as in these previous works. Nevertheless, it is an interesting avenue for future work to investigate how to combine the best of both approaches, especially with guarantees on variance reduction of policy gradient estimators (Nota et al., 2021) and hierarchical policies (Wulfmeier et al., 2021).

heading: 5 EXPERIMENTS
text: We now provide a wide array of empirical evaluations of our method, PGIF, encompassing tasks with delayed rewards, sparse rewards, online access to the environment, offline access to the environment, and partial observability. In the appendix, we include further demonstrations of PGIF applied to the challenging AntMaze environment (Sec. H) with substantial performance improvements, online RL with full observability (Sec. E) with improvements over a SAC baseline, and numerous ablation analyses (Sec. F, I) identifying the components of PGIF that are responsible for performance.
In our online RL experiments, we compare against Soft Actor Critic (SAC) (Haarnoja et al., 2018) and Proximal Policy Optimization (PPO) (Schulman et al., 2017) with an LSTM layer in the policy network. The comparison with PPO is particularly important since this method also leverages a form of artificial memory of the past (but not forward-looking like in PGIF). SAC is a state-of-the-art model-free RL method that employs a policy entropy term in a policy gradient objective and shows optimal performance and stability when compared with other online deep RL benchmarks. We show the hyper-parameters for each experiment in the Appendix D. In addition, we explore using our method with a transformer as opposed to an LSTM as the backwards network in Appendix G.

heading: 5.1 CREDIT ASSIGNMENT AND SPARSE REWARDS
text: PGIF aims to provide a better representation for credit assignment by accounting for downstream information rather than current. This can be used as additional input to the value function or policy. We find this helps address two very related problems, optimal credit assignment and learning with sparse rewards. The former problem is exasperated by the latter since rewards are rarely obtained, possibly only at the end of the episode. The relevance of each action (credit) must be assigned using only these rare rewards.
Method Mean Bsuite-score
PPO-LSTM 0.33± 0.09 SAC 0.41± 0.03 DQN (3-step return) 0.44± 0.06 DQN (1-step return) 0.51± 0.03 PGIF-PPO (VPN) 0.46± 0.09 PGIF-PPO (Force) 0.26± 0.10 PGIF-SAC (VPN) 0.60± 0.04 PGIF-SAC (Force) 0.58± 0.08 PGIF-DQN (Force) 0.58± 0.04
the final reward to the initial timestep, more so than either the one-step backups used in SAC or the multi-step return regressions used in PPO-LSTM can.
We continue to the Gym-MiniGrid (Chevalier-Boisvert et al., 2018) set of partially-observable environments, for which rewards are sparse and there is a non-zero reward only when the agent completes the task (with the reward proportional to the time taken). The agent is only given a local ego-centric view of its environment. For example, in DoorKey a key must first be obtained and then a door opened leading to another room with the goal state. It is difficult, but essential, to assign credit to obtaining the key and opening the door in the correct sequence. We compare against PPO-LSTM and an implementation of HiMo, which also suggests to use information from the future during agent learning (Guez et al., 2020). We show the results of these experiments and examples of the grids in Figure 2, with PGIF-PPO outperforming both baselines. The difference is especially stark on the Unlock Environment where credit assignment is the key challenge and the standard PPO baseline
fails to converge to any rewards. In LavaGap, there is no significant improvement. We hypothesize that this is because the most common trajectories generated during early training are ones where the agent enters the lava. There is also no negative reward for entering the lava, and the episode is simply terminated, making it more difficult to avoid. This creates a situation where our future representation is highly suboptimal, especially during early training

heading: 5.2 PARTIAL OBSERVABILITY
text: We now aim to show that our method is not only effective in fully-observed Markovian settings, but also in environments with partial observability. This set of experiments uses the MuJoCo robotics simulator (Todorov et al., 2012) suite of continuous control tasks. These are a set of popular environments used in both online and offline deep RL works (Fujimoto et al., 2019; 2018) and provides an easily comparable benchmark for evaluating algorithm sample efficiency and reward performance. As in previous work (Yang & Nachum, 2021), we introduce an easy modification to these tasks to make the environment partially observable thereby increasing the difficulty: We zero-out a random dimension of the state space at each data collection step. This helps us test partial observability, a feature common in robotic agent tasks.
We compare a PGIF-style SAC implementation to vanilla SAC and PPO-LSTM on these domains. We show the results of these experiments in Fig. 3. We find that PGIF can provide improved performance on these difficult tasks, suggesting that PGIF is able to leverage future information in the trajectory to appropriately avoid uncertainties about the environment, more so than when only conditioning on the immediate state (vanilla SAC) or even when conditioning on the entire past trajectory (PPO-LSTM). Interestingly, we find that the simple state based forcing (Force) performs more consistently better than the more sophisticated VPN based forcing. See Appendix E for online evaluations without partial observability.

heading: 5.3 OFFLINE RL EVALUATIONS
text: To asses if our method is effective in an offline RL setting, we evaluate our proposed algorithm in several continuous control offline RL tasks (Fu et al., 2020) against Behavior Regularized Actor
Critic (BRAC) (Wu et al., 2020) and Batch-Constrained Q-learning (BCQ) (Fujimoto et al., 2019). BRAC operates as a straightforward modification of SAC, penalizing the value function using a measure of divergence (KL) between the behaviour policy and the learned agent policy. For our PGIF algorithm, we use BRAC as the starting point. For these offline MuJoCo tasks, we examine D4RL datasets classified as medium (where the training of the agent is ended after achieving a "medium" level performance) and medium expert (where medium and expert data is mixed) (Fu et al., 2020). Datasets that contain these sub-optimal trajectories present a realistic problem for offline RL algorithms. We also include an offline version of the AntMaze, which is particularly challenging due to sparse rewards. We show the results of these experiments in Table 2. We find that our method outperforms the baselines in all but one of the tasks in terms of final episodic reward. We hypothesize in the medium-expert setting that we perform slightly worse due to the lack of action diversity which makes learning a dynamics representation difficult. Interestingly, in contrast to the online results, we find that VPN based forcing performs better than state based forcing, although even state based forcing usually performs better than the baseline methods.

heading: 6 DISCUSSION
text: In this work, we consider the problem of incorporating information from the entire trajectory in model-free online and offline RL algorithms, enabling an agent to use information about the future to accelerate and improve its learning. Our empirical results attest to the versatility of our method. The benefits of our method are apparent in both online and offline settings, which is a rare phenomenon given that many previous offline RL works suggest that what works well in online RL often transfers poorly to offline settings, and vice versa (Fujimoto et al., 2019). Beyond just online and offline RL, our results encompass partial observability, sparse rewards, delayed rewards, and sub-optimal datasets, demonstrating the ability for PGIF to achieve higher reward faster in all settings.
We also wish to highlight potential risks in our work. Specifically, the use of future information in PGIF may exacerbate biases present in the experience or offline dataset. For example, it is known that expert datasets lack action diversity. Further conditioning on the future in this dataset could force these biases to be incorporated more easily into the learning algorithm. Some other biases that may arise are a trajectory that contains some marginal reward and is therefore incorporated into our policy/value function with PGIF. This could hamper exploration of the agent and prevent discovery of states that yield optimal reward. It may be interest to combine our method with exploration strategies. Perhaps there are benefits in decreasing access to future information during initial exploratory agent steps so the agent is better able to explore. Furthermore, we find that our method is slower to train than the baselines we compare to, due to the fact that the architecture requires training an LSTM, with at times, very long trajectories as input. We must use the entire trajectories to train this LSTM. These issues can largely be alleviated by using a transformer, with minimal difference in performance.

Review summaries:
The paper proposes to incorporate future information for more accurate policy and Q-function estimation, and consequently to derive variants of policy gradient based algorithms such as PPO, SAC, and BRAC. The authors argue that the ability to condition on the future information enables better credit assignment. The authors empirically demonstrate that their approach improves performance on a range of tasks that feature various challenges, including sparse rewards credit assignment, partial observability, and offline RL.

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.

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.

The paper proposes a method for incorporating future trajectory information when training model-free RL methods. Two technical challenges were reducing over-reliance on future information and allowing the policy to make predictions without future trajectory information at deployment time. These are addressed by introducing an information bottleneck. Optimization issues with latent variable models are tackled using two different Z-forcing approaches. The method shows favorable empirical performance in both online and offline settings.