Skill or Luck? Return Decomposition via Advantage Functions
Authors: Hsiao-Ru Pan
Bernhard Schölkopf
Venue: ICLR 2024 poster (2024)
Keywords: reinforcement learning
off-policy
advantage function
Primary area: reinforcement learning

Abstract: Learning from off-policy data is essential for sample-efficient reinforcement learning. In the present work, we build on the insight that the advantage function can be understood as the causal effect of an action on the return, and show that this allows us to decompose the return of a trajectory into parts caused by the agent’s actions (skill) and parts outside of the agent’s control (luck). Furthermore, this decomposition enables us to naturally extend Direct Advantage Estimation (DAE) to off-policy settings (Off-policy DAE). The resulting method can learn
from off-policy trajectories without relying on importance sampling techniques or truncating off-policy actions. We draw connections between Off-policy DAE and previous methods to demonstrate how it can speed up learning and when the proposed off-policy corrections are important. Finally, we use the MinAtar environments to illustrate how ignoring off-policy corrections can lead to suboptimal policy optimization performance.

Published as a conference paper at ICLR 2024
SKILL OR LUCK? RETURN DECOMPOSITION VIA AD- VANTAGE FUNCTIONS
Hsiao-Ru Pan & Bernhard Schökopf
Max Planck Institute for Intelligent Systems, Tübingen, Germany
ABSTRACT
Learning from off-policy data is essential for sample-efficient reinforcement learn- ing. In the present work, we build on the insight that the advantage function can be understood as the causal effect of an action on the return, and show that this allows us decompose the return of trajectory into parts caused by the agent's 's actions (skill) and parts outside of the agent's ' control (luck). Furthermore. this decomposition enables us to naturally extend Direct Advantage Estimation (DAH) to off-policy settings (Off-policy DAE). The resulting method can learn from off-policy trajectories without relying on importance sampling techniques truncating off-policy actions. We draw connections between Off-policy DAE and previous methods to demonstrate how can speed up learning and when the proposed off- off-policy corrections important. Finally, ronments to illustrate how ignoring off- off-policy corrections can lead to suboptimal policy optimization performance.
the MinAtar envi-
INTRODUCTION
Imagine the following scenario: One day, and B both decide to purchase a lottery ticket, hoping to win the grand prize. Each of them chose their favorite set of numbers, but only A got lucky and won the million-dollar prize. In this story, we are likely to say that A got lucky because, while A's of numbers) led to the reward, the expected rewards are the same for both A and B (assuming the lottery is fair), and A was ultimately rewarded due to something outside of their control.
action (picking
This shows that, in a decision decision-making problem, the return is not always determined solely by the actions of the agent, but also by the randomness of the environment. Therefore, for an agent to correctly distribute credit among its actions, it is crucial that the agent is able to reason about the effect of its actions on the rewards and disentangle it from factors outside its control. This is also known the problem of credit assignment (Minsky) 1961). While attributing luck to the drawing process in the lottery example may be easy, it becomes much more complex in sequential settings, where multiple actions involved and rewards are delayed.
The key observation of the present work is that we can treat the randomness of the environment as actions from an imaginary agent, whose actions determine the future of the decision-male making agent. Combining this with the idea that the advantage function can be understood as the causal effect of an action the return (Pan et al. 2022), we show that the return can be decomposed into parts caused by the agent (skill) and parts that are outside the agent's control (luck). Furthermore, we show that this decomposition admits natural way to extend Direct Advantage Estimation (DAE), an on- policy multi-step learning method, to off-policy settings (Off-policy DAE). The resulting method makes minimal assumptions about the behavior policy and shows strong empirical performance.
Our contributions can be summarized as follows:
We generalize DAE to off-policy settings.
We demonstrate that (Off-policy) DAE can be seen as generalizations of Monte-Carlo (MC) methods that utilize sample trajectories more efficiently.
We verify empirically the importance of the proposed off-policy corrections through exper- iments in both deterministic and stochastic environments.

---

Published as a conference paper at ICLR 2024
2 BACKGROUND
In this work, we consider a discounted Markov Decision Process A,,,, with finite state space S, finite action space A. transition probability p(s'|s, a), expected reward function S x A + RR and discount factor E [0, 1). A policy is a function π S + ((A) which maps states to distributions over the action space. The goal of reinforcement learning (RL) is to find policy that maximizes the expected return, x = arg max E, [G], where G
Et=n 7'r and r. = r(st, at). Ex [o the Q-function of state-action pair is similarly defined by Q (²( a) = Ex [ -0-8, a0=a (Sutton et al. 1998). These functions quantify the expected return of a given state or state-action pair by following given policy æ. and are useful for policy improvements. They are typically unknown and are learned via interactions with the environment.
The value function of state is defined by V (8)
Direct Advantage Estimation The advantage function, defined by (F another function that is useful to policy optimization. Recently, Pan et al., (2022) showed that the advantage function can be understood as the causal effect of an action on the return, and is more stable under policy variations (under mild assumptions) compared the Q-function. They argued that it might be an easier target to learn when used with function approximation, and proposed Direct Advantage Estimation (DAE), which estimates the advantage function directly by
no
27
{ Σ 1 ((.s)))((a))) DEA
Σ 7 (.)--A)))
AT = arg min E AEF,
Fr
where A. = A(s,, a.). The method can also be scamlessly combined with a bootstrapping target to perform multi-step learning by iteratively minimizing the constrained squared error
27
n-1 Σ Mrr - At) + 7" " Viarget - (50)
L(Ā IT
subject to Å E Fr,
(2)
where Vurget is the bootstrapping target, and (V, A) are estimates of the value function and the ad- vantage function. Policy optimization results were reported to improve upon generalized advantage estimation (Schulman et al., (2015b), a strong baseline for on-policy methods. One major drawback of DAE, however, is that can only estimate the advantage function for on-policy data (note that the expectation and the constraints share the same policy). This limits the range of applications of DAE to on-policy scenarios, which tend to be less sample efficient.
Multi-step learning In RL, we often update estimates the value functions based on previous estimates (e.g., TD(0), SARSA (Sutton al., 1998)). These methods, however, can suffer from excessive bias when the previous estimates difter significantly from the true value functions, and was shown that such bias can greatly impact the performance when used with function approxima- tors (Schulman et al. 2015b). One remedy is to extend the backup length, that is, instead of using one-step targets such as r(Sg. 0o) + Qurges ($1,01) a1) (Quuget being our previous estimate), we include more rewards along the trajectory, i.e., r(50, an) +-r((11 a1) ++)+++(22222 a2) +... +.. Quarrge (Sn; )- This way, we can diminish the impact of Qtarget by the discount factor y However, using the re- wards along the trajectory relies on the assumption that the samples are on-policy (i.c., the behavior policy is the same as the target policy). To extend such methods to off-policy settings often requires techniques such as importance sampling (Munos al., 2016 Rowland et al. 2020) or truncating (diminishing) off- off-policy actions (Precup et al.] 2000: Watkins 1989), which can suffer from high variance or low data utilization with long backup lengths. Surprisingly, empirical results have shown that ignoring off-policy corrections can still lead to substantial speed-ups and is widely adapted in modern deep RL algorithms (Hernandez- Garcia & Sutton 2019 Hessel et al. 2018: Gruslys et al.
2017).
3 RETURN DECOMPOSITION
From the lottery example in Section I we observe that, stochasticity of the return can come from two sources, namely, (1) the stochastic policy employed by the agent (picking numbers), and (2)
2

---

Published as a conference paper at ICLR 2024
the stochastic transitions of the environment (lottery drawing). To separate their effeet, begin by studying deterministic environments where the only source of stochasticity comes from the agent's s policy. Afterward, we demonstrate why DAE fails when transitions are stochastic, and introduce simple fix which generalizes DAE to off-policy settings.
3.1 THE DETERMINISTIC CASE
First, for deterministic environments, we have 8t+11 h(s),a)))) where the transition probability replaced by a deterministic transition function h S A S. As consequence, the Q-function becomes (st, ar) = r(5t, oz) + (V-(s)))), and the advantage function becomes (st, az) = r(st, (4) (st+1) - V (st). Let's start by examining the sum of the advantage function along a given trajectory (Su, ao $1, a1,...) with return G.
20
20
o
(s,
Σ ar)
Σ r (s--)) - G
VT (sa):
(3)
telescupingg
or, with a simple rearrangement, G = (sp) + Et=o (st, One intuitive interpretation of this equation is: The return of a trajectory equal to the average return (V*) plus the variations caused by the actions along the trajectory (A*). Since Equation 3 holds for any trajectory, the following equation holds for any policy μ
27 =0 O.
<--
G- Σ Y'A -VA-V (50))
8
E
(4)
t=o
This means that (V*, A*) is solution to the off-policy variant of DAE
)
27
L(A =
Σ 7(rr A.) - V(00)
s.t.
= =)
(5)
aEd
where the expectation now taken with respect to an arbitrary behavior policy l instead of the target policy и in the constraint (Equation with n + co). We emphasize that this is a very general result, as we made no assumptions on the behavior policy μu, and only sample trajectories from μu are required to compute the squared error. However, two questions remain: (1) Is the solution unique? (2) Does this hold for stochastic environments? We shall answer these questions the next section.
3.2 THE STOCHASTIC CASE
The major difficulty in applying the above argument to stochastic environments is that the telescoping sum (Equation 3) no longer holds because As)s,, = r(s;, a.) + ylR.g/p((-eeee [V*(8 (V" (5*())) az] (st) * r(st,a)) + (8++1) - V" (s:) and the sum of the
advantage function becomes
7 (s
Σ
7 (r(s), at) + (V" [V* (*(())
(6)
t=
1=0
= G
Σ
BT V*(),,
where B (st,,, 0t, 81+1) = (*(++)) -B ae) V* 1s,, a.). This shows that V and A are not enough to fully characterize the return G (compared to Equation and B is required. But what exactly is DT? understand the meaning of D", we begin by dissecting state transitions into two-step process, see Figure] [ In this view, we introduce
imaginary agent nature, also interacting with the environment, whose actions determine the next states of the decision-makin agent. In this setting, nature follows stationary policy म equal to the transition probability, i.e., i (s'](s, a)) = p(s'|e a). Since मे is fixed, we omit it in the following discussion. The question are interested in is, how much do nature's actions affect the return? We note that, while there
3

---

Published as a conference paper at ICLR 2024
Step 1: Agent picks an action az (-188)
Step 2: Nature picks a state 81+1 ~ p(11 a+)
St. dg
8t.
S1++1
A two-step view of the state transition process. First, we introduce an imaginary agent nature, which controls the stochastic part of the transition process. In this view, nature lives world with state space S = S x A and action space A = S. At each time step t, the agent chooses its action a based on 8t, and, instead of transitioning directly into the next state, it transitions into an intermediate state denoted (st, ax) E S, where nature chooses the next state St+1 E A based on (8:, az). We use nodes and arrows to represent states and actions by the agent (red) and nature (blue).
Figure 1:
are no immediate rewards associated with nature's actions, they can still influence future rewards by choosing whether we transition into high-rewarding states or otherwise. Since the advantage function was shown to characterize the causal effect of actions on the return, we now examine nature's 's advantage function.
By definition, the advantage function is equal to Q* a) - (s).. We first compute both Q" and V from nature's 's point of view (we use the bar notation to differentiate between nature's view and the agent's 's view). Since S S x A and A = S, V is now a function of = S A, and Q is a function of S A = S A x S, taking the form
(" a) E
r r. s0= a0=a] = E, (p |V =7 a0==]
(8)
(T(s,,,
2 rz
-0-ass==s11 Vs)).
(9)
Q' a, exactly Bd as, as introduced at the beginning of this section. Now, if we rearrange Equation 6 into
We thus have Az (8 as,
a) = V= (s Ept (V* (s')ss,, a). which is
00
(s0) + r' (A ac) + + at 81+1)) G,
(10)
1=0
then an intuitive interpretation emerges, which reads: The return of a trajectory can be decomposed
into the average return V* (so), the causal effect of the agent's 'ss actions A ("t,,, ax) (skill), and the causal effect of nature's 's actions 3 (st, a, St-1) (luck).
Equation 10 has several interesting applications. For example, the policy improvement lemma Kakade & Langford, 2002). which relates value functions of different policies by VP(() = (*(() + E,, E1>0 7 Ap S0 = s), is immediate consequence of taking the conditional expec-
tation Ey [=|[===]] of Equation 10 More importantly, this equation admits natural generalization
of DAE to off-policy settings:
Theorem 1 (Off-policy DAE). Given a behavior policy μ. targel policy v, and buckup length
O. Let A, A(se, at), B, B(st, des St41), and the objective function
27
L(A,
E 1
A -BB) +
x rr A.
( ((nt)))) vV
(II)
E AlEdA(s, -0 VsES Eves B(8,a, 8')(s'ss..))=== ((x
subject to
4

---

Published as a conference paper at ICLR 2024
then (A",B", V is minimizer of the above problem. Furthermore, the minimizer is unique if μ
sufficiently explorative (i.e., non-zero probability of reaching all possible transitions (s, a, s')).
See Appendix A for a proof. In practice, we can minimize the empirical variant of Equation I1
from samples estimate (V*, , B²), which renders this an off- off-policy multi-step method. We
highlight two major differences between this method and other off off-pooicy multi-step methods. (1) Minimal assumptions on the behavior policy are made, and no knowledge of the behavior policy is required during training (in contrast to importance sampling methods).
(2) It makes use of the full trajectory instead of truncating or diminishing future steps when off-policy actions are encountered [Watkins] 1989 Precup al. 2000 We note, however, that applying this method in practice can be non-trivial due the constraint Eves [ss, a, s' (')pps''ss Is, a) 0. This constraint is equivalent the À constraint DAE, in the sense that they both ensure the functions satisfy the centering property of the advantage function (i.e., Ba-- (A (s, = 0). Below, briefly discuss how deal with this.
Approximating the constraint As first step, we note that similar constraint Eaea A(s, ())(s = 0 can be
Pa(z|s, a)
Z
enforced through the following parametrization Åo(s, a) fo(s,a) - EasA So(s, (a(s), where Se is the underlying unconstrained function approximator (Wang al., 2016b). Unfortunately, this technique cannot be applied directly to the B constraint, because (I) it requires sum over the state space, which is typically too large, and (2) the transition function p(s'ls, a) is usually unknown.
E
SxA
Pa(s (1,51)8,,,2)
Figure 2: Latent variable model of transitions; Z is discrete latent space, which can be understood as actions from nature.
To overcome these difficulties,
a Conditional Varia- tional Auto-Encoder (CVAE) (Kingma & Welling 2013 Sohn
al. 2015) encode transitions into diserete latent space Z such that the sum can be efficiently approximated, see Figure 2 The CVAE consists of three components: (1) an approximated conditional posterior 45(2|1, s') (encoder), (2) conditional likelihood Pg(s'|s, a, 2) (decoder), and (3) a conditional prior P5(2|: ss, a). These components then be learned jointly by maximizing the conditional evi- dence lower bound (ELBO),
ELBO DRL. (ag (2)8, a, ('llg((ssss a)) + smag(11s,,ss))
E
[logp(6''ss a, =)].
(12)
Once a CVAE is learned, we can construct ((s, a, s') from an unconstrained function go a, 2) by B(s, s') Bx (=s a, 2))s, a, s'| Ez=pa(1x,)) 190 z)]s, a). which has the prop- erty that Σ. p(s B(s, 8') es 0 because 43 (2|8, 8') Ro p3(=1s, 8').
RELATIONSHIP TO OTHER METHODS
In this section, first demonstrate that (Off-policy) DAE can be understood generalization of MC methods with better utilization of trajectories. Secondly, we show that the widely used uncorrected estimator can be seen special case of Off-policy DAE and shed light when it might work.
4.1 MONTE-CARLO METHODS
To understand how DAE can speed up learning, let us first revisit Monte-Carlo (MC) methods through the lens of regression. In a typical linear regression problem, we are given dataset {(x,, y:) E RR x R], and tasked to find coefficients w E R" minimizing the error E, (w - I; - y: In RL, the dataset often consists of transitions sequences of transitions (as in multi-step methods) and their returns, that is, (T,, G.) where T, has the form ($o. op, $1, a a1....) and G; is the return associ- ated with 7, However, 7 may be abstract object which cannot be used directly for regression, and
5

---

Published as a conference paper at ICLR 2024
0.6
V(1)
p(50=1) = 009
1
r=0
Tu T==1
0.4
TD(0)
MC
V(2)
0.6
DAE
2(00-2)=00 0.1
2
r=0
ra=0
0.4
25
75
Trajectories
Figure 3: Left: An MDP with = (1,2,34), Both states and have only single action with immediate rewards 0 that leads to state 3. State 3 has two actions, and d, that lead to the terminal state 4 with immediate rewards I and respectively. Right: We compare the values estimated by Batch TD(0), MC, and DAE with trajectories sampled from the uniform policy. Lines and shadings represent the average and one standard deviation of the estimated values over 1000 random seeds.
The dashed line represents the true value ((1)= V(2) 0.5. See Appendix B for details.
we must first map 7 to a feature vector ((T) E R For example, in MC methods, we can estimate the value of a state by rolling out trajectories using the target policy starting from the given state and averaging the corresponding returns, i.e., E 7'relso s) Σ k. This is equivalent to linear regression problem, where we first map trajectories vector by (T( [(so = s) (vector of length [S] with elements if the starting state is s or 0 otherwise), and minimize the squared error
I(v) = Σ Σ (ada((t)
(13)
where v is the vector of linear regression coefficients V/2. Similarly, we can construct feature maps such as (T) = [[s0=s,, (g0=s,aaa))) and solve the regression problem to arrive at Q" (s, a). This view shows that MC methods can be seen linear regression problems with different feature maps. Furthermore, it shows that MC methods utilize rather little information from given trajec- tories (only the starting state(-action)). An interesting question is whether it is possible construct features that include more information about the trajectory while retaining the usefulness of the co- efficients. Indeed, DAE (Equation 2 with
utilizing two different feature maps (T) Ei= 7T,(, aa) and (T) = T(sp=)) which results in a vector of size [S] [A|] that counts the multiplicity of each state-action pair in the trajectory and vector of size (S) indicating the starting state. This suggests that DAE can be understood as a generalization of MC methods by using more informative features.
xc) can be
To see how using more informative features can enhance MC methods, let consider an example (see Figure adapted from Szepevvári (2010). This toy example demonstrates a major drawback of MC methods: it does not utilize the relationship between states 2 and 3, and therefore, an accurate estimate of (3) does not improve the estimate of V (2). TD methods, utilize this relationship to achieve better estimates. DAE, similar to TD methods, also utilizes the relationship between (2) and Á(3, -) to achieve faster convergence on ((2).. In fact, in this case, DAE converges faster than TD(O) it exploit the sampling policy efficiently estimate Á(3,-), whereas TD(0) has rely on sample means estimate V(3).
the other hand, can
Similarly, we can compare DAE to Off-policy DAE, which further utilizes Φx.a, (T) >'(x, a,=a, 81+1 =')), ), in stochastic environments. See Figure 4 for another example. Here, we observe that both Off-policy DAE variants can outperform DAE even in the on-policy set- ting. This is because Off-policy DAE can utilize (4, across different trajectories to account for the variance caused by the stochastic transitions at state 4.
'This is not to be confused with the features of states, which are commonly used approximate value functions.

---

Published as a conference paper at ICLR 2024
17(1)
1.0
1
0=
Yo
0.5
DAE OfIDAE
1.5
OfiDAE(oracle)
V(2)
1.0
2
r=
To To=0
6
r=
0.5
0
25
50
75
Trajectories
Figure 4: Left: An MDP extended from Figure 3 Instead of terminating at state the agent transitions randomly to state 5 or 6 with equal probabilities. Both states 5 and 6 have a single action, with rewards and 0, respectively. State 7 is the terminal state. Right: We compare the values (with uniform policy) estimated by DAE, Off-policy DAE (learned transition probabilities), and Off- policy DAE (oracle, known transition probabílities.. Lines and shadings represent the average and one standard deviation of the estimated values over 1000 random seeds. The dashed line represents the true value V(1) ((2) 1.
4.2 THE UNCORRECTED METHOD
The uncorrected method (simply "Uncorrected" in the following) updates its value estimates using the multi-step target == 7'r, + - (mrger(Sn+11 without any off-policy correction. Hernandez- achieve performance competitive with true off-policy methods deep RL, although it was also noted that its performance may be problem- specific. Here, examine how Off-policy DAE, DAE, and Uncorrected relate each other, and give possible explanation for when Uncorrected can be used.
Garcia & Sutton 2019) showed that Uncorrected
We first rewrite the objective of Off-policy DAE (Equation 11 into the following form:
(V(r))
Σ
7'r, + a1 V Vargen (sn-1)
A'
( B.)
(14)
Uncorrected
DAE
where the underbraces indicate the updating targets of the respective method. We can see now there clear hierarchy between these methods, where DAE special case Off-policy DAE by assuming B and Uncorrected is special case by assuming both À 0 and B 0.
The question is, then, when is A = or B a good assumption? Remember that, deterministic environments, we have 8x meaning that DAE can be directly applied to off-polic data when the environment is deterministic. Next, see when Á = is useful, remember that the advantage function can be interpreted as the causal effect of an action on the return. other words, if actions in the environment tend to have minuscule impacts the return, then Uncorrected work with a carefully chosen backup length. This can partially explain why Uncorrected worked environments like Atari games (Bellemare) al. 2013; Gruslys et al. 2017; Hessel al. 2018 for small backup lengths, because the actions are line-grained and have small impact (AA s 0) in general. In Appendix [C we provide a concrete example demonstrating how ignoring the correction can lead biased results.
for any policy æ; therefore, B O is a correct estimate of B
EXPERIMENTS
We now compare (1) Uncorrected, (2) DAE, (3) Off DAE, and (4) Tree Backup (Precup al. 2000) in terms policy optimization performance using simple off-policy actor-critic algorithm. By comparing (1), (2), and (3), we test the importance of Å and & as discussed in Section 4.2 Method (4) serves baseline of true off-policy method, and Tree Backup was chosen because,

---

Published as a conference paper at ICLR 2024
Deterministic Environments
Stochastic Environments
3.0
Uncorrected
Uncorrected
2.5
Normalized Score
2.5
DAE
Normalized Score
DAE
OFF-potic DAE
2.0
OFF-policy DAE
Tree
Tree
1.5
1.5
1.0
1.0
0.5
0.5
0.0
0.0
10
10
Frames (Million)
Frames (Million)
Figure 5: Normalized training curves aggregated over deterministic (left) and stochastic (right) en- vironments. Scores were first normalized using the PPO-DAE baseline and then aggregated over 20 random seeds, environments, and backup lengths. Lines and shadings represent the means and standard error of the means, respectively. The dotted horizontal lines shows the PPO-DAE baseline.
like Off-policy DAE, it also assumes no knowledge of the behavior policy, in contrast to importance sampling methods. We compare these methods in a controlled setting, by only changing the critic objective with all other hyperparameters fixed.
Environment We perform our experiments using the MinAtar suite Young & Tian 2019). The
MinAtar suite consists of 5 environments that replicate the dynamics of subset of environments from the Arcade Learning Environment (ALE) (Bellemare et al.) 2013) with simplified state/action spaces. The MinAtar environments have several properties that desirable for study: (1) Actions tend to have significant consequences due to the coarse discretization of its state/action spaces. This suggests that ignoring other actions' effects (A), as done in Uncorrected, may have larger impact on its performance. (2) The MinAtar suite includes both deterministic and stochastic environments, which allows us to probe the importance of B.
Agent Design We summarize the agent in Al- gorithm [ Since (Off-policy) DAE's loss func- tion depends heavily on the target policy, we found that having smoothly changing target pol- icy during training critical, especially when the backup length is long. Preliminary experi- ments indicated that using the greedy policy, i.e., arg max A(s, a), as the target policy can lead to divergence, which is likely due to the phe- nomenon of policy churning (Schaul et al. 2022). To mitigate this, distill policy by maximiz- ing Ea [Â(s, a)), and smooth it with exponen- tial moving average (EMA). The smoothed policy TEMA then used as the target policy. Addition- ally, to avoid premature convergence, we include KL-divergence penalty between me and TEMA- similar to trust-region methods (Schulman et al. 2015ap. For critic training, we also use an EMA of past value functions as the bootstrapping tar- get. For Off-policy DAE, we additionally learn a CVAE model of the environment. Since learning the dynamics of the environment may improve sample efficiency by learning better representations
Algorithm 1 A Simple Actor-Critic Algorithm
Require: backup Initialize Ag. Vo. Bg. Tg
2: Initialize CVAE 93. Pō
3: Initialize D t- (). θEMA 0
4: for = 1,
do
Sample (s, a, a, r') T
6: D + D U {(8 a, r
Sample batch B trajectories from D
8:
if backup is Off- DAE then
Train CVAE (Eq [2] using B
10: Approximate Ba (')))
11:
end
12:
Compute Leritie (Eq.
14
13: Compute Lactor
= EuNno [A] + 3KL 1 (T
14:
Train Tucrttcc Tractor using B
15:
OEMA +- TθEMA + (1 +)
16:
end for
(Gelada et al. 2019: Schwarzer et al. 2020 Hafner et al.) 2020), we isolate this effect by training a separate network for the CVAE such that the agent can only query p(2) s, a, s') and p(2[s, a). See Appendix D for more details about the algorithm and hyperparameters.
8

---

Published as a conference paper at ICLR 2024
Results Each agent is trained for 10 million frames, and evaluated by averaging the undiscounted scores of 100 episodes obtained by the trained policy. For comparison, we use the scores reported on-policy baseline, which were trained using PPO and DAE (denoted
by Pan et al. (022) PPO-DAE).. The results are summarized in Figure Additional results for individual environments and other ablation studies can be found in Appendix [D] We make the following observations: (1) For deterministic environments, both DAE variants performed similarly, demonstrating that ß is ir- relevant. Additionally, both DAE variants converged to similar scores as Tree backup, albeit slightly slower, suggesting that they can compete with true aff-po methods. Uncorrected, on the other hand, performed significantly worse than DAE, suggesting that Á is crucial in off- off-poolcy settings, as the two methods only differ in Å. (2) For stochastic environments, we see a clear hierarchy between Uncorrected, DAE and Off-policy DAE, suggesting that both Á and / corrections are important. Notably, Tree backup performs significantly worse than both DAE variants in this case, while only being slightly better than Uncorrected.
RELATED WORK
Advantage Function The advantage function was originally proposed by Baird (1994) to address small time-step domains. Later, it was shown that the advantage function be used to relate value functions of different policies (Kakade & Langford 2002) or reduce the variance of policy
gradient methods (Greensmith al. 2004). These properties led to wide adoption of the advantage
function in modern policy optimization methods (Schulman et al. 201 b 2017] Mnih et al., 2016). More recently, the connection between causal effects and the advantage function was pointed out by Corcoll & Vicente (2020), and further studied by Pan et al. (2022). who also proposed DAE.
1989 Sutton 1988 have been widely adopted in recent deep RL research and shown to have a strong effect performance (Schulman iet al. 2015b Hessel et al., 2018 Wang et al.) 2016a Gruslys et al. 2017 Espeholt et al.) 2018 Hernandez-Garee Garcia & Sutton 2019), Typical off off-policy multi-step methods include importance sam- pling (Munos et 2016 Rowland et al., 2020 Precup et al. (00T). truncating (diminishing) off-
Multi-step Learning Multi-step methods (Watkins)
policy actions (Watkins| 1989 Precup et al. 2000), a combination of the two (De Asis et al. 2018).
or simply ignoring any correction.
Afterstates The idea of dissecting transitions into a two-step process dates at least back to Sutton jet al. (1988,, where afterstates (equivalent to nature's states in Figure [ were introduced. was shown that learning the values of afterstates can be easier in some problems. Similar ideas also appeared the treatment of random events in extensive-form games, where they are sometimes referred to as "move by nature" (Fudenberg & Tirole 1991).
Luck Mesnard et al. 1 proposed to use future-conditional value functions to capture the effect of luck, and demonstrated that these functions can be used as baselines in policy gradient methods to reduce variance. In this work, we approached this problem from causal effect perspective and provided a quantitative definition of luck (see Equation TOp.
DISCUSSION
In the present work, we demonstrated how DAE can be extended off-policy settings. We also re- late Off-policy DAE to previous methods to better understand how it can speed up learning. Through experiments both stochastie and deterministic environments, we verified that the proposed off- policy correction is beneficial for policy optimization.
One limitation of the proposed method lies in enforcing the á constraint in stochastic environments. In the present work, this was approximated using CVAEs, which introduced computational over- head and additional hyperparameters. One way reduce computational overhead and scale to high dimensional domains is to learn a value equivalent model (Antonoglou et al. [021] Grimm et al. 2020). We will leave it as future work to explore more efficient ways to enforce the constraint.

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGMENTS
HRP would like to thank Nico Gürtler for the constructive feedback. The authors thank the Inter- national Max Planck Research School for Intelligent Systems (IMPRS-IS) for supporting Hsiao-Ru Pan.
REFERENCES
Ioannis Antonoglou, Julian Schrittwieser, Sherjil Ozair, Thomas K Hubert, and David Silver. Plan- ning in stochastic environments with a learned model. In International Conference on Learning Representations, 2021.
Leemon C Baird. Reinforcement learning in continuous time: Advantage updating. In Proceedings of 1994 IEEE International Conference on Neural Networks (ICNN"94 994,, volume 4, pp. 2448 2453. IEEE, 1994.
M. G. Bellemare, Y. Naddaf, J. Veness, and M. Bowling. The arcade learning environment: An evaluation platform for general agents. Journal of Artificial Intelligence Research, 47:253-279, jun 2013.
Oriol Corcoll and Raul Vicente. Disentangling causal effects for hierarchical reinforcement learning. arXiv preprint arXiv:2010.01351, 2020.
Kristopher De Asis, Hermandez- Garcia, G Holland, and Richard Sutton. Multi-step reinforcement learning: A unifying algorithm. Proceedings of the AAAI Conference on Artificial Intelligence, volume 32, 2018.
Lasse Espeholt, Hubert Soyer, Remi Munos, Karen Simonyan, Vlad Mnih, Tom Ward, Yotam Doron, Vlad Firoiu, Tim Harley, lain Dunning, et al. Impala: Scalable distributed deep-rl with im- portance weighted actor-learner architectures. In International conference on machine learning, pp. 1407-1416. PMLR, 2018.
Drew Fudenberg and Jean Tirole. Game theory. MIT press, 1991.
Carles Gelada, Saurabh Kumar, Jacob Buckman, Ofir Nachum, and Mare Bellemare. Deepmdp: Learning continuous latent space models for representation learning. In International Conference on Machine Learning, pp- 2170-2179, PMLR, 2019.
Evan Greensmith, Peter L Bartlett, and Jonathan Baxter. Variance reduction techniques for gradient estimates in reinforcement learning. Journal of Machine Learning Research, 5(9), 2004.
Christopher Grimm, André Barreto, Satinder Singh, and David Silver. The value equivalence princi- plc for model-basee reinforcement learning. Advances in Neural Information Processing Systems, 33:5541-5552. 2020.
Audrunas Gruslys, Will Dabney, Mohammad Gheshlaghi Azar, Bilal Piot, Marc Bellemare, and Remi Munos. The reactor: fast and sample-efficient actor-critic agent for reinforcement learn- ing. arXiv preprint 1704.
2017.
Danijar Hafner, Timothy P Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering atari with discrete world models. In International Conference on Learning Representations, 2020.
J Fernando Hernandez-Gaariia and Richard S Sutton. Understanding multi-step deep reinforcement learning: systematic study of the dqn target. arXiv preprint 10XIvv:90010755 2019.
Matteo Hessel, Joseph Modayil, Hado Van Hasselt, Tom Schaul, Georg Ostrovski, Will Dabney, Dan Horgan, Bilal Piot, Mohammad Azar, and David Silver. Rainbow: Combining improvements in deep reinforcement learning. In Proceedings of the AAAI conference on artificial intelligence, volume 32, 2018.
Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterizzaion with gumbel-softmax. arXiv preprint 166 1.01 144, 2016.
10

---

Published as a conference paper at ICLR 2024
Sham Kakade and John Langford. Approximately optimal approximate reinforcement learning. In In Proc. 19th International Conference on Machine Learning. Citeseer, 2002.
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint 1411.6:412.6980 2014.
Diederik Kingma and Max Welling. Auto-encoding variational bayes. 131.v:1112.61 2013.
arXiv preprint
Chris J Maddison, Andriy Mnih, and Yee Whye Teh. The concrete distribution: A continuous relaxation of discrete random variables. arXiv preprint arXiv: (000:///1..0011,, 2016.
Thomas Mesnard, Theophane Weber, Fabio Viola, Shantanu Thakoor, Alaa Saade, Anna Haru- tyunyan, Will Dabney, Thomas S Stepleton, Nicolas Heess, Arthur Guez, et al. Counterfactual credit assignment in model-ffee reinforcement learning. In International Conference on Machine Learning. pp. 7654-7664. PMLR, 2021.
Marvin Minsky. Steps toward artificial intelligence. Proceedings of the IRE, 49)1))8-22 30, 1961.
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 International conference on machine learning, pp. 1928-19299 1928-1937. PMLR, 2016.
Rémi Munos, Tom Stepleton, Anna Harutyunyan, and Marc Bellemare. Safe and efficient off-policy reinforcement learning. Advances neural information processing systems, 29, 2016.
Hsiao-Ru Pan, Nico Gürtler, Alexander Neitz, and Bernhard Schölkopf. Direct advantage estima- tion. Advances in Neural Information Processing Systems, 35: 1869-118881 2022.
Doina Precup, Richard S. Sutton, and Satinder P. Singh. Eligibility traces for off- off-policy policy evaluation. In Proceedings of the Seventeenth International Conference on Machine Learning, ICML '00, pp. 759-766, San Francisco, CA, USA, 2000. Morgan Kaufmann Publishers Inc. ISBN 1558607072.
Doina Precup, Richard S Sutton, and Sanjoy Dasgupta. Off-policy temporal-difference learning with function approximation. In ICML, pp. 417 424,,44,2 2001.
Mark Rowland, Will Dabney, and Rémi Munos. Adaptive trade-off's in off-policy learning- In International Conference on Artificial Intelligence and Statistics, pp. 34-44. PMLR, 2020.
Tom Schaul, André Barreto, John Quan, and Georg Ostrovski. The phenomenon of policy chum.
arXiv preprint arXiv:2206.00730) 2022.
John Schulman, Sergey Levine, Pieter Abbeel, Michael Jordan, and Philipp Moritz. Trust region policy optimization. In International conference on machine learning, pp. 1889-1897. PMLR, 2015a.
John Schulman, Philipp Moritz, Sergey Levine, Michael Jordan, and Pieter Abbeel. High- dimensional continuous control using generalized advantage estimation. arXiv preprint arXiv:1506.02433, 2015b.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347, 2017.
Max Schwarzer, Ankesh Anand, Rishab Goel, R Devon Hjelm, Aaron Courville, and Philip Bach- man. Data-efficient reinforcement learning with self preficredi representations. arXiv preprint NwXi::2007.05599, 2020.
Kihyuk Sohn, Honglak Lee, and Xinchen Yan. Learning structured output representation using deep conditional generative models. Advances in newral information processing systems, 28, 2015.
Richard S Sutton. Learning to predict by the methods of temporal differences. Machine learning, (1):9 44, 1988.
Richard S Sutton, Andrew G Barto, et al. Introduction reinforcement learning. 1998.

---

Published as a conference paper at ICLR 2024
Csaba Szepesvári. Algorithms for reinforcement learning. Synthesis lectures artificial intelli- gence and machine learning, 4(1):1-103, 2010.
Aäron van den Oord, Oriol Vinyals, and Koray Kavukcuoglu. Neural discrete representation learn- ing. In NIPS, 2017.
Ziyu Wang. Victor Bapst, Nicolas Heess, Volodymyr Mnih, Remi Munos, Koray Kavukcuoglu, and Nando de Freitas. Sample efficient actor-critic with experience replay. arXiv preprint 11102:461..01124. 2016a.
Ziyu Wang, Tom Schaul, Matteo Hessel, Hado Hasselt, Mare Lanctot, and Nando Freitas. Ducling network architectures for deep reinforcement learning. In International conference machine learning, pp. 1995-2003. PMLR, 2016b.
Christopher John Cornish Hellaby Watkins. Learning from delayed rewards. PhD thesis, 1989.
Kenny Young and Tian Tian. Minatar: An atari-inspired testbed for thorough and reproducible reinforcement learning experiments. arXiv preprint 1903vv1903.03111 2019.
12

---

Published as a conference paper at ICLR 2024
A PROOF OF THEOREM
Theorem (Off-policy DAE). Given a behavior policy P and a warett policy a and backup length n O. Let A, = A(st, ar), B, = B(se, de, 8t41), and the constrained squared error
2
L(Á, BV E J
Σ 7 (re A. - + +
(15)
=0 =0 A
subject to Eves B(s, a, ('ss","))('''sss)) as)
then VV) is minimizer of the above problem. Furthermore, the minimizer is unique if sufficiently explorative (i.e., non-zero probability of reaching all possible transitions (s, a, 8')).
Proof. Since
)
27
0 < B"A",B", V")= = E Σ r( A -B,,) + +y n (sn+1) - V-' (50) = 0,
(16)
and both E Eeea A(als) A (s, 0 and Eves p(8'1s, a) B" 13T(( (s, aa)')) constraints satisfied, (A",, D", VT) is minimizer of the problem. For uniqueness, sufficiently explorative such that any sequence (8o, ag To. ro...))111) has non-zero probability of being visited. Now, suppose there exists (A', B', V*) that also minimizes L, i.e., L(A', B', V`) 0, then for any sequence (s0, 4o S14 1). we must have
assume the behavior policy is
Σn' A4 BB) + " (8n+1) V'(s)) 0,
(17)
otherwise L(A', B', V') + O. If we take the conditional expectation over (ao, 81, ai .14411)) con- ditioned on 80 using the target policy æ, then
('(s)) = E
('r Ai -B,) + 7 v
(18)
1=0
7 rr + 7 n+ ('sn+)(s)))))
(19)
which means that V' satisfies the Bellman Equation. Therefore V v uniquely. Similarly, if take the expectation over (81, (1,11))11) conditioned on Su, ag then
1
A'(so. on) r(so, ao) +E Yr n+ V V ool V (so) A (sp.
(20)
Finally, if we take the expectation over (a1, 82, ..,,... s1+1) conditioned on so, Go, $1, then
BB'(so, an. s1) = r(s0. an) ( ao) + 2 re + + V (S1+1) ag. s1] (so)
(21)
= -s)
ool) Bd (so: ao 81)-
(22)
Similarly, we get ('',B',V') = (A= BT, V*) for all (sy, ae, Sp41) with 0 < 1' S n by re- peatedly taking the conditional expectations over the sequence. By the assumption that μ has non-zero probability of visiting any sequence, we have (A', A','',,V (8, a, S S.
(A", BT V*) for all
Remarks: (1) While we used the squared error the objective function in the theorem, it can be replaced with an arbitrary metric in R, as the proof does not rely on properties of the squared error. (2) For uniqueness, the condition the behavior policy states/actions covered by the target policy æ. In that case, we only need the coverage of μ to include the coverage of
be relaxed if we only care about
13

---

Published as a conference paper at ICLR 2024
B DETAILS OF FIGURE 3 AND FIGURE
In this example, we compare the sample efficiency of MC, Batch TD(O) and DAE. We note that there only possible trajectories in this environment, depending the starting state or 2) and the action chosen at state (u or d). We denote the number of trajectories starting from i and choosing action a E {u. d} by ni.a. and n, Ri,a + ni.a- The trajectories were sampled using the uniform policy, i.e., (T]33) = (d]3) = 0.5. In the following list, we summarize the estimates from each method.
MC: ((1)
(22) 72
B1
Batch TD(0): V(1) ((2) V(3)
ni n1
DAE: The minimizer of
n1, A(3, v(1)) -1,a A(3. V(1)²
+ n2, A(3, u) V(2))² + A(3.d) V(2))²
subject to A(3,u) + A(3, d) (since the sampling policy is uniform).
One the method of Lagrange multiplier solve the DAE problem and arrive the following linear equations:
A(3,u) u) V(1)
a1
11,dd
711
11,
n22, - 122
0
=
722
(23)
(11,u + 122,u)
11,u
T22,
((2)
f1, T22,
Note that there are only 3 equations, since Á(3, + A(3, d) 0. Additionally, the solution is unique only when both nı and n2 > O, otherwise the first row or the second row of the matrix would be O. For simplicity, we use the pseudoinverse to compute the solution:
A(3. u) (11)
1i,, 111,d
17
R1,a
T22,u
722
0
722
T22,a
(24)
(2)
(ni, + 12,
1,,
122
11,a + 122
where denotes the pseudoinverse. This explains why the DAE estimates in Figure 3 are slightly skewed towards the beginning. F be obtained similarly by adding the B terms the DAE loss. One difference that, in the case of learned transition probabilities, the / constraint was enforced based on the estimated transition probabilities instead of fixed distribution.
14

---

Published as a conference paper at ICLR 2024
2
p = 0.5
Fa
p = 0.5
r= 0
Figure 6: A simple MDP with S = 2}. State 1 has only single action, which can lead to state 2 or the terminal state with equal probabilities. State has two actions u and d with rewards I and O, respectively, and both actions lead to the terminal state.
C COUNTEREXAMPLE
In this section, construct example demonstrate that naively applying DAE to off-policy data can lead to biased results. Consider the environment in Figure Suppose the data is collected with behavior policy f and we wish to estimate values for target policy If apply DAE directly to this problem without any off- aff-policy correction, then the loss is equal to
L(Á, V)
1 ju(ul2) (1 A(2, u) - -v()) ((((((2) (o A(2 v(1))
2 1
+ (0 (())²
(25)
since there only three possible trajectories in this environment. Now, if we include the constraint that E. (()l)) 0, then the problem can be solved by the method of Lagrange multiplier using the following Lagrangian:
L(
Σ A(2, (a)a))
(26)
The minimizer (A', V*) arg min L(Ā, V) is given by:
V V*(1) =
x(u22)
1+ (1/2)4 [0
A \=
(27)
=1 1 1- V (1)))(()
wlul2)
A' d2 V*(1) v (d))
m(d
meaning that (1( # (1) =(4)) and A + A if + p. One can also verify that, if π (on-policy), then (V", A²) = (V", 2n).
15

---

Published as a conference paper at ICLR 2024
D DETAILS OF THE MINATAR EXPERIMENTS
D.1 PSEUDOCODE
The more detailed pseudocode of the proposed actor-critic method is provided Algorithm Here we note some details about the training process. Unlike typical methods that store 1-step transitions in the replay buffer, our buffer consists of n-step trajectories. When computing the critic loss, we also compute the loss for each sub-trajectory as in Pan et (2022). For example, if (5o, 40, .... 8.) a sample trajectory, we accumulate the critic loss for all sub-trajectory (8s, da. 8n. for 1}. Also, to speed up training, we use parallel actors to sample transitions from the
{0, 1,2, .... environments.
Algorithm A simple Oollpy Actor-Critic Algorithm
Require: backup E (Uncorrected, DAE, Off Off-policy DAE, Tree} Require: n (backup length)
1: Initialize actor-critic components Ao(s, a). Vo(s), Bo(s, a, z), no(ass)
2: Initialize CVAE gz(z|s, a, Pa(z]s, a), (s(s'|s, a.
3: OEMA t-
4: D ()
5: Dn =
6: fort= =
do 7: Sample transition a. 8') with policy To
8:
Dn < D ((s,
9:
if is terminal or n-stop then
10:
D+ DU (concatenate(D.)) f
11:
12:
end if
13:
if + 1 mod steps per. updat 0 then
14:
Sample batch of trajectories B from D if baackup Off-policy DAE then
15:
16:
Train (z(2|s, a,)), Pa(z|s, a). pa(s'|s, 2) using B by Equation 12
17:
Ba(s, 8') + Exngg(-ls.a (a a, 2),,,,, (x Bo 88
18:
end if
19:
A0(5,a) + Ab(s,a) Ea [Ao(s, a)]
20:
Compute critic loss Leritic according to backup (Eq. or Eq, 28
21:
Anormalized f- sttsstgpaaeeeeeee cerp-rradiennAAAAA a Var As]
22:
Compute actor loss Lactor Easxo pormalized
8kL DuL (Te (T
23:
Train Tenne + Tractor with Adam (Kingma & Ba 2014
24:
θEMA TθEMA (1 (1-7)0 +)
25:
end if
end for
The n-step Tree backup Q Q-target is defined recursively by:
(sngerSr ae) r((e,
Σ
T(alse+1 10 (tagga(((1411 + TTa44 1) 1urget furgll (8t+1 1 0e+1)
(28)
where Qtarget Vieu + Agpa in our case.
16

---

Published as a conference paper at ICLR 2024
state 10 x 10 n
3x3 Conv2d, 128
3x3 Conv2d, 128
Linear, 1024
Linear, 1
Linear, |A|]
Linear, |A] Taass)
Linear, |A| 2
V(s)
A(s,a)
a(s,
Figure 7: Network architecture for the actor-critic. Each hidden layer is followed by ReLU activa- tion function. The dashed line indicates that gradients are stopped.
D.2 ACTOR-CRITIC NETWORK
In our experiments, we use a convolutional neural network followed by multiple heads to approxi- mate Ag, Bg, Vo, and Ty (see Figure (Mihh et al. 2016 Wang al. 20166). Since we train both the actor and the critic using single network simultaneously. to avoid interference between the two losses (Leritic and Lactor), we simply use the representation learned from by the critic to train the actor by stopping the gradients from Lactor to the shared network. This eliminates the need to balance for the different loss functions.
17

---

Published as a conference paper at ICLR 2024
D.3 CONDITIONAL VARIATIONAL AUTOENCODER (CVAE) NETWORK
We illustrate the training process and the network architecture in Figure First, we use small discrete latent space 2 which allows us compute expectations over Z efficiently. This also alleviates the need to use heuristics such as VQ-VAE (van den Oord et al.] (017))
or Gumbel-Softmax (Jang et al. 2016 Maddison et al. 2016), because we can now compute 2) exactly. Second, to elim- inate the need to balance between KL-divergence loss and the reconstruction loss, the conditional prior trained using the representation from the encoder with gradients stopped. This is similar to the approach in VQ- VAE where prior is trained separately. Finally, can sometimes collapse early in training. To mitigate this, we add a small entropy penalty for the posterior. Combining everything together, we have the loss function for CVAR:
Bz-qz(-1saa,s)) [log (3 Is, a, 2) = LE 9z(2|s, a, s' Pg s (s,
observed that the posterior
LovwE(O, 83,,0, 8)) = DxI (xa,(g((3ss, (opo(=)),))
Dog Po (8'1o Is, a, =)]
-ee H (1g(-1s,,,
where H()) is the entropy, and Bent controls the strength of the entropy penalty.
D.4 HYPERPARAMETERS
In Table I summarize the hyperparameters used in the MinAtar experiments. In general, the agent was designed to have very few hyperparameters to reduce potential confounding when com- paring different backup methods. Our preliminary experiments found that the effects of the hyper- parameters to be agnostic to backup methods, except for which tend to have a larger impact on Off-policy DAE and DAE, which is likely due to the heavy dependence the policy when training with DAE-like loss functions.
For CVAE training, found the Adam's suggested 3 (0.9, 0.999) can sometimes lead to diver- gence, and lowering it to (0.5, 0.9) can greatly improve stability.
Table I:
List of hyperparameters. Note that for Off-policy DAE, there are two separate optimizers used for actor-critic and CVAE training. 1: not used in Pan et al. (2022).
Group
Parameter
Value
Sticky Action
False
Environment setting
Difficulty Ramping
False
Maximum Episode Length
108000 frames
Discount y
0.99
Parallel actors
128
Initial steps before training
25000 frames
Replay Buffer Size
1000000 frames
Backup Length
8/16/32
Optimizer
Adam(Kingma Ba 2014)
Shared (actor-critic training)
Learning rate
0.00025 (linearly annealed to 0)
Adam
(0.9, 0.999) 10-4
Adam
Env. steps per update Batch Size
32
1024 frames
3KL
3.0
T (EMA weight)
0.999
Latent size 2
Optimizer
Adam
Off-policy DAE only (CVAE training)
Learning rate
0.00025
Adam 9
(0.5, 0.9)
Adam
10-8
Bent
0.0001
18

---

Published as a conference paper at ICLR 2024
state 10 x 10
state s' 10 x
Encoder
Encoder
3x3 Conv2d, 64
tile
action
concatenate
Residual Block
3x3 Conv2d, 128
3x3 Conv2d, 128
Residual Blocks x2
Residual Block
Residual Blocks x2
concatenate
avg. pooling
Liveooa (s/ (5
Residual Block
3x3 Conv2d, 128
Linear, |2|
Po(2|s,a)
Residual Blocks x2
avg. pooling
Batch Norm
SiLU
Linear, |2|
3x3 Conv2d
DKi (az(22,, ('pg(()s a))
( [2|s, a, s')
Batch Norm
sample
SiLU
Decoder
3x3 Conv2d
state g 10 x 10 x n
Figure 8: Network architecture and the training graph for the CVAE. Dashed lines indicate that gradients are stopped. The Decoder the architecture the encoder with orders reversed and convolutions replaced with transposed convolutions. Softmax is applied to both q3 (=)s, a, s') and p3(2|s, a)) to ensure they are probability distributions over Z (note the Z is discrete space this case). We use binary cross entropy for the reconstruction loss, since states in the MinAtar environments are binary images.
19

---

Published as a conference paper at ICLR 2024
Uncorrected
3.0
Unncerracedd
3.0
Uncorrected
DAE
DAE
DAE
2.0
Nomalted Score
DAE
Normaradd Srene
OFF pollcy
Narmateed
1.5
2.0
2.0
1.5
1.0
1.0
0.5
0.0
10
10
Frames (Million)
Froes (Million)
Promes (Millin
Figure 9: Normalized training curves for each backup method and backup length (from left right:
8, 16, 32). The dashed horizontal line represents the PPO-DAE baseline.
D.S MORE RESULTS
In the MinAtar experiments, normalized score were calculated by score score - baseline we use the scores reported by Pan et al. (2022) as baselines. We note that in the MinAtar suite, BreakouBland Space Invaders are deterministic environments, while Asterix, Freeway, and Seaquest are stochastic environments. Interestingly, we found that Off-policy DAE tends to slightly outper- form DAE on Space Invaders (see Figure ITI. This is likely because Space Invaders is partially observable, and modeling the transitions as stochastic can be helpful.
baseline where
Table 2: Score comparison between different methods and backup lengths. Results were aggregated over 20 random seeds. Numbers represent (mean)+t( standard error)
Backup Method
Environment
N
Uncorrected
DAE
Off-policy DAE
Tree
8
4.5 + 0.2
155.6 + 5.4
161.5 + 5.9
44.3 + 1.1
Asterix
16
2.5 1 0.1
194.2 1 4.8
207.0 1 9.2
39.0 + 1.4
32
2.0 0.1
215.4 8.1
268.9 +:: 14.0
39.0 + 1.3
8
5799.9 + 611.3
9573.0 + 250.1
9423.3 = 351.9
9411.4 = 505.5
Breakout
16
5220.1
448.7
8506.1 476.6
8887.9 429.0
10069.0
288.8
32
3179.3 661.8
8119.8 + 617.5
7372.1 + 582.0
10139.8 338.1
8
5.5 1.8
61.9
0.4
Freeway
16
16.8 1.3
57.6 0.1
62.3 +
4.1 1.0
32
8.2 0.8
58.8 + 0.1
62.9 + 0.3
5.1 + 0.8
8
13.5 1.5
413.5 25.1
839.4 + 48.3
312.3 + 16.4
Seaquest
16
4.1 0.1
594.3 + 36.3
1171.6 66.6
286.4 + 11.6
32
4.1
1225.3 + 63.7
266.2 = 19.1
8
5561.5 + 452.3
15116.1 +: 377.4
18086.9 + 419.1
15615.3 = 563.0
SpaceInvaders
16
1180.3 + 145.0
13478.7 391.2
16756.8 460.4
16009.6 508.1
32
307.7
32.0
12560.4 441.4
14970.3 348.9
17374.2 584.9
D.6
WITH AND WITHOUT TARGET NETWORK
While Equation II suggests that separate target network for critic training is necessary, in practice, we have found that using a target network leads to better performance. See Figure [2 for results regarding the effect of target networks. In general, we found the critic loss be lower when using target networks. One possible explanation is that using target networks results in biased, but lower variance estimates. which in turn makes the loss easier optimize. Further investigation is required to understand the tradeoffs.
TThe initial states in Breakout random, but the transitions stochastic.
20

---

Published as a conference paper at ICLR 2024
Backup Length
Backup Length 16
Backup Length 32
500
Astertk
200
300
300
200
100
100
100
5
10
5
10000
10000
10000
8000
8000
Breakout
6000
4000
4000
4000
2000
2000
2000
0
10
10
10
60
60
Fresway
40
Score
20
20
0-
2000
2000
1500
1500
1500
Seaquest
1000
500
500
500
0
5
10
10
10
20000
20000
Spacelrvaders
15000
15000
15000
10000
10000
www
5000
5000
5000
5
10
Uncorrected
DAE
Off-palicy DAE
Tree
Frames (Million)
Figure 10: Training curves of each environment.
21

---

Published as a conference paper at ICLR 2024
Asterix
Breakout
Freeway
300
60
10000
200
5000
100
20
32
16
32
16
Score
Backup Length
Backup Length
Backup Length
Seaquest
Spacelnvaders
1500
20000
15000
Uncorrected
1000
DAE
10000
Off-polis policy DAE
Tree
500
5000
8
16
32
8
16
32
Backup Length
Backup Length
Figure 11: Distributions of the scores obtained by each method under different backup lengths.
Asterix
Breakout
Freeway
500
10000
60
400
8000
300
6000
200
4000
100
2000
20
0
10
5
10
5
10
Score
Seaquest
Spacelnvaders
2000
20000
1500
15000
1000
5000
0
10
10
With Target
Without Target
Frames (Million)
Figure 12: Off-policy DAE with and without target networks.
22

---

Published as a conference paper at ICLR 2024
D.7 COMPUTATIONAL RESOURCES
All experiments were performed on an internal cluster of NVIDIA A100 GPUs. Training an agent takes approximately 2 hours, depending on the backup method and the environment. For Off-policy DAE, the training time is significantly longer (approx. 15 hours) due to CVAE training. The increase in training time is largely due the use of a large residual network for the CVAE, which we found to be easier to optimize compared to smaller convolutional networks.
23


Reviews:
[{"review_id": "d1kYLVaxDk", "reviewer": "Reviewer_rop6", "summary": "This paper extends Direct Advantage Estimation (DAE), an algorithm designed to improve credit assignment by directly learning the advantage function. The original DAE formulation was limited to the on-policy case and this paper derives an off-policy version. This new algorithm is shown to be beneficial in a toy example and with experiments on the MinAtar environments.", "questions": "- Off-policy DAE requires an additional neural network to estimate the transition distribution. This seems slightly unelegant since actor-critic algorithms are usually model-free. Is it possible to avoid this by converting the constraint into a loss function instead? i.e. optimize the Lagrangian of the constraint with SGD? Perhaps we would have to use an approximate loss here to make it tractable. \n\n- If we do learn a model to estimate $B_\\pi$, is it possible to adapt the algorithm so the model is sampled-based only? Requiring a sum over all discrete latent states seems a bit restrictive in the choice of model and it seems to prevent easy scaling of the model size.\n\n- Could you clarify how equation (10) reduces to the policy improvement lemma? In particular, what happens to the $B^\\pi_t$ term? \n\n- Why was Tree Backup chosen as the baseline? It seems like a less popular or effective choice compared to, say, Retrace [1] or its successors. \n\n- It could be interesting to try to incorporate off-policy DAE with model-based RL methods since those algorithms would already have a learned model to use. E.g. Dreamer-v3\n\nMinor point:\n- To improve the clarity of the notation, I would suggest using capital letters for the random variables. In certain places, the lowercase letters are used to denote both fixed values and random ones. E.g. below equation (7), in the definition of $B^\\pi_t$, $s_{t+1}$ in the expectation should be uppercase.\n\n\n[1] \"Safe and efficient off-policy reinforcement learning\" Munos et al.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- The decomposition of the return in terms of advantages of both the agent and the environment's \"actions\" is very intriguing and novel. It also leads to a practical algortihm, which could potentially outperform the standard approaches to learning advantage functions, a crucial part of actor-critic algorithms.\n\n- The paper is well-written and this makes the derivation much easier to follow.\n- The toy environments are well-designed to demonstrate the differences between the different approaches, DAE vs. off-policy DAE.\n- There's some nice insights into why the uncorrected off-policy n-step return may work well in practice due to certain environmental properties.", "weaknesses": "- The larger-scale experiments are fairly limited with MinAtar being the most complex domain. Other environments could be considered.\n- Other baselines could be more approrpirate for the MinAtar experiment. See questions.", "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": "", "extra_scores": {}}, {"review_id": "lvmEPQ4Jlv", "reviewer": "Reviewer_y5tB", "summary": "This paper presents Off-policy DAE, a framework for off-policy learning that decomposes returns into two components: 1. those controllable by agents (skills) and 2. those beyond an agent's control (luck). Specifically, \"luck\" refers to stochastic transitions that the agent can't control. By explicitly modeling the advantage attributed to this luck factor, the agent can more effectively discern the impact of its own actions, leading to quicker generalization through enhanced credit assignment. Evaluations conducted on 5 MinAtar environments demonstrate improvements over baselines, particularly in scenarios with stochastic transitions.", "questions": "Q1: Can the method be directly applied to offline RL, and if so, what additional challenges might arise in that context?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "S1: This offers a methodical approach to incorporate the effects of non-controllable factors, leading to enhanced credit assignment in off-policy learning. It's a novel concept, one I haven't encountered in other papers.\n\nS2: The paper is very well-written and straightforward. I'm especially impressed by the intuitive examples provided to make the reader understand the main concept.\n\nS3: Experiments are conducted on a standard benchmark suite using 20 random seeds. The results are notable, demonstrating that in stochastic environments, the agent performs better compared to methods not leveraging the proposed advantage function decomposition.", "weaknesses": "The paper is largely well-composed, with the proposed idea articulately presented. While I couldn't pinpoint any specific areas needing improvement,\n\nW1: Conducting experiments in more domains might bolster the paper's claims even further.", "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": "", "extra_scores": {}}, {"review_id": "sDiKZL90vc", "reviewer": "Reviewer_NYwT", "summary": "This paper proposed a new return decomposition method via advantage functions, which contains the average return, the skill term, and the luck term.", "questions": "- Why the decomposition of the return could benefit learning? Could the author give more insights into it?\n- What's the probability of the sticky action for stochastic environments?", "limitations": "", "rating": 5, "confidence": 5, "soundness": 3, "presentation": 1, "contribution": 3, "strengths": "The formulations of the paper are quite clear and easy to follow. Although I didn't get the main idea of the paper through the text, the formulations are quite clear for me to follow.\n\n\n\n\n\nGenerally, I think the paper contributes to giving some insight into the decomposition of the return.", "weaknesses": "Although I can get the main ideas from the formulations of the paper, the writing of the paper is poor.\n\n- The introduction is not quite clear. Although it starts with an intuitive example, it doesn't contribute too much to get the idea of the paper.\n- Section 3 lacks motivation, making it kind of hard to follow.\n\n\n\nThe experimental results are not very convincing.\n\n- For Figure 9, Table 2, Figure 10, 11, 12, are the results in deterministic environments? Are there any results in stochastic environments?  I believe this is very important as the evaluation in stochastic environments is the main contribution of the paper.\n- The authors used very large $N$ (8) for Uncorrected. However, in existing literature, it's known 3 or 5 are the best. Therefore, the comparison is unfair.\n- What's the performance of 1-step DQN? \n- The paper doesn't compare to other state-of-the-art multi-step off-policy method, such as Retrace($\\lambda$) [1].\n\n\n\n[1] Munos R, Stepleton T, Harutyunyan A, et al. Safe and efficient off-policy reinforcement learning[J]. Advances in neural information processing systems, 2016, 29.\n\n\n\nThe below claim is exaggerated. I didn't find a theorem that clearly proves the properties of faster convergence.\n\n> We demonstrate that (Off-policy) DAE can be seen as generalizations of Monte-Carlo (MC) methods that utilize sample trajectories more efficiently to achieve faster convergence.\n\n\n\n\n\nMinor comments:\n\n- Near eq.7: The second term $E{[V^\\pi(s_{t+1})|s_t,a_t]}$. What are the random variables of expectation? If it's $s_{t+1}$, please use another notation to differentiate it from the existing one using $s_{t+1}$.\n- The complexity of the method seems very large (15 hours vs. 2 hours, as stated in D.7)", "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": "", "extra_scores": {}}, {"review_id": "DjmZMHIkMD", "reviewer": "Reviewer_K7qX", "summary": "The proposes to decompose the advantage function into two parts: 1) return due to the agent's action selection and 2) return due to transition dynamics of the environment. They then use this decomposition to extend an existing algorithm DAE to the off-policy setting.", "questions": "1. It is not clear to me why off-policy corrections are not necessary given that the sequence of rewards was generated by a different policy; it feels like frequency at which the sequences appear must be corrected for.\n2. Related to above, it appears that the only \"off-policyness\" in Eqn 5 is the $\\pi$-centered constraint. Why is this sufficient for the off-policy correction?\n3. The decomposition seems related to exogenous and endogenous stochasticity. Is there a way to phrase the current work in that context? I'd also refer the authors to this paper that seems relevant: https://people.csail.mit.edu/hongzi/var-website/content/iclr-19.pdf\n4. I am curious if once the advantage function is decomposed into skill and luck, is there a benefit to weighing each component differently? I would suspect that this leads to some bias in the policy ordering, but I am wonder if say the skill related component is too small, it may get overshadowed by the luck component, and the agent may not learn efficiently.\n5. Related to above, I am curious how off-policy DAE performs as a function of environment stochasticity. \n6. In Figure 3 and 4, it is unclear to me why all methods are able to produce similar mean estimates? Of course each is different in terms of their variance, but all are centered around the same mean which is a bit surprising.\n7. What were the number of trials for the Figure 5 results? These should be mentioned.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The proposed decomposition is straightforward, and I think is a useful idea to spark newer ideas.\n2. The toy examples and Section 4 is useful.", "weaknesses": "See questions", "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": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)