Direct Advantage Estimation in Partially Observable Environments

Is it necessary to incorporate reward signal $r_{0:t−1}$ to the information vector $h_t$? Since $r_t$ info seems not be incorporated in the $h_t$.

Whether reward signals should be incorporated into the information vector depends on how the reward function is defined. Intuitively, if the reward function depends on the state (i.e., $r: \mathcal{S} \times \mathcal{A} \rightarrow \mathbb{R}$, the setting used in the present work) [4], then rewards should be incorporated into the information vector as they might contain additional information about the state. On the other hand, if the reward function is defined as $r: \Omega \times \mathcal{A} \rightarrow \mathbb{R}$ [5], then rewards are not needed as past observations/actions already form a sufficient statistic [6]. $r_t$ is not incorporated into $h_t$ because, under our notation, $r_t$ is observed after taking action $a_t$.

The reviewer is a little bit confusing on the mixing use of state $s, h, z,$ and $x$. It would be great to have a better clarification on these symbols and definitions of corresponding $\hat{A},\hat{B}$, and $\hat{V}$.

We have updated Section 3 to make the notations more consistent, and Figure 1 to better illustrate the relationships between variables.

[1] Pan, Hsiao-Ru, and Bernhard Schölkopf. "Skill or luck? return decomposition via advantage functions." arXiv preprint arXiv:2402.12874 (2024).

[2] Hafner, Danijar, et al. "Mastering atari with discrete world models." arXiv preprint arXiv:2010.02193 (2020).

[3] Hafner, Danijar, et al. "Mastering diverse domains through world models." arXiv preprint arXiv:2301.04104 (2023).

[4] Kaelbling, Leslie Pack, Michael L. Littman, and Anthony R. Cassandra. "Planning and acting in partially observable stochastic domains." Artificial intelligence 101.1-2 (1998): 99-134.

[5] Liu, Qinghua, et al. "When is partially observable reinforcement learning not scary?." Conference on Learning Theory. PMLR, 2022.

[6] Bertsekas, Dimitri P. "Dynamic programming and optimal control 3rd edition, volume ii." Belmont, MA: Athena Scientific 1 (2011).

We thank the reviewer for the detailed and constructive feedback, please see our answers to your questions below.

line 247, 255, and 266 would be better to refer to Eq 9 instead of Eq 14 back in the appendix.

We fixed this in the latest revision.

Though technical sound, this paper seems a bit incremental comparing with previous works[4] with a different dynamic reparameterization for high-dim state space.

We would like to emphasize that, this seemingly incremental change to the dynamics model can have far-reaching consequences. Previously, Off-policy DAE’s empirical results were limited to environments with low-dimensional states, precisely because learning a stochastic dynamics model was too computationally expensive. As a result, it remained unclear whether Off-policy DAE would also be effective in high-dimensional domains. The approach presented in this work –combining SPR, a representation learning method, with the WTA loss, a method for multi-modal predictions, to approximate the dynamics–represents, to the best of our knowledge, a novel approach to this problem that can be easily applied to high-dimensional domains. Furthermore, our empirical results not only validated the effectiveness of this approach, but also demonstrated the viability of Off-policy DAE in high-dimensional domains, which, we believe, is a valuable contribution.

Section 3.2 states the reduction of computational complexity comparing to LSTM and RNN based methods, but it would be great to have some empirical or theoretical demonstrations.

The method of truncating sequences, presented in Section 3.2, is not intended to replace RNN-based methods, but rather to reduce their computational complexity. This reduction in computational complexity is straightforward, as parts of the trajectories are discarded when forwarding them through an RNN. While this method is widely adapted in deep RL, Section 3.2 highlights a commonly overlooked problem (i.e., confounding) that can happen when applying this method. We have revised Section 3.2 to better explain how this method is typically implemented in practice and its connection to confounding.

The authors claim that off-policy correction is crucial for stochastic environment and performance improvement. It would be better to explain how the stochastic part is validated.

We have expanded the explanation regarding the off-policy correction in the latest revision. Essentially, the idea is that, for deterministic environments, we have $B^\pi\equiv 0$ for arbitrary policy $\pi$, and this term only comes into play when the environment is stochastic. As such, we can view the $\hat{B}$ terms in the DAE objective function as corrections for stochastic dynamics, and by forcing $\hat{B}\equiv 0$, we can see the importance of off-policy corrections in stochastic environments. This trick was also used in [1] to demonstrate the effectiveness of the corrections in stochastic environments.

Reviewer is also interested in seeing the comparison with other recent works with support for pomdp such as [1,2,3]. Current comparisons with Rainbow and DQN might not be sufficient.

As TD-MPC2 and FORBES were originally developed for continuous action domains, it remains unclear whether they would also work well for discrete action domains like Atari games. As a remedy, we have updated our list to include more recent SOTA model-based recurrent algorithms DreamerV2 [2] and Dreamer V3 [3] (evaluated at 20M training frames, same training frames as ours). We summarize the results in the following table:

