Policy Gradients Incorporating the Future
Authors: David Venuto
Elaine Lau
Doina Precup
Ofir Nachum
Venue: ICLR 2022 Poster (2022)

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 \emph{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 \emph{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. 

Published as a conference paper at ICLR 2022
POLICY GRADIENTS INCORPORATING THE FUTURE
David Venuto Elaine Lau2, Doina Precup 12.3, Ofir Nachum Mila, McGill University, 'DeepMind, *Google Brain david. venuto@mail mmcg .mcgill.ca ill.ca
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 (RI.), especially in highly-stochastie
partially observable environments. While predicting the future directly is hard, in this work we introduce method that allows 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 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 dynamies in addition 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 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.
INTRODUCTION
Fundamentally, reinforcement learning (RL) is composed of gathering useful information (explo- ration) 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 RI. 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 starting state and action. The value function learned using experience of the agent acting in the environment via temporal difference (TD) methods (Sutton, 1988), which regress the value function groundtruth returns achieved in the environment and the approximate value function itself. The need to bootstrap learning of the value function its 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 al., 2016; Mnih et al.). If 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 pronounced that algorithms must apply strong regularizations on both learned policy and value function to achieve stable performance (Wu et al., 2020; Kumar al., 2019: Zhang al., 2021; Nachum al., 2019).
target based on combination of
The model-free approach plays dual to the model- basell approach, where an agent learns 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 appropriate planning algorithm that can perform multiple rollouts the model, effectively allowing the agent to "look into the future" (Racanière et al., 2017; Pascanu et al., 2017) to determine cause-and-effee (Sutton, 1991; Peng et al., 2018; Abbas et al.). While model-based RL may appear more straightforward, learning an accurate model problems (Wang al., 2020) well
is challenge in practice, presenting its sample-efficiens as memory and computational issues (Lukasz Kaiser et al., 2020). Model-based approaches are thus most beneficial when the environment exhibits some level of regularity (Fra, 2019).

---

Published as a conference paper at ICLR 2022
Beyond these issues, credit assignment in both model-free and model-basee 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) 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 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. 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-foreing 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 RI. algorithms, including RNN- based 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 brief overview of the notation and preliminary concepts that will use in our later derivations.
Markov Decision Processes (MDPs) MDPs are defined by a tuple (S,A,P, po. 7) where S is 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, Po is an initial state distribution, and _ [0, 1) is discount factor. An agent in this MDP is a stationary policy π giving a probability 7 (als) over actions any state S. A policy π interacts with the MDP by starting at 80 Pa and then at time 2 O sampling an action a, T(::) at which point the MDP provides an immediate reward R(st, 0t) and transitions a next state 8141 ~ 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 (T(s) E, E y're|so = s). where E, denotes the expectation of following N in the MDP and a random variable denoting when terminal state is reached. Similarly, the state-action value function Q" S x A + R is defined as (=((s Er T- s'r,|s0
s, ag a). The advantage A is then given by A (s, (s, - (*)... We denote P.r as the distribution over trajectories T = (30, Co, ro, ST) sampled by क when interacting with the MDP.
During learning, π is typically parameterized (c.g., by a neural network), and in this case, we use Te to denote this parameterized policy with learning parameters given by 0. The policy gradient theorem (Sutton et al., 1999) states that, in order to optimize the RL objective Exg=ao [V (so)). parameterized policy should be updated with respect to the gradient of the following loss (ignoring
2

---

Published as a conference paper at ICLR 2022
any gradients through Prg. where we denote stopping the gradient with A),
[ = B, r 1 ff mg
where Q. is an unbiased estimate of (se, ar). In the simplest case, Qr is the empirically observed future discounted return following 8t, de- In other cases, 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 parameterized, we will use & to denote parameters. For example, the policy gradient loss with a parameterized Qu is given by,
[PG Q+) = Ey Eu-ep
Σ r2 (u ai) log lo (az[ss)))
(2)
The value function Qs is typically learned via some regression-based temporal differencing method. For example,
Ei ac))"
JmD(Qs) = Brrppro
(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 latent space Z, typically real-valued vector space. We define parameterized policy that is conditioned on this latent variable as my (a)a,, 2) for a E A, s E S, 2 E Z; i.c. ny takes in states and latent variables and produces distribution over actions. In this way, the latent variable z as modulating the behavior of Ty in the MDP. During interactions with the MDP during training, the latent variables themselves are generated by some stochastic process, thus determining the behavior of Ty. For example, in the simplest case may be sampled from latent prior (z|s), parameterized by u(2), Thus, during interactions with the MDP actions a, are sampled as a₃ ~ (s, 21), 2 p. (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 Putn (u)s).
consider
3 POLICY GRADIENTS INCORPORATING THE FUTURE
Our method aims allow a policy during training leverage future information for learning control. We propose to utilize stochastic latent variable models to this end. Namely, we propose to train me and Qs with latent variables (z, u) = {(z,, wt) T-1 sampled from learned function (o(T) which has access to the full trajectory. For example, the PGIF form the policy gradient objective in (2) may be expressed as
JYGIE na (To, Qu. = (2.)) (r)
r' Qu ae, u.) log To To 2x)
(4)
The PGIF form of the temporal difference objective in (3) may be expressed analogously as,
T-
TkaenoQQus go) = Er P (T)
Σ
(Q. Qu 0e, u)))
is clear that any RL. objective which trains policies and/or value functions on trajectories can be adapted PGIF form in straightforward manner. For example, in experiments we will apply PGIF to an LSTM-based PPO (Schulman et al., 2017), SAC (Haarnoja et al., 2018), and BRAC (Wu al., 2020).
While the PGlf-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 (T), since the full trajectory T is not yet observed. Therefore, while we want to give Tg and Q2 the ability to look 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 qo in terms of a KL divergence from prior distribution P.. (T) := {p. (z,, u ( which conditions (2,, u.) only on st. Thus, in the case of policy gradient, the full loss is,
J KL. (mo, : Q 4s) JeGI, (o, Qu, 40) + BT (Pu.(qs())llp(T)))
(6)
3

---

Published as a conference paper at ICLR 2022
feftib))
Robols.5)
Training
Inference
mylak.
Polisaltsc)
7A-1
polals))
21
Figure 1: The architecture of the model. Our inference model lae backwards hidden state b to approximate dependencies of 2 on the future of the trajectory. The blue line separates the data collection and policy gradient training steps our algorithm and the red lines represent information flowing into the backwards RNN. Grey variables are used during training and white variables are used during data collection. Top: we show the training model where the policy gradient loss is calculated with backwards RNN hidden state information. Bottom: we show the data collection phase of the algorithm utilizing latent variables sampled from the latent prior.
where 3 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 Po (s) when interacting with the environment.
Parameterization of 96 In our implementation, reverse order on т. Specifically, we use an LSTM network to process the states in backwards, to yield LSTM hidden states = {b, 1T-1 The function q& is then given by Gaussian distributions with mean and variance at time t derived from the backwards state b, In practice, avoid potential interfering gradients from the objectives of Ty and Qv. we use separate RNNs with independent parameters 4530, 9a(u) for žt, uy, respectively. See Figure for 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 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.
parameterize (T( RNN operating in
3.1 VARIATIONAL INFORMATION BOTTLENECK INTERPRETATION
The KL regularization we employ above may be interpreted as variational information bottleneck, constraining the mutual information between the latent variable distribution and the trajectory T. Here provide brief derivation establishing this equivalence.
For simplicity, we consider a specific timestep E N and a starting state 8, = s. Let Tz denote the random variable for all information contained after and including timestep t in trajectory induced by т. Let U, be the random variable for latents u induced by g (T2+) conditioned on all steps in the trajectory after and including t. Consider constrained objective minimizing JEGIETD while enforcing an upper bound Jmux on the mutual information between the distribution of trajectory steps and the distribution of latent variables I(T2,, Uels, = s). This objective is given by,
min JPGIETD (s = 8) == Er-p, ( I(Q) Qu az 41) 2.
s.t. T(T2e. 8) mmax
Recall the definition of mutual information:
p(T21, ue|s: = s) dundrzr
(T22, s) = p(Tzu ulst =s)log p( = s)p(un up(uul s)
(8)
/ Pro (Tgcle =
su p log (u T2t) duidT2n
=
p(uels: 5)
4

---

Published as a conference paper at ICLR 2022
where plulls s) is the marginal distribution over the latent variable p(ur|s; s) I qo( (ul(uTTze)) [2 (Tycls: = s)dT2r. This marginal is intractable to compute directly, and so to approx- imate this marginal we introduce a variational distribution h(u, (s = s). By definition we know that DKL. (allpuusss = 6) [/([u) = *)] 2 O. We can then see that f ppuell< |st = s) log p(ue = s)du 2 J pluy|s: = 8) log h(u/ = s)du. We therefore derive the upper bound for use in equation as,
I(T2;,44|s = 2 / (TT(1e = () log
t (o(u)(T2)) dundTze
h(ue St 8)
S SEr=p_gtte 11,, ( = 8))
We subsume the constraint into the objective as,
min JPGIET((Qu s0 + (E (1 [DLL (4s (Ts(T)))l((u(ss s))l Lmax P.
By taking h to be our learned prior Pvs we recover the single step (st = 8), TD analogue of the PGIF objective in (6), offset by constant Imax. as desired.
3.2 Z-FORCING WITH AUXILIARY LOSSES
While our proposed training architecture enables the policy and value function to look at the full trajectory т, in practice 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 force the latent variables to encode information about the future. We denote this loss as (C) 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 total loss function.
We emphasize that these auxiliary losses are applied exclusively to the backwards encoder (T( and have direct impact the learned policy or value functions. In fact, conduct ablations in Appendix 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ę (b. (2) over the backwards states given the inferred latent variables 2 qacz (z+ be). and similarly for the latents u4, We can write this auxiliary objective a maximum log-likelihood loss JAx(() = -Ba (z.(6,) [log Pc (belzz)). This way, we enforce the noisy mapping
be 24 defined by 4am to not be too noisy so as to completely remove any information from by-
Value prediction networks (VPN) A more sophisticated approach to force information to be propagated use an autoencoder-like model-based auxiliary loss. To this end, we take inspiration from VPNs (Oh et al., 2017), and apply auxiliary loss that uses b, to predict future rewards, values, and discounts. Note that, in principle, b, already has access to this information, by virtue of the backwards RNN transformer conditioned 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 the state based forcing described above, this approach only enforces by to contain the relevant information, and it is up to the RI. loss whether this information should be propagated to the latents 2t, My. We give detailed explanation of VPNs in Appendix B.
3.3 FULL ALGORITHM
The full learning objective for PGIF is thus composed of three components: First, a latent-varian variable augmented RL objective, e.g., policy gradient as shown in (4). Second, KI. regularizer, e-g., as shown in (6). Finally, 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.
5

---

Published as a conference paper at ICLR 2022
Algorithm 1 PGIF Algorithm with State-Action Value Function Estimation
Require: Initial parametens: 0, vn) vl o0t,, ,00, v. (re, ÇIN, Weights: org, OND, SPG, 9rv 1: for policy-step k = 0, ........ N do
2: Collect set of Trajectories D {:.....
repeat
4: 2: ~ PPpco (2)
5: Execute: To (a [t, 2) and observe re; 8++1 from environment.
6: until episode termination
7: for Trajectory: 7, D do
b² BackwardsLSTMM (T:) (Operates over the entire trajectory)
9:
b" U Backwardsl
10:
: 7. ub2 b
11:
V(s, r,,, } D:
12:
Derive (APG) JAL-T((M) according to any auxiliary loss.
13:
DTD = Dxi(qsm (u/b Ip, (u[s))
14:
JID = ErED [Jranm (Qw, Iawi) anDdAsm(GGi + 8Iv Div]
15:
DPG = (ra (z/b² (llb")ll,,,r)) ee (z|s))
16:
JIG = ETED [Jrau, Qu (sa) + ((rc) + Dro
17:
Update all parameters w.r.t: JPG and JID
4
RELATED WORK
We review relevant works in the literature in this section, with additional discussions 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 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 KI. 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 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-dimeniional representation, which is distilled to a non- hindsihhh 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 RI. loss to the representation loss (only the separate value prediction loss 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 the representation is through its dimension size and look-ahead, while using 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- on-policy RL, in these previous works. Nevertheless, it is 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).
found our algorithm performs significantly worse

---

Published as a conference paper at ICLR 2022
5 EXPERIMENTS
We now provide 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 SAC baseline, and numerous ablation analyses (Sec. F, 1) identifying the components of PGIF that are responsible for performance.
In our online RL experiments, 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 form of artificial memory of the past (but not forward-looking like in PGIF). SAC is a state-of-the-ar model-free 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 syper-parametes 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.
5.1 CREDIT ASSIGNMENT AND SPARSE REWARDS
PGIF aims to provide 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.
We first aim to show that our method is effective in simple environment where credit assignment is the paramount objective. We examine the Umbrella- Length task from BSUITE (Osband et al., 2020), task involving long sequential episode where only the first observation (a forecast of rain or shine) and action (whether to take an umbrella) matter, while the rest of the sequence contains random unrelated information. A reward of +1 is given at the end of the episode if the agent chooses correctly to take the umbrella or not depending on the forecast. This dif- ficult task is used to test the agent's s ability to assign credit correctly to the first decision. We evaluate PGIF-versions of SAC and PPO to vanilla discrete- SAC (Haarnoja et al., 2018; Christodoulou, 2019) as well as a PPO-LSTM (Schulman et al., 2017) baseline. Results are presented in Table 1. We see that our method is able to achieve best performance. This suggests that the PGIF agent is able to effi- ciently and accurately propogate information about the final reward 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.
Method
Mean Bsuite-score
PPO-LSTM
0.33 + 0.09
SAC
0.41 + 0.03
DON (3-step return)
0.44 + 0.06
DQN (I-step return)
0.51 = 0.03
PPO (VPN)
0.46 1 0.09
PGIF-PPO (Force)
0.26 1: 0.10
PGIF-SAC (VPN)
0.60 +.0+0.004
PGIF-SAC (Force)
0.58 + 0.08
PGIF-DQN (Force)
0.58 + 0.04
Table 1: Performance the Umbrella-Length environment. We run our model for 1000 episodes steps over random seeds. The BSuite score is calculated in terms of the regret normalized [random, optimal] -> [0, 1] (higher is better). The number after + is the standard deviation.
We continue the Gym-MiniGrid (Chevalier- Bedvvalie-BBoiveertt et al., 2018) set of curtially-observabll 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 only given local ego-centric view of its environment. For example, in DoorKey key must first be obtained and then 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
7

---

Published as a conference paper at ICLR 2022
SinpleCrosaingeeN
Un
LavaGap55
DcorRey SQF-PE0 PP04SITM
par PRO
HBM0-PP0
PRF-RRG RFOASTM
H8V0 PRO
Trieodo Donart
Edtocio Mean Roward
Eplests Mean Dwaard
Fpleott:: Moxrs Dowward
Plarnos
Frames
Figure 2: The online episodic mean reward evaluated over 5 episodes every 250 steps for MiniGrid RL tasks. We show the average over random seeds. 2M environment step interactions are used. The shaded area shows the standard error. PGIF always uses state based forcing in these environments.
22
2000
1000
2000
Vian
Veear
1000
2000
500G 20000 500-30 REOIE MOX0
-100
10003 20000 500000 A0000 30000)
10000 JU000 S00G SOXIO (000)))
Figure 3: The online episodic mean reward evaluated over 5 episodes every 500 steps for MuloCo continuous control RL tasks with partial observability. We show the average over random seeds. 500, 000 environment step interactions are used. The shaded area shows the standard error.
fails to converge to any rewards. LavaGap, there no significant improvement. We hypothesize that this is because the most common trajectories generated during early training 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 situation where our future representation is highly suboptimal, especially during early training.
5.2 PARTIAL OBSERVABILITY
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 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 random dimension of the state space each data collection step. This helps us test partial observability, feature common in robotic agent tasks.
We compare a Ptlle 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 than when only conditioning 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 than the more sophisticated VPN based forcing. See Appendix for online evaluations without partial observability.
consistently better
5.3 OFFLINE RL EVALUATIONS
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
8

---

Published as a conference paper at ICLR 2022
Critic (BRAC) (Wu et al., 2020) and Betch-Consttained Q-learning (BCQ) (Fujimoto et al., 2019). BRAC operates straightforward modification of SAC, penalizing the value function using 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 "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 contrast to the online results,
makes learning a dynamics representation difficult. Interestingly, we find that VPN based forcing performs better than state based forcing, although even state based forcing usually performs better than the baseline methods.
Environment
BRAC
PGIF (VPN)
PGIF (Force)
BCQ
ant-medium
2731 + 329
3250 = 125
2980 + 164
1851 = 94
ant-medium-expert
2483 + 329
3048 + 362
2431 + 417
2010 + 133
hopper-medium
1757 183
2327 + 399
1930 + 44
1722 4 166
walker2d medidum
3687 + 25
3989 = 259
3821 341
2653 = 301
salfiheectah-mediin
5462 198
6037 324
6231 + 303
4722 : 206
silichhctta-mmdduu---me
5580 + 105
5418 76
5491 143
4463 = 88
untmaze-umaze
0.5 0.16
0.95 + 0.0
0.1
0.8 + 0.13
Table 2:
Performance on the offline RL tasks showing the average episodic return. The final average return is shown after training the algorithm for 500, 000 episodes and then evaluating the policy over 5 episodes. Results show an average of 5 random seeds. The value after + shows the standard error.
6 DISCUSSION
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 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 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 RI, 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 known that expert datasets lack action diversity. Further conditioning force these biases to be incorporated more easily into the learning algorithm. Some other biases that may arise trajectory that contains some marginal reward and is therefore incorporated into our policy/value function with PGIR This could hamper exploration of the agent and prevent discovery of states that yield optimal reward. It may be interest combine method with exploration strategies. Perhaps there are benefits in decreasing access to future information during initial exploratory agent steps so the agent better able to explore. Furthermore, we find that our method is slower train than the baselines we compare to, due to the fact that the architecture requires training an LSTM, with times, very long trajectories as input. We must use the entire trajectories to train this LSTM. These issues can largely be alleviated by using transformer, with minimal difference in performance.
the future in this dataset could
REFERENCES
Combined reinforcement learning via abstract representations. Proceedings of the AAAI Conference on Artificial Intelligence, 2019,
Zaheer Abbas, Samuel Sokota, Erin Talvitie, and Martha White. Selective synalstyle planning under limited model capacity. In Proceedings of the 37th International Conference Machine Learning.

---

Published as a conference paper at ICLR 2022
Justin Bayer and Christian Osendorfer. Learning stochastic recurrent networks, 2015.
Samy Bengio, Oriol Vinyals, Navdeep Jaitly, and Noam Shazeer. Scheduled sampling for sequence prediction with recurrent neural networks. Proceedings of the 28th International Conference on Neural Information Processing Systems, 2015.
Samuel R. Bowman, Luke Vilnis, Oriol Vinyals, Andrew Dai, Rafal Jozefowicz. and Samy Bengio. Generating sentences from a continuous space. In Proceedings of The 20th SIGNLL Conference Computational Natural Language Learning, 2016.
Jacob Buckman, Danijar Hafner, George Tucker, Eugene Brevdo, and Honglak Lee. Sample- efficient reinforcement learning with stochastic ensemble value expansion. In Advances in Neural Information Processing Systems, 2018.
Keith Bush and Joelle Pineau. Manifold embeddings for model-based reinforcement learning under partial observability. Advances in Neural Information Processing Systems, 2009.
Xi Chen, Diederik P. Kingma, Tim Salimans, Yan Duan, Prafulla Dhariwal, John Schulman, Ilya Sutskever, and Pieter Abbeel. Variational lossy autoencoder, 2017.
Maxime Chevalier- for OpenAI Gym. https: //gth/ttb//u com/ maammmammeeeee ob/ /gwg//ggmmminnggggeggeiiggg minigrid. 2018.
Lucas Willems, and Suman Pal. Minimalistic gridworld environment
Petros Christodoulou. Soft actor-critic for discrete action settings. CoRR, abs11910002007,, 2019.
URL huttp://arxivy wwtpxiii.rrg..oooooo orgoo//g//aoooo /tbs/1910.. 07207.
Carlos Florensa, David Held, Markus Wulfmeier, Michael Zhang, and Pieter Abbeel. Reverse curriculum generation for reinforcement learning. Proceedings of Machine Learning Research, 2017.
Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4rl: Datasets for deep data-driven reinforcement learning, 2020.
Scott Fujimoto, Herke van Hoof, and David Meger. Addressing function approximation error in actor-critic methods. In Proceedings of the 35th International Conference Machine Learning, 2018.
Scott Fujimoto, David Meger, and Doina Precup. Off-policy deep reinforcement learning without exploration. In Proceedings of the 36th International Conference on Machine Learning. 2019.
Anirudh Goyal, Alessandro Sordoni, Marc-Alexandre Côté, Nan Rosemary Kc, and Yoshua Bengio. Z-forcing: Training stochastic recurrent networks. In Proceedings of the 31st International Conference Neural Information Processing Systems, 2017.
Anirudh Goyal, Alex Lamb, Jordan Hoffmann. Shagun Sodhani, Sergey Levine, Yoshua Bengio, and Bernhard Schölkopf. Recurrent independent mechanisms. In International Conference on Learning Representations, 2021.
Arthur Guez, Fabio Viola, Theophane Weber, Lars Buesing, Steven Kapturowski, Doina Precup. David Silver, and Nicolas Heess. Value-driven hindsight modelling. In Advances in Neural Information Processing Systems 33: Annual Conference, 2020.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. In Proceedings of the 35th International Conference Machine Learning, 2018.
Anna Harutyunyan, Will Dabney, Thomas Mesnard, Mohammad Gheshlaghi Azar, Bilal Piot, Nicolas Heess, Hado van Hasselt, Gregory Wayne, Satinder Singh, Doina Precup, and Remi Munos. Hindsight credit assignment. In Advances in Neural Information Processing Systems, 2019.
Michael Janner, Qiyang Li, and Sergey Levine. Reinforcement learning big sequence modeling problem. arXiv preprint 2eXiv:20060020:: 2021.
Philip N. Jchnson-Laird Mental models and human reasoning. Proceedings of the National Academy of Sciences, 2010.
10

---

Published as a conference paper at ICLR 2022
M. Karl, Maximilian Sölch, J. Bayer, and P. V. D. Smagt. Deep variational bayes filters: Unsupervised learning of state space models from raw data. ArXiv, htss/6605.04433 2017.
Nan Rosemary Ke, Amanpreet Singh, Ahmed Touati, Anirudh Goyal, Yoshua Bengio, Devi Parikh, and Dhruv Batra. Modeling the long term future in model-based reinforcement learning. In International Conference on Learning Representations, 2019.
Aviral Kumar, Justin Fu, Matthew Soh, George Tucker, and Sergey Levine. Stabilizing off-policy q-learning via bootstrapping error reduction. In Advances Neural Information Processing Systems, 2019.
Alex M Lamb, Anirudh Goyal, Ying Zhang, Saizheng Zhang, Aaron C Courville, and Yoshua Bengio. Professor forcing: A new algorithm for training recurrent networks. In Advances in Neural Information Processing Systems, 2016.
Thomas Mesnard, Theophane Weber, Fabio Viola, Shantanu Thakoor, Alaa Saade, Anna Harutyunyan, Will Dabney, Thomas S Stepleton, Nicolas Heess, Arthur Guez, Eric Moulines, Marcus Hutter, Lars Buesing, and Remi Munos. Counterfactual credit assignment in model-free reinforcement learning. In Proceedings of the 38th International Conference on Machine Learning, 2021.
Volodymyr Mnih, Adria Puigdomenech Badia, Mehdi Mirza, Alex Graves, Timothy Lillicrap, Tim Harley, David Silver, and Koray Kavukcuoglu. Asynchronous methods for deep reinforcement learning. In Proceedings of The 33rd International Conference on Machine Learning.
Ofir Nachum, Bo Dai, Ilya Kostrikov, Yinlam Chow, Lihong Li, and Dale Schuurmans. Algaedice: Policy gradient from arbitrary experience, 2019.
Chris Nota, Philip Thomas, and Bruno C. Da Silva. Posterior value functions: Hindsight baselines for policy gradient methods. In Proceedings of the 38th International Conference on Machine Learning, 2021.
Junhyuk Oh, Satinder Singh, and Honglak Lee, Value prediction network. In Advances in Neural Information Processing Systems, 2017.
Ian Osband, Yotam Doron, Matteo Hessel, John Aslanides, Eren Sezener, Andre Saraiva, Katrina McKinney, Tor Lattimore, Csaba Szepesvári, Satinder Singh, Benjamin Van Roy, Richard Sutton, David Silver, and Hado van Hasselt. Behaviour suite for reinforcement learning. In International Conference on Learning Representations, 2020.
Georg Ostrovski, Mare G. Bellemare, Aäron van den Oord, and Rémi Munos. Count-based exploration with neural density models. In Proceedings of the 34th International Conference on Machine Learning, 2017.
Razvan Pascanu, Yujia Li, Oriol Vinyals, Nicolas Heess, Lars Buesing, Sebastien Racanière, David Reichert, Théophane Weber, Daan Wierstra, and Peter Battaglia. Learning model-based planning from scratch, 2017.
Deepak Pathak, Pulkit Agrawal, Alexei A. Efros, and Trevor Darrell. Curiosity-driven exploration by self-supervised prediction. In Proceedings of the 34th International Conference on Machine Learning, 2017.
Baolin Peng, Xiujun Li, Jianfeng Gao, Jingjing Liu, and Kam-Fai Wong. Deep Dyna-Q: Integrating planning for task-completion dialogue policy learning. In Proceedings of the 56th Annual Meeting of the Association for Computational Linguistics), 2018
Sébastien Racanière, Theophane Weber, David Reichert, Lars Buesing, Arthur Guez, Danilo Jimenez Rezende, Adrià Puigdomènech Badia, Oriol Vinyals, Nicolas Heess, Yujia Li, Razvan Pas- canu, Peter Battaglia, Demis Hassabis, David Silver, and Daan Wierstra. Imagination-augmented augmented agents for deep reinforcement learning. In Advances in Neural Information Processing Systems, 2017.
Andrew Michael Saxe, Yamini Bansal, Joel Dapello, Madhu Advani, Artemy Kolchinsky, Bren- dan Daniel Tracey, and David Daniel Cox. On the information bottleneck theory of deep learning. In International Conference on Learning Representations, 2018
11

---

Published as a conference paper at ICLR 2022
Juergen Schmidhuber. Reinforcement learning upside down: Don't predict rewards just map them to actions, 2020.
John Schulman, Philipp Moritz, Sergey Levine, Michael Jordan, and Pieter Abbeel. High-dimensional continuous control using generalized advantage estimation. In Proceedings of the Fourth Interna- fional Conference of Learning Representations, 2016.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv, 1707.06347, 2017.
Satinder P. Singh. Learning without state-estimation in partially observable markovian decision processes. In In Proceedings of the Eleventh International Conference on Machine Learning. 1994.
Nobuo Suematsu and Akira Hayashi. A reinforcement learning algorithm in partially observable environments using short-term memory. In Advances Neural Information Processing Systems, 1999.
Richard S. Sutton. Learning to predict by the methods of temporal differences. Machine Learning. 1988.
Richard S. Sutton. Dyna, an integrated architecture for learning, planning, and reacting. SIGART Bull, 1991.
Richard S. Sutton and Andrew G. Barto. Reinforcement Learning: An Introduction. Bradford Book, 2018. ISBN 0262039249.
Richard S. Sutton, David McAllester, Satinder Singh, and Yishay Mansour. Policy gradient methods for reinforcement learning with function approximation. In Proceedings of the 12th International Conference on Neural Information Processing Systems, 1999.
Naftali Tishby, Fernando Pereira, and William Bialek. The information bottleneck method. 1999.
E. Todorov, T. Erez, and Y. Tassa. MuloCo: A physics engine for model-based control. In 2012 IEEE/RSJ International Conference on Intelligent Robots and Systems, 2012.
John N. Tsitsiklis and Benjamin Van Roy. Analysis of temporal-difference learning with function ap- proximation. In Proceedings of the 9th International Conference on Neural Information Processing Systems, 1996.
Hado van Hasselt, Yotam Doron, Florian Strub, Matteo Hessel, Nicolas Sonnerat, and Joseph Modayil. Deep reinforcement learning and the deadly triad, 2018.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in neural information processing systems, pp. 5998-6008, 2017.
Tingwu Wang, Xuchan Bao, Ignasi Clavera, Jerrick Hoang, Yeming Wen, Eric Langlois, Shunshi Zhang, Guodong Zhang, Pieter Abbeel, and Jimmy Ba. Benchmarking model-based reinforeement learning. 2020.
Lex Weaver and Nigel Tao. The optimal reward baseline for gradient- reinforcement learning. In Proceedings of the Seventeenth Conference on Uncertuinty in Artificial Intelligence, 2001.
Yifan Wu, George Tucker, and Ofir Nachum. Behavior regularized offline reinforcement learning, 2020. URL
net/ forumm th
Markus Wulfmeier, Dushyant Rao, Roland Hafner, Thomas Lampe, Abbas Abdolmaleki. Tim Hertweck, Michael Neunert, Dhruva Tirumala, Noah Siegel, Nicolas Heess, and Martin Riedmiller. Data-efficient hindsight off off-po option learning, 2021.
Mengjiao Yang and Ofir Nachum. Representation matters: Offline pretraining for sequential decision making, 2021.
12

---

Published as a conference paper at ICLR 2022
Chi Zhang, Sanmukh Rao Kuppannagari, and Viktor Prasanna. {BRAC}+: Going deeper with behavior regularized offline reinforcement learning, 2021. URL https:://ppeneeiie...rrt
Eorum?2d-bMCCPeppXX
Lukasz Kaiser, Mohammad Babaeizadeh, Piotr Milos, Blažej Osiński, Roy H Campbell, Konrad Czechowski, Dumitru Erhan, Chelsea Finn, Piotr Kozakowski, Sergey Levine, Afroz Mohiuddin, Ryan Sepassi, George Tucker, and Henryk Michalewski. Model based reinforcement learning for atari. In International Conference on Learning Representations, 2020.
13


Reviews:
[{"review_id": "8U-Yd0YvrTt", "reviewer": "Reviewer_EPSM", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "# Method\n- I see why the proposed method would help in partially observed environments: conditioning on some representation of the future allows the policy to condition on a function of the hidden information and give a lower variance policy gradient. However, in fully observed online settings it is not obvious to me why this method improves performance, and I don't see much intuition provided in the paper. \n- Equations 4 and 5: why are $u$ and $z$ separate? Did you try learning a single latent variable for both policy and $Q$ function? \n- This latent variable $u$ seems to contain information about both the dynamics and logging policy since these interact to form the future. If we know the logging policy, the information bottleneck means we artificially throw away information that we know. This seems like an undesired side-effect of this information bottleneck. \n- 3.2: \"Especially in settings with highly stochastic learning signals\". What does this mean?\n- I didn't follow how $p_{v^{(u)}}$ was updated. It is described as a \"learned prior\" but it only appears in the KL term in the loss. If the network for learning this has sufficiently many parameters will it not just overfit to tracking $q_{\\phi}$? \n\n# Experiments\n- It seems that Table 1 might address partial observability. If the first state and action are not maintained in the state for the rest of the task (if I have misunderstood this please clarify), we do not have an MDP. If this is true, I don't see a convincing case for this method improving performance much outside of partially observable environments (in online RL) besides the minor improvement in appendix E, and the great results in appendix H. If the focus of improvement is in partially observable or offline settings, this should be made clear in the motivation. \n- In Appendix H, I did not find the explanation for improvement convincing, even though the results themselves are impressive. \"We hypothesize our algorithm has benefits in these environments since as soon as it obtains a reward signal it can adapt quickly and make use of the signal by incorporating it in both policy and value optimization, therefore accelerating learning.\" With enough Bellman updates, SAC should also be able to propagate this information back to states earlier in the trajectory. How many rounds of Bellman updates are being used in SAC? Is it possible that in fully observed environments, PGIF is giving an optimization benefit by skipping this reward information directly to early-trajectory states instead of indirectly through Bellman updates? \n- Figure 2: main text should highlight LavaGapS5 as a case where we don't see improvement. Is there some intuition as to why?\n- Section 5.1 is supposed to evaluate the method on sparse rewards, testing credit assignment. But a partially observed environment is used for some of the experiments in this section. Partial observability is the focus of 5.2.\n- Table 1 shows great results when using PGIF-SAC, but then in the next experiments in Figure 2 the PPO version is used? \n- Why is HIMO not compared to in Table 1? Same for Figure 3. It is explained in the appendix but could perhaps be a footnote in the main paper. \n- 5.2: \"more realistic scenario\" is probably not an accurate way to describe this particular modification to the environment. \n- Experiments in appendix E. The environment has no partial observability but we are in a finite horizon setting, meaning that the number of remaining steps in an episode is not known to the agent. It could be that the slight improvement in performance over SAC is because PGIF can provide the policy with the number of remaining steps. This can be tested by including the number of remaining episode steps into the state and rerunning the experiment. \n- How does the method perform with no Z-forcing? This seems like a straightforward ablation.\n- Mesnard et al. (2021) is a very similar approach that is not compared to.  \nWhy is the improvement in offline continuous control much greater than the online equivalent? Is it because in the offline case the future contains information about the logging policy? \n\n# Minor points\n- The proposed method is called PPO-PGIF in the text but in Figure 2 the legend uses PGIF-PPO. \n- Appendix B: \" Attention augment agents utilize dynamics\" Typo\n- \"PGIF agent is able to efficiently\" typo\n- Row 1 of Table 7 in the appendix: typo?\n- The text on figures (eg legends) is too small. \n\nI'm not too familiar with the related work on hindsight, so my review takes the paper's comparison with this prior work at face value. ", "summary_of_the_paper": "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. ", "summary_of_the_review": "Here I review the key constructive points:\n- Beyond the online ant setting in the appendix, I don't see much evidence that this method improves performance outside of partially observable environments in the online setting. The intuition behind why it would help in fully observed environments is also not clear from my reading of the paper. Given this, I think that the paper should focus the motivation more heavily on the POMDP and offline setting. This was the main reason for my score. I'm happy to discuss if I've misunderstood something about the utility of the method in the fully observed setting. \n- From my description above I think there are a few ways in which the experiments could be more convincing. \n\nWhy I liked the paper and recommended to accept:\nThe core idea in this paper is very elegant and the variational inference + z forcing to implement this idea is well done. As far as I can tell the use of Z-forcing and the variational bottleneck is novel in this setting. The results on partially observable and offline settings are impressive.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "4RVH0nxUh4a", "reviewer": "Reviewer_9hJX", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Overall, this is a solid paper, that combines an interesting and novel theoretical idea and convincing theoretical results. The clarity of the technical content is a particular strength, and the authors' clear prose helps to communicate non-trivial theoretical concepts and provides intuition for their design decisions. There are still a few places in which I believe the clarity of the paper could be enhanced that would serve to make the paper both easier to follow and the results more impactful.\n\nThe main change I would suggest is to rely less heavily on Appendices to communicate conclusions about the various other activities. For example, while the reference to Appendix.F in Sec. 3.2 is quite informative (and mentions that applying PGIF to the policy and value functions degrades performance), for other content, the reader *must* go to the appendices to  understand the conclusions of supplementary experiments: e.g., appendix G in which Transformers are used instead of an RNN; there should be couple sentences in the main text summarizing those results. While I do not think it is an issue that there is so much content in the appendices---and, in fact, it may be a benefit so that more time is devoted to the clear descriptions of the technical content of the paper---mentions of the appendices containing results (including both Appendix G and E) should be expanded on.\n\nSome other smaller questions and suggestions are as follows:\n- [Sec 3.2] It would be worth mentioning that the Force and VPN objectives are separate from one another. It was only by the time that section was over that it was clear that only one would be used at a time.\n- [Sec 5.1] Some readers may not be as familiar with the Gym-MiniGrid environments. Mention that the agent is given only given local ego-centric view of its environment, which is what gives rise to the partial-observability.\n- [Sec 5.2] While not essential, it would be helpful to understand why Fig. 2 includes results only from the State-based forcing of PGIF. Would the results from VPN look quite different? Its omission is especially odd considering that the VPN results are clearly better for the experiments from Table 2.\n- [Sec 6] I appreciate the potential risks highlighted in the Discussion section. It might be helpful to include how much slower \"much slower to train\" implies. Could the authors include an order-of-magnitude number for the difference in wall-clock training time for one of the experimental environments?\n", "summary_of_the_paper": "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.", "summary_of_the_review": "This seems to be a solid paper with an interesting and well-executed novel theoretical contribution tackling a difficult problem of relevance to many in the RL and planning-under-uncertainty communities. The experiments are also convincing, and compare against popular RL strategies. I believe that the clarity of the paper could be further enhanced, yet point out that the paper is already well-written and easy to follow.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "3bkMIvqxi0m", "reviewer": "Reviewer_GXWR", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths\n------------------------------------\n- The clarity was mostly good.\n- Experimental results appeared good (in terms of improving over the baselines, but I am not sure the baselines were well-implemented).\n- The method is general and could be applied together with most RL algorithms.\n- The method improved performance for both online and offline RL algorithms.\n\nWeaknesses\n------------------------------------\n- It wasn't clear to me why the method should improve the performance.\n- The experiments only looked at reward scores without examining any other statistic or property of the learning.\n- The computational time is much larger compared to other methods (they said it's much slower, but unfortunately did not say exactly by how much).\n- It was not clear to me that the baseline algorithms were solid, as it seems that the experimental results are new, and not reproductions of previous results in other publications.\n- Some ablation studies are missing (in particular, adding a prior network to the models essentially increases the capacity of the Q and policy models, so I thought it should be checked whether the baseline algorithms may also benefit from using larger models).\n\nSuggested decision\n------------------------------------\nCurrently, I am leaning toward recommending to reject, as I cannot recommend accepting a paper unless I am sure the empirical results are solid. Potentially, this can be cleared up in the rebuttal.\n\nSupporting arguments\n------------------------------------\n**Experiments**\n\nI can't think of any good theoretical reason why one would want to do what the authors did in their paper, so I think that the experimental results are crucial to convey usefulness and validity of their approach. Currently, I was not convinced that the experiments are solid (see below), and I also think that reporting only reward scores is quite a limited method of experimentation as it does not explain why the performance may have improved. One of the clearest ways to show results that can be trusted is to reproduce existing results from another publication, then outperform those results, but it didn't seem like this was done.\n\nThe Bsuite umbrella-length experiment seems fine. The result of ~0.6 is higher than the result of ~0.5 for DQN reported in the bsuite paper. However, the new baselines in this paper (PPO and SAC) only had a score of ~0.4. Also, why were none of the other bsuite environments tested?\n\nIn minigrid there is also no reference to any previous publication showing the same results as produced in this paper.\n\nThe experiments with partial observability were based on the work of Yang & Nachum, (2021), but it appears executed in a different way, and it's not easy to see that the baselines reproduced any results from their work.\n\nThe offline RL experimental results also are not clear to me. For example, on halfcheetah your baselines achieve around 4.5k score, but in the BRAC paper, the algorithms achieve around 6k score. Where does the discrepancy come from?\n\nAlso, in the antmaze results in Table 10, your baselines achieve ~0, whereas the in the original paper\nby Florensa et al. the baseline algorithms reached around 0.6 reward.\n\nIt is quite possible that the experiments are performed solidly; however, this is not explained sufficiently for me to understand this. I would appreciate if the authors can comment on the correctness of their implementations, and provide evidence that the results achieved by their baseline algorithms are competitive (e.g., by providing references to papers where the scores are similar to the baselines implemented in this paper).\n\n**Concept of algorithm**\n\nThe algorithm was motivated from the point of view of allowing to \"look into the future\" to improve the optimization performance; however, typical policy gradient algorithms already do this when using Monte Carlo returns. Instead of using your approach, one can give perfect information of the future by just using the empirical return. Then what is the motivation of the method, and why should it help with learning?  One could think of some possible reasons such as variance reduction (perhaps it is doing something similar to variance elimination https://arxiv.org/abs/1811.06225 or http://proceedings.mlr.press/v100/cheng20a.html), or better representation learning; however, such points are not explored.\n\n**Ablation studies**\n\nThe prior network $p(z|s)$ is essentially increasing the capacity of the policy network. I think it would be better to also compare to a version where the backward RNNs are not trained, but the $p(z|s)$ is just trained end-to-end as if it were a part of the policy with some additional latent noise $z$. It would also be interesting to know what happens if you increase the number of parameters of the critic (to account for the greater model capacity achieved by using the backward RNN).\n\nOther notes\n------------------------------------\n\nIs the prior network also trained using the policy gradient loss (the action depends on the sampled $z$ from the prior, so it could be trained end-to-end)?\n\nThe discussion around equations 1 and 2 is not wrong, but\nI think it would be better to explain that this is a surrogate loss\nthat allows easily implementing the policy gradient using\nautomatic differentiation, and that only the $\\log \\pi$ term will\nbe differentiated. Originally, the policy gradient theorem was not\npresented in the way that you did.\n\nOn page 5 in the VPN section, what do you mean by predicting the discounts?\nIsn't the discount factor just a constant? Also, is the encoding network\ntrained end-to-end?\n\n\"Our method is more versatile that previous works,\"\n-> \"...versatile than...\"\n\nThe figure axis labels are small and hard to read. Also, you used a\nnon-vector figure format, which makes them blurry. I prefer vector\nformat figures so that they look sharp when you zoom in. Moreover,\nit would be better to use more than color to differentiate between\nthe lines.\n\n\"Furthermore, we find that our method is much slower to train than the\nbaselines we compare to, due to the fact that the architecture\nrequires training an LSTM, with at times, very long trajectories as\ninput.\"\nIt is good that you mentioned this, but can you roughly quantify it?\nHow much slower was it?\n\nTable 7 top right, are you missing one digit in the PGIF score?", "summary_of_the_paper": "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.\n(Note that this z-forcing is done for both $z$ and $u$).\n\nThey 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.", "summary_of_the_review": "The experimental results look impressive, but I was not convinced that the baselines were well-implemented. The experimentation also only looked at reward scores, so it is not clear why the performance improved. It was also not clear to me that there is any good theoretical reason why the method should improve the performance. Currently, I am recommending to reject. I hope the authors will provide strong arguments and evidence in their rebuttal.\n\n---------------------------------------------------------------------------------------------------------------------\nUpdate after author rebuttal\n----------------------------------------\nThe authors have provided additional evidence on the correctness of their baselines.\nThey also performed additional ablations, such as checking the change in performance when increasing\nthe model capacity. I was satisfied with these changes, so I have increased the correctness to 3, and the score the 6.\nMy recommendation is a borderline accept, as I was satisfied that the method empirically improves the performance; however,\nI think the authors could have done more to examine why their method improves the performance.\n\nI still have some concerns with the paper:\n- The experiments still only look at the reward curves. They could have explored points such as the prediction accuracy of the value function, the variance and accuracy of the gradients, or tried to do some experiment to examine the quality of the representations (although this one is more difficult, as I don't think there is a standard way to look into this).\n- I am still not convinced by the authors' explanation of why \"looking into the future\" is useful. When the information bottleneck is removed, the value function would ideally just predict the Monte Carlo return. So the authors claim appears to be that \"looking into the future\" is good, but if it is done too much, then it stops being useful. I think this point could have been examined in more detail, e.g. by doing a study on how the information bottleneck coefficient affects the performance (admittedly, it would have been better for me to say this in my initial review). Previous works (e.g., https://arxiv.org/abs/2011.09464) doing related ideas provided longer more principled explanations of why their methods would be useful by arguing that it might reduce the policy gradient variance by removing the random fluctuations in the future.\n\nSome more minor concerns:\n- I still think the discount prediction requires more clarification in the appendix B.\n- Regarding the discussion around equations 1 and 2. I think you still\nneed to give more clarification. If one differentiates this \"loss\"\ntwice, one does not obtain the correct Hessian of the RL objective\n(e.g., https://arxiv.org/abs/1802.05098).  I think this section should\nbe explained as an automatic differentiation trick to obtain first\norder policy gradients.\n", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "-y6fbClIz2E", "reviewer": "Reviewer_UbCM", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n1) The paper attempts to address the problem of credit assignment in RL which is very important. The authors build off of the existing idea of learning future conditioned policy/Q function and introduce several tricks (informational bottleneck and z-forcing) to make it work.\n2) The paper demonstrates that the proposed approach can be applied to several RL algorithms, such as PPO, SAC, and BRAC.\n3) The paper features an extensive empirical study across several tasks that feature various challenge, such as sparse reward credit assignment, partial observability, and learning from offline datasets.\n\nWeaknesses:\n1) The novelty of the method is limited as it largely builds off of existing work from Harutyunyan et al, 2019 and Mesnard et al., 2021. The the authors attempt to differentiate their work from the prior methods, but the difference mostly comes from the fact that PGIF can be made off-policy, unlike the prior work (i.e. Hindsight Credit Assignment) which is on-policy.\n2) It is not clear how much information about the future one can force into the prior distribution that is only conditioned on the current state.  I'm not convinced that this does any better than for example SAC (or DDPG/TD3) + n-step returns. In general, the paper omits comparison to any methods that uses n-step return.\n3) The authors feature z-forcing as an advantage of their method, but I believe that the need for z-forcing stems from the fact that there is very little mutual information between the prior and posterior distributions, so the backward model without z-forcing will simply collapse to  a trivial solution in order to minimize KL-div to the prior distribution.\n4) The final method ends up being extremely complicated featuring several auxiliary losses and the recurrent model. The empirical results are encouraging but not ground breaking. Thus I'm not sure the algorithm is practical and can enjoy wide adoption considering its complexity.\n\nQuestions:\n1) Re the umbrella-length experiment: what episode length has been actually used? or this presents an average results across multiple horizons?\n2) Re the umbrella-length experiment: it would be interesting to see a baseline that uses n-step returns, for example DDPG/TD3 + nstep. It would be much stronger baseline than SAC in these sparse reward/long credit assignment settings.\n3) Re the MuJoCo partial-observability experiments: I don't think it is a fair comparison to compare SAC-based PGIF against vanilla SAC (that doesn't use recurrent policy/Q). Obviously PGIF, that has a way to deal with partial observability, will outperform SAC. Instead, the authors compare to PPO-LSTM, which can deal with partial observability, but not as strong baseline as SAC on MuJoCo. It would be interesting to also see performance of SAC-LSTM and PPO-based PGIF to make this experiment valid. ", "summary_of_the_paper": "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.", "summary_of_the_review": "While the paper studies an interesting direction of learning future conditioned policy/Q-function, which is quite exciting, I have several issues with the paper, namely:\n1) The incremental nature of the work.\n2) The complexity of the method.\n3) Several concerns regarding empirical evaluation.\n\nTaking these into account, I'm not sure the paper contributes enough in its current state to justify an acceptance. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper presents a new reinforcement learning algorithm for POMDPs that specifically deals with the credit assignment problem. The proposed algorithm consists in using at each time-step t of a training trajectory the subsequent future trajectory that starts at time t+1 as additional inputs to the policy and value networks. Instead of using the trajectories directly, two RNNs are used to encode the trajectories into latent two variables that are then given as inputs to the policy and value networks. A key novel contribution of this work is the use of "Z-forcing" to help the RNNs learn the relevant information. Since future trajectories are not available during testing, a "prior" network is trained to predict the latent variable given a state. During testing, the latent variable is sampled from the network. Empirical experiments on simple simulated environments show that the proposed algorithm outperforms several baselines.

Key issues raised by the reviewers include the complexity of the proposed algorithm, the fact that several interesting results are in the appendix rather than the main paper, and the weakness of certain baselines. The authors responses helped clarify these issues, and additional experiments (such as a comparison to a DQN with n-step value updates) were performed and added to the paper. The reviews are updated accordingly.

In summary, the paper contains several novel ideas in the context of learning in partially observable environments. It is not entirely clear similar effects of the proposed algorithm can be obtained by using simpler tricks, but the evidence provided by the authors supports the claim that the algorithm outperforms several SOTA techniques in the context of POMDPS.
Decision: Accept (Poster)