Long-Term Credit Assignment via Model-based Temporal Shortcuts
Abstract
This work explores the question of long-term credit assignment in reinforcement learning.
Assigning credit over long distances has historically been difﬁcult in both reinforcement learning and recurrent neural networks, where discounting or gradient truncation respectively are often necessary for feasibility, but limit the model’s ability to reason over longer time scales.
We propose LVGTS, a novel model-based algorithm that bridges the gap between the two ﬁelds.
By using backpropagation through a latent model and temporal shortcuts to directly propagate gradients, LVGTS assigns credit from the future to the possibly distant past regardless of the use of discounting or gradient truncation.
We show, on simple but carefully-designed problems, that our approach is able to perform effective credit assignment even in the presence of distractions.
Introduction
In reinforcement learning (RL) [33], an agent executes actions in an environment to maximize future rewards.
An arbitrarily long period of time can occur from the execution of a particular sequence of actions to the reception of the corresponding reward; the problem of understanding the relationship between actions and rewards regardless of this delay is called temporal credit assignment [34].
Effective assignment of credit over long time spans is an issue for most current RL methods and is exacerbated in realistic settings with partially observable environment states.
Zooming out, temporal credit assignment is also a fundamental problem is sequence-based supervised learning tasks.
Recurrent Neural Networks (RNNs) [10, 29] are often the go-to solution to solve them but, in a similar way to basic RL algorithms, vanilla RNNs struggle with long-term temporal dependencies [4, 14, 26].
The root causes of this phenomena have been thoroughly investigated and numerous intuitive solutions have been designed over the years to more accurately capture these temporal dynamics, with considerable performance gains [4, 19, 21, 26, 35].
In this paper we propose to build on the connection between temporal credit assignment in RNNs and RL tasks, by leveraging ideas developed for the ﬁrst as a natural solution to design algorithms for the second.
In particular, we exploit the similarity between the computational graph in which
Backpropagation Through Time (BPTT) operates for computing the gradients for RNNs and the one implied by the interactions of an agent with its environment.
As core techniques in RNNs allow them to effectively propagate gradients in a more direct way to the distant past, their RL counterpart will allow agents to assign credit to actions even when they yield rewards in a non-immediate future.
We propose Latent Value Gradients with Temporal Shortcuts (LVGTS), a novel model-based policy optimization algorithm, which learns a latent model of the dynamics and uses it, together with reparameterization techniques, for the analytical computation of the policy gradient.
During dynamics learning, we impose a speciﬁc attentive structure which allows the model to create shortcuts between potentially distant states and actions; then, during policy optimization, the gradient is computed by backpropagating through real trajectories of experience using the learned model, implicitly injecting nuanced credit assignment patterns into policy learning.
LVGTS is designed by starting from a Partially Observable Markov Decision Process formulation; in that context, we highlight that a blind application of an attention mechanisms to model-based policy optimization does not necessarily result in improved credit assignment capabilities.
We then showcase, in small but focused experiments, that LVGTS is able to deal with delayed rewards even in the presence of distractions.
To the best of our knowledge, this is the ﬁrst model-based algorithm that performs effective long-term credit assignment without the need for heuristics or reward shaping, being simultaneously fully compatible with a discounted reinforcement learning formulation.
Background
Problem Deﬁnition.
A continuous ﬁnite Partially Observable Markov Decision Process (POMDP) is a tuple M =
( S , O , A , f, O, r, µ, T ) , where S is the state space, O ⊆ R x is the observation space,
A ⊆ R u is the action space, f : S ×
A → ∆( S ) is the state transition function, O : S → ∆( O ) is the observation function, r : S ×
A → R is the reward function, µ ∈ ∆( S ) is the initial state distribution and T is the problem horizon.
A sequence τ t =
( o 0 , a 0 , . . .
, s t , a t ) of t observations and actions can be seen as taken from a space T t of trajectories .
At each discrete time step, the agent interacts with the environment by receiving an observation and acting according to a policy π θ : T t → ∆( A ) , which potentially depends on the entire history, and which comes from a space Π Θ = { π θ : θ ∈ Θ ⊂ R d } of parameterized policies.
To dissect the fundamental features of credit assignment in POMDPs, we are interested in solving so-called type 1 and type 2 information acquisition tasks as introduced in [20].
Type 2 tasks test the agent’s capability to learn adequate state representations in a partially observable setting where past observations, however distant, are necessary for future reward predictions.
The notion of passive memory is thus crucial to acting optimally in such an environment: to reach the optimal policy, the agent must be able remember a speciﬁc observation o i .
Conversely, type 1 tasks test an agent’s ability to simultaneously learn adequate memory-based state representations and perform long-term temporal credit assignment.
In this setting, the memory state is not passively observed, but needs to be sought out by the agent who is only rewarded for it in the future when it needs to be recalled.
Gradient-based Policy Optimization.
The goal of the agent is to maximize the expected cumulative reward of trajectories induced by its policy:
[Equation]
where V π θ is a state-value function, which gives for every state the expected cumulative reward of policy π θ .
To maximize this metric, we are interested in estimating ∇ θ V .
Under non-restrictive smoothness assumptions the reparameterization trick [13, 22] allows us to rewrite the above expectation, which concerns the distribution induced by policy and environment, with respect to another random variable, for instance (cid:15) ∼ N (0 , 1) , to compute the gradient by direct differentiation of a transformed function.
In other words, ﬁnding ∇ θ V becomes a simple matter of estimating ∇ θ r t .
If the transition and reward functions of the environment are known and differentiable, then the gradient of r t can be found by backpropagating along the trajectory.
The resulting gradient is estimated not unlike the traditional backpropagation through time (BPTT) mechanism in recurrent neural networks.
The methods, either model-based or value-based, which build upon this techniques for policy optimization are referred to as value gradients [11, 18].
Discounting and Truncation.
When dealing with large horizons, the practice of discounting the rewards during the computation of a value function has been found to be practically beneﬁcial even in episodic reinforcement learning, where it would not be necessary on a theoretical level.
By constraining a problem to an effective horizon of 11 − γ , the discount factor has been seen to reduce the variance of model-free RL.
While the advantages of discounting is not as well understood in model-based learning, it is still employed in almost all prior works [16, 18, 32], and can intuitively control the bias of propagating a learned model through time.
Similar to discounting, RNNs often employ truncated backpropagation through time (TBPTT) for long sequences.
Again, by effectively reducing the horizon of the backwards pass, some of the major issues in long sequences can be
[Figure]
mitigated.
In both cases, long-term credit assignment becomes difﬁcult or impossible beyond the effective horizon implied by TBPPT and discounting in vanilla RNNs and standard RL algorithms.
Credit Assignment with Model-based Temporal Shortcuts
Value gradient methods lie at the intersection of RNN training and reinforcement learning.
Given the popularity of discounting in RL and gradient truncation in RNNs, and their direct mathematical connection (see Appendix), it seems natural to propose an algorithm compatible with these methods, but at the same time featuring long-term credit assignment capabilities.
In this section, we present a value gradient method which directly exploits this connection, by employing gradient truncation together with a strategy to circumvent the myopia it implies.
Our Latent
Value Gradients with Temporal Shortcuts algorithm takes advantage of two main building blocks: reparameterization-based stochastic value gradients (SVG)
[18] and attention-memory augmented
RNNs [21].
Value gradients are used for policy optimization, and an attention-memory augmented
RNN is used for state representation in POMDPs.
Given its novel architecture, LVGTS has the capability to leverage the rich structure learned during model learning, by directly incorporating it into credit assignment.
We ﬁrst introduce a naive implementation of SVG with the specialized RNN, called latent-memory stochastic value gradients (LM-SVG).
This will serve both as an introduction and as a baseline for LVGTS.
3.1 A Value Gradients Baseline
In the rest of the paper, we assume for clarity of presentation deterministic dynamics, but then use reparameterization in our implementation to accommodate for stochasticity.
When value gradients are applied to the fully observable setting, the policy is directly improved given an (approximate) transition model f : S ×
A → S and reward model r : S ×
A → R .
In the partially observable setting, the true underlying states are not seen, and must be instead approximated through a latent (or belief)
state h t :=
b ( o 1 ..t a 1 ..t − 1 )
≈ s t , where b ( · ) is some function approximator.
It is also beneﬁcial to introduce a decoder d ( h t )
≈ o t , to be employed as a regularizer.
To summarize, the models learned for latent planning are therefore:
[Equation]
Transition model
[Equation]
Reward
[Equation]
Given such models, we can train a policy π θ ( h t ) → a t following the value gradient methodology outlined in the previous section.
LM-SVG implements this architecture by employing an encoder based on an attention-augmented recurrent network architecture.
This form of attention-augmented recurrent network is similar to the one introduced in Sparse Attentive Backtracking [21].
The parameterization then becomes:
[Equation]
where att ( h i<t ) uses a dot-product key-value based attention mechanism [35] to select k latent states from the past, summarize them, and passes it along the encoder b to better predict h t .
A more detailed description of this attention mechanism can also be found in Algorithm 2.
Despite being perhaps the most natural architecture for memory-based latent value gradients, this approach presents two important drawbacks:
The resulting latent space is non-markovian.
The transition model used for value gradients assume that the given states are Markovian, such that p ( h t | h
i<t )
=
p ( h t | h t − 1 ) .
This is not the case given (3), which may make the transition model difﬁcult or impossible to learn.
Credit assignment remains unchanged during policy optimization.
Truncating the gradients will still result in a myopic agent, no matter the accuracy of the model employed.
Long-term dependencies that are captured by the dynamics through repeated transitions are cut short with gradient truncation, preventing long-term reasoning.
The memory in this case is only used for representation learning, and does not directly address the problem of credit assignment in RL.
In the following, we will present an algorithm able too overcome both of these shortcomings.
3.2 Latent Value Gradients with Temporal Shortcuts
We can solve the two problems of LM-SVG by implementing a simple change: LVGTS includes h tm as an argument to the transition model f .
The approximate state transitions can then be rewritten with ˆ h tm which recycles the attention mechanism learned in (3), but collects the approximate states predicted by the transition model instead:
[Equation]
The advantages of this formulation are two-fold, conveniently addressing the two problems of
LM-SVG:
The resulting latent space is Markovian .
Since h t − 1 and h tm are the only past states used for the encoder b to predict h t , then the combination of the previous state and memory state provides all the necessary past information for the transition model to predict h t .
Credit assignment is improved during policy optimization.
The addition of ˆ h tm into the argument of f acts as a temporal shortcut in the computational graph of r i ≥ t , allowing credit to directly ﬂow among distant time steps as seen in Figure 1.
Let us further develop the intuition behind the second point through the idea of temporal shortcuts.
Assume that ˆ h tm is comprised of a single state at time step m < t .
The attention mechanism itself has no bias towards selecting states that are as close to t as possible, so m can be arbitrarily far from t as long as it was deemed important in predicting h t .
Traditionally, BPTT is only able to reason of the effects of m on t by ﬁrst reasoning about m + 1 , then m + 2 , then m
+ 3 until ﬁnally t is reached.
Similar to a game of broken telephone, the true effects of m on t are often difﬁcult to discern because of this phenomenon, whether it is due to noise or model approximations.
The situation is made even worst when employing TBPTT, where the model has no way of understanding the effects of m on t if t − m > N .
If a temporal shortcut exists between m and t , as is the case in LVGTS, a single step of reasoning separates h t and h tm instead of t − m , allowing for a more direct path for credit to be assigned to distant events.
Implemented with TBPTT, this allows the agent to reason about past important events, and their neighboring states, just as Figure 1 suggests.
The complete algorithm for LVGTS, including the reparameterization trick for stochastic dynamics, is detailed in Algorithm 1.
In what follows, we will see that LVGTS is capable of solving both typeand type 2 information acquisition tasks even with gradient truncation, while LM-SVG fails in the former.
Experiments
We test LVGTS and LM-SVG on a few simple type 1 and type 2 information acquisition tasks.
The goal of these experiments are to verify the algorithms’ capabilities to perform long-term credit assignment beyond what gradient truncation permits.
First, we formalize the two types of tasks into two respective concrete environments.
4.1 Environments
Input: Episode length T , truncation length N , and models: [Table] , w t | o [Table]
, a
[Table] [Table] − 1 ) , d ( o t | h t ) , f ( h t +1 | h
[Table]
, h m , a
[Table]
[Table]
[Table]
[Table]
|
[Table]
, a t ) , π θ ( a
[Table]
|
[Table] [Table] while not converged do Observe initial observation o 1 ; for t ← 1 to T do h t ← b ( o 1 ..t , a 1 ..t − 1 ) from Algorithm 2; a t ← π θ ( h t , η t ) , η t ∼ ρ ( η ) ; Take action a t and observe o t +1 , r t from the environment ;
Insert ( o t , a t , r t , o t +1 ) into D ; end
Train b, d, f, r with (7) using D ;
Initialize latent state ˆ h 1 ← b ( o 1 a 0 ) ;
Initialize memory: M ∈ R T ×| h | , M [1] ← ˆ h 1 ;
Initialize return: G ← 0 ; for t ← 1 to T do
Infer ξ, w t | o t − 1 , a t − 1 , o t following Algorithm 2; h m ← (cid:80) w t [ i ] ∈ w t w t
[ i ] M [ i ] ; ˆ a t ← π
[Equation] (ˆ h t , η t ) , ˆ , ˆ a ; ˆ a
Passive Memory Task (Type 2 Information Acquisition).
The passive memory task (PMT) is a simpliﬁcation of the grid world environments in [20, 36], generalised to the continuous state-action space.
Given two constant time stamps t 1 and t 2 , the task is separated into three phases: remember, distractor, recall.
The remember phase ( t < t 1 ): A random observation o m is generated for the agent to see.
Actions in this phase do not affect future observations.
In the simple case, this phase can be characterized by the observation generation process: o m ∼ N (0 , 1) .
The distractor phase ( t 1 < t < t 2 ): Any desired continuous task can be inserted here.
It serves to distract the agent with an unrelated task, temporally distancing the ﬁrst and second phase.
The recall phase ( t > t 2 ): Random unrelated observations are shown here, but a reward is given to the agent for adequately remembering the randomly generated observation o m seen during the remember phase.
For some maximum reward c , the reward can be deﬁned as: r t>t 2 = ( a t − o m ) [-
-]2 + c .
Active Memory Task (Type 1 Information
Acquisition).
The active memory task (AMT) uses the same three phases deﬁned above for
PMT, with a slight modiﬁcation to the remember phase.
Instead of always observing the correct memory observation, o m is only fully observed when performing the correct action.
Given this action to be a m , deﬁned in the environment, the observations in the remember phase for AMT are o t<t {+
+}
= o m +
( a m − a t ) (cid:15) , where (cid:15) ∼ N (0 , σ )
We consider three different instances of these two tasks: dummy-PMT, dummy-AMT, and pendulumAMT.
The preﬁx speciﬁes what environment is used for the distractor phase.
The dummy environment always returns o t = 0 , r t = 0 for any t 1 < t < t 2 .
The pendulum environment is the Pendulum-v environment for the OpenAI Gym library [6].
4.2 Results
LM-SVG and LVGTS are tested on all three instances of PMT and AMT with t 1 = 3 , t 2 = 10 , and T = 15 .
In all cases, c = 0 .
5 , therefore the maximum possible return for each episode is 2.
Both algorithms are subjected to a truncation length of N = 4 , which means that long-term credit assignment from the recall phase to the remember phase is impossible without skip connections.
We set k = 3 for all experiments, allowing for a maximum of three skip connections to be made to the past at every time step.
The results in Figure 2 show the top ﬁfteen best performing runs out of twenty
[Figure]
n
[Figure]
While both algorithms are able to solve the passive memory task, indicating a competency towards representation learning, only LVGTS is able to reliably solve both AMT instances, even in the presence of a distraction.
Conversely, truncated LM-SVG struggles to surpass the performance of an ideal myopic agent even with minimal distractions (dummy-AMT).
When distrations are present, in the form of the pendulum task, LM-SVG completely fails under truncation.
The maximum myopic return is calculated based on a policy that acts perfectly in phases 2 and 3 according to the noisy observation in phase 1, but acts completely randomly in the ﬁrst phase.
Additionally, the bottom row of Figure 2 shows the agent forming the correct shortcuts to the past for credit to ﬂow through in the third phase.
These results demonstrate that even though we truncate the gradients of LVGTS, a model-based algorithm, it is still able to perform long-term credit assignment in a POMDP well beyond its effective horizon without the need for heuristics.
Related Work
Long-Term Credit Assignment in RNNs.
A fundamental problem faced while training recurrent neural networks is the assignment of credit across timesteps.
The goal of facilitating this task inspired the development of nuanced architectures and training algorithms, which go well-beyond the primordial forms of RNNs, based on the simple ﬂow of a latent state through time [10].
Indeed, vanilla RNNs are susceptible to exploding and vanishing gradients during backpropagation [19, 26], which hinder the ability to model dynamics spanning over long time windows.
To counterbalance this, attention [4] has been recently proposed as a particularly successful paradigm: instead of being constrained to only look at a single, most recent, state, the model is augmented with the capability of attending time frames from the past.
Methods such as transformers [35] and SAB [21], which we use as a backbone for our algorithm, extensively leverage attention mechanisms.
Value Gradients.
The intuitive approach of improving a control policy by differentiating through the dynamics and reward function was one of the ﬁrst to be conceived in modern RL [30, 37].
Based on this idea, the resulting family of algorithmic tools encompasses both direct differentiation through the learned dynamics [1, 8, 11] and the use of the gradient from a learned value function [12, 15, 23, 31].
Value gradients have been recently revamped [2, 7, 18] due to their synergy with automatic differentiation, the workhorse of deep learning [5], which allows automatic computation of the policy gradient from the computational graph constructed by repeatedly evaluating dynamics, policy and reward.
It is possiible to extend value gradients to partially observable and complex environments, by using latent models [16] and computing gradients through the latent space.
Overcoming delayed rewards in RL.
Delayed rewards pose a great challenge to long-term credit assignment in reinforcement learning [3, 20].
Traditional use of the discount factor is especially problematic in these settings, as it is a direct contradiction to the reality of delayed rewards in many real world problems [24, 27, 36].
Previous work in addressing this issue has mostly been modelfree, and includes methods around meta-learning [28, 38], reward reshaping [3, 28], re-weighting discounting factors [38], and hindsight reasoning [17, 25].
Of particular relevance is the work on
Temporal Value Transport (TVT) [20].
TVT is a model-free approach that shares one important similarity with our work: they use an attention-based architecture for latent state representation.
However, there are two major differences with LVGTS.
Our algorithm is model-based, and directly reuses the attention weights for policy optimization, not requiring any reward shaping heuristic.
We take full advantage of automatic differentiation and let this process naturally handle the value transport that is manually implemented in TVT.
Conclusion
In this paper we presented a model-based method for long-term credit assignment, based on performing backpropagation through time in latent space.
By allowing the agent to construct temporal shortcuts during model learning, credit can be assigned to distant events during policy optimization, while taking full advantage of truncated backpropagation through time.
We show, through a simple set of experiments that LVGTS is able solve both type 2 and type 1 information acquisition tasks.
Value gradient algorithms have historically been shown to have difﬁculty scaling to larger more complex environments, and can beneﬁt greatly from using an approximate value function [2, 18].
A natural followup to our work is to scale it to more difﬁcult problems, and to see how critics and value functions might translate into this framework
A Appendix A.1 Truncated Value Gradients and Discounting
We formalize here the relationship of truncated value gradients with discounting.
For the remainder of this section, we suppose a deterministic environment with Markov states s t .
The truncated value function can be written in the following form, where the complete proof can be found in [18]:
[Equation]
We can rewrite the recurrent relationship iteratively to better understand how it compares to the original value function:
[Equation]
The original un-truncated and undiscounted value function being:
[Equation]
and an n-step value function being:
[Equation]
We can now appreciate the equivalence between equations 5 and 6, demonstrating that the ntruncated value gradient is equivalent to a full value gradient using an n-step value function .
Given that in traditional actor-critic methods [Equation] the critic ˆ V i ( s i ) is only used derivative estimation [9], the n-truncated value gradient is not so different than policy gradient algorithms with γ = 1 − 1 n , where future rewards are discounted smoothly instead of discretely
[Equation]
A.2 Training the Latent Models
Given a dataset of transition observations sampled from the environment ( o 1 ..t a 1 ..t − 1 , a t , o t +1 , r t ) the environment model approximations are trained on the following loss function:
[Equation]
The total loss sums all the above terms, and can be controlled by the constants α rew , α decoder , α trans ,
L models =
α rew L rew + α decoder L decoder + α trans L trans .
The decoder model acts as a regularizer, and provides a training signal beyond simple reward prediction for the latent model.
A.3 Attention Augmented RNN
The attention augmented RNN that uses key-value dot product attention, inspired by the work in [21] is detailed here.
We assume a sparse recall to memory with k memory recalls, where k is a hyper-parameter set during training.
Given a new observation o t a t − 1 , and the previous latent state h t − 1 the procedure returns the next latent state h t , and the attention weights attn k used for predicting h t .
Algorithm
[Table]
Input :
History
[Equation]
Sparsity k
Attention Key-gen
[Equation]
Multilayer Perceptron
[Equation]
K ← o
[Equation]
a
[Equation]
W ; ˆ h t ← MLP ψ ( o t a t − 1 ) ; attn w ← ˆ h t K T ; attn k ← sorted ( attn w )
{+[ k + 1] ; attn k ← ReLU ( attn w − attn k ) ; h m ← (cid:80) ti =
+}
attn k
[ i ] o i a i
[Equation]
1 (cid:80) i attn k [ i ] ; return ˆ h t + h m , attn k