We note that DreamerV3 uses a significantly larger model (~200M parameters) compared to our m=8 model (~50M parameters). To enable a closer comparison, we also include Dreamer V2 (~20M parameters). The results show that our method (m=8) achieves performance close to DreamerV3 in 3 out of 5 environments, while our smaller models (m=2, ~35M parameters) outperforms DreamerV2 in 4 out of 5 environments. Finally, while we also learn dynamics models for our methods, we differ from Dreamer in that we do not use these models for rollouts. Instead, we leverage them solely for the off-policy corrections in the DAE objective (the B constraint).

We thank the reviewer for the detailed and constructive feedback, please see our answers to your questions below.

The paper is hard for me to follow, and the writing could benefit from further polishing. For instance, the parameter notation $θ$ is used in various forms, such as $f_θ(h,a), g_θ(s,a,z)$, and $p_θ(z|h_t,a_t)$. This raises questions about whether these functions share the same parameters, which the paper does not clarify. Additionally, the notation for $z$ in Eq. (10) is undefined until $\mathcal{Z}$ is introduced before Eq. (11), making the notation confusing.

We have updated Section 3 to make the notations more consistent. In general, $f_θ$, $g_θ$ and $p_θ$ need not share the same parameters, although in practice we do share parts of the parameters (the CNN and the RNN backbones) to make the forward pass more efficient. We have also updated the experiment section to reflect this design choice.

The authors appear to consider only the simplest case, with finite state and action spaces. Under these conditions, why is DAE necessary, and why wouldn’t a TD loss with a Bellman operator suffice? The motivation for introducing DAE in this context is unclear and could be better presented.

While the convergence properties of TD algorithms are well-understood in tabular settings (finite state/action spaces), such analyses become much more challenging in the function approximation setting, especially within deep RL, due to the complexity of neural networks. As such, much of the research in deep RL has been empirical, and it remains unclear to what extent results from tabular settings can be transferred to this more complex setting. For instance, [1] demonstrates that convergent algorithms in tabular settings can sometimes diverge in deep RL settings.

The motivation for introducing DAE, as outlined in the abstract (line 10) and the introduction (line 34-40), stems from previous works that have demonstrated the effectiveness of DAE in deep RL settings [2, 3]. We view the theoretical results in the tabular setting from the present work as initial steps toward justifying DAE’s application in the POMDP setting, with our experimental results providing empirical evidence that DAE can indeed perform effectively in the deep RL setting.

The benchmark algorithms used in this paper date back 6-9 years, which weakens the strength of the comparison. I recommend that the authors incorporate more recent, robust baselines to more convincingly showcase the advantages of their proposed method.

We have updated our list to include more recent SOTA model-based recurrent algorithms Dreamer V2 [2] and Dreamer V3 [3] (evaluated at 20M training frames, same training frames as ours). We summarize the results in the following table:

We note that DreamerV3 uses a significantly larger model (~200M parameters) compared to our m=8 model (~50M parameters). To enable a closer comparison, we also include DreamerV2 (~20M parameters). The results show that our method (m=8) achieves performance close to DreamerV3 in 3 out of 5 environments, while our smaller models (m=2, ~35M parameters) outperforms DreamerV2 in 4 out of 5 environments. Finally, while we also learn dynamics models for our methods, we differ from Dreamer in that we do not use these models for rollouts. Instead, we leverage them solely for the off-policy corrections in the DAE objective (the B constraint).

[1] Van Hasselt, Hado, et al. "Deep reinforcement learning and the deadly triad." arXiv preprint arXiv:1812.02648 (2018).

[2] Hafner, Danijar, et al. "Mastering atari with discrete world models." arXiv preprint arXiv:2010.02193 (2020).

[3] Hafner, Danijar, et al. "Mastering diverse domains through world models." arXiv preprint arXiv:2301.04104 (2023).

We thank the reviewer for the detailed and constructive feedback, please see our answers to your questions below.

... Besides that they report only mean and 1 standard error of their agents, whereas reporting the median agent performance and either a 2 standard deviation interval or better yet a tolerance interval would have improved the visualization of their algorithm performance.

Currently, we reported performance measures for individual environments instead of the aggregated median performance because we only focus on a small subset of the full Atari suite. This allows us to more closely examine the impact of various corrections on specific environments (e.g., some environments are more partially observable, where LSTMs are more effective). As a remedy, we have included an additional performance profile as suggested by [1] in Figure 10. This profile displays the distribution of human-normalized scores along with bootstrapped confidence intervals.

Would it be possible to compare the performance with methods using similar networks to the ones you used, like Impala and R2D2? Why did you think this comparison was irrelevant?

We certainly do not regard other baselines as irrelevant. Originally, we used DQN/Rainbow as baselines because they are standard for Atari games. To strengthen our results, we have included more recent SOTA model-based recurrent algorithms DreamerV2 [2] and DreamerV3 [3] (evaluated at 20M training frames, same training frames as ours) to our list of baselines. We summarize the results in the following table:

We note that DreamerV3 uses a significantly larger model (~200M parameters) compared to our m=8 model (~50M parameters). To enable a closer comparison, we also include DreamerV2 (~20M parameters). The results show that our method (m=8) achieves performance close to DreamerV3 in 3 out of 5 environments, while our smaller models (m=2, ~35M parameters) outperforms DreamerV2 in 4 out of 5 environments. Finally, while we also learn dynamics models for our methods, we differ from Dreamer in that we do not use these models for rollouts. Instead, we leverage them solely for the off-policy corrections in the DAE objective (the B constraint).

Would it be possible to report ablations on using a Exponential Moving Average network vs a periodically updated network (DQN-style) as your target network? If not would it be possible to justify the choice a bit more?

We have included an additional ablation study regarding EMA vs periodic copy in the latest revision (see Appendix C.5). In summary, we find EMA to outperform periodic copy in 4 out of 5 environments. We believe this is partially due to SPR style self-supervised learning methods being more stable when trained with EMA updates [4],

[1] Agarwal, Rishabh, et al. "Deep reinforcement learning at the edge of the statistical precipice." Advances in neural information processing systems 34 (2021): 29304-29320.

[2] Hafner, Danijar, et al. "Mastering atari with discrete world models." arXiv preprint arXiv:2010.02193 (2020).

[3] Hafner, Danijar, et al. "Mastering diverse domains through world models." arXiv preprint arXiv:2301.04104 (2023).

[4] Ni, Tianwei, et al. "Bridging State and History Representations: Understanding Self-Predictive RL." arXiv preprint arXiv:2401.08898 (2024).

I agree with Reviewer nsQ7 that further evidence of the importance of partial observability for the algorithm's ability to outperform others being necessary, as such my score will also remain unchanged

We thank the reviewer for the detailed and constructive feedback, please see our answers to your questions below.

In my opinion, there is only one weakness worth focusing on. The main contribution is supposed to be the extension of DAE to the partially observable setting, but the Atari environments chosen are arguably not challenging from the partially observable perspective. In particular, the fact that DQN and Rainbow, algorithms that are typically paired with non-recurrent architectures, can maximize rewards in them suggests that partially observability is not a big issue on them. In this sense, I think that ablation studies in these environments, for example, one comparing frame-stacking vs. LSTMs, are not as effective as choosing tasks where partial observability is a crucial factor for good performance.

Now, while I consider that the paper should be accepted in its current version, I think that a higher score requires considering more appropriate environments. For example, see Pašukonis et al., Evaluating Long-Term Memory in 3D Mazes, in ICLR, 2023.

We would like to note that while agents with non-recurrent architectures (such as Rainbow) have achieved strong performance in Atari games, previous research has demonstrated that recurrent architectures can significantly improve performance across various environments (e.g., R2D2 [1]). This strongly suggests that partial observability in Atari games is a non-trivial factor. As a more concrete example, in Battle Zone, players control a tank to eliminate enemies and avoid attacks while having a limited view of its surrounding at each time step, which aligns closely with the definition of partial observability.

To the best of our knowledge, the extent of partial observability in Atari remains an open question. We believe our ablation studies provide valuable evidence supporting the presence of partial observability in these games, while also offering new insights, such as the confounding bias, to this problem.

What is the computational overhead of the predictive model?

For our models, the computational overhead of the predictive model, which consists of 3 MLPs (each with a single hidden layer), is similar to the value heads ($\hat{A}, \hat{B}, \hat{V}$) and almost negligible compared to other parts (convolutional layers or recurrent layers) of the network.

In line 246, why is $\hat{x}$ being called an information vector?

We have corrected this typo and made the notations more consistent in the latest revision.

[1] Kapturowski, Steven, et al. "Recurrent experience replay in distributed reinforcement learning." International conference on learning representations. 2018.

As the DreamerV3 scores reported in [1] are based on a much bigger model (~200M parameters) compared to ours (m=8, ~50M parameters), we conducted additional experiments using the 50M-parameter variant of DreamerV3 (see table 3 in [1]) for a closer comparison. In addition to model capacity, we also examine the impact of replay ratio, which specifies how much training was done (see page 18 of [1] for further details). All DreamerV3 experiments were based on the official implementation provided in [2] and were repeated over 3 random seeds.

The following table summarizes the result after 20M training frames (RR: replay ratio, values represent mean$\pm$1 standard error):

(by default, DreamerV3 uses replay ratio = 32)

In general, we find our method outperforms DreamerV3 (RR=16), and stays competitive with DreamerV3 (RR=32).

We hope these additional experiments address any questions regarding the baselines, and thank the reviewers once again for their constructive feedback.

[1] Hafner, Danijar, et al. "Mastering diverse domains through world models." arXiv preprint arXiv:2301.04104 (2023).

[2] https://github.com/danijar/dreamerv3