A Differentiable Sequence Model Perspective on Policy Gradients

Thank you for your feedback!

To me, state conditioning is necessary to predict the next states. Only conditioning on the actions does not provide complete information about the environment.

In deterministic environments, which are the object of study of our work, a sequence of actions is always sufficient to exactly determine future states, provided the inital state is given. Empirically, our experimental results show that they are sufficient for effective prediction and credit assignment. Even in certain stochastic environments, action sequences are often sufficient for next state prediction, as is seen from the latent models in [1] [2] and [3].

The connection between policy gradients and RNNs/sequence models seems obvious in the literature.

We agree that this connection has been intuitively been brought up multiple times, in works we cited in our paper [1] [4] [5]. It can be seen as a contribution of our work to transform this intuitive connection into an actual theoretical and algorithmic framework, that allows to build algorithms with better credit assignment properties by just employing more powerful sequence models.

I don’t quite see what brings the novel insights from the proposed understanding.

We encourage the reviewer to read our general response and re-evaluate the novelty of our paper. Our proposed understanding allows us, for the first time, to show how any arbitrary advanced sequence model can be used to directly improve credit assignment in policy gradient methods.

[1] Amos, Brandon, et al. "On the model-based stochastic value gradient for continuous reinforcement learning." Learning for Dynamics and Control. PMLR, 2021.

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

[3] Rezende, Danilo J., et al. "Causally correct partial models for reinforcement learning." arXiv preprint arXiv:2002.02836 (2020).

[4] Metz, Luke, et al. "Gradients are not all you need." arXiv preprint arXiv:2111.05803 (2021).

[5] Heess, Nicolas, et al. "Learning continuous control policies by stochastic value gradients." Advances in neural information processing systems 28 (2015).

Thank you for your positive feedback!

I would recommend adding some clarification between the proposed method and various temporal credit assignment methods using sequence models. [...] It would be beneficial to establish a connection between this work and these references. Further clarification of the connections, differences, and novelties would be appreciated.

We have added to the updated version of our paper a more detailed account of related work. See general response.

It would be beneficial to include comparisons in the benchmarks and experiments as well

We appreciate this suggestion. Unfortunately, as far as we know, other more common benchmarks in credit assignment are usually POMDPs (at least the reward is history-dependent) and/or discrete in nature [1]. In our paper the main focus is on MDPs with continuous action spaces, where policy optimization via backpropagation through time can be applied. We are excited to generalize our method to such environments in future work.

The results would be more convincing if some visualizations of what sequence models have learned are provided.

Thanks for the suggestion! We added in Figure 9 of the updated paper a visualization of the attention weights learned by an Action-Sequence Model transfom in the Toy Credit-Assignment problem. For predicting future states, the transformer attends to the only important action, the first one: in this way, when computing the policy gradient, the credit can directly flow from a reward at the end of the sequence to an action at the beginning of the sequence. This happens naturally, just by virtue of backpropagation through time and the particular architecture of the Action-Sequence Model.

[1] Ni, Tianwei, et al. "When do transformers shine in rl? decoupling memory from credit assignment." arXiv preprint arXiv:2307.03864 (2023).

Thank you for your feedback!

The idea of treating RL problems as sequence modeling tasks is not new and has been extensively covered in prior work, specifically in [1] and [2]

Our work does not treat the RL problem as whole as a sequence modeling task. Instead, it is geared towards a problem formulation and general methodology based on policy gradients, and relies on sequence modeling as a tool to learn a model of the dynamics, and to connect the theory of deep learning to RL. Indeed, a large part of the conceptual appeal of our proposed framework compared to e.g., decision transformers, is that we are able to leverage powerful sequence models, at the same time being grounded in the traditional theory of model-based policy gradients. To state these differences in an even clearer way, in addition to the discussion about the work mentioned by the reviewer already present in our submission, we have included in the updated paper a detailed account of the relationship between our work and other work employing sequence models in RL.

This paper does not clearly establish its unique contributions to the field, and the related work section is insufficiently detailed.

The length of our related work section was only due to space constraints. We included more discussion on the related work in the main text, the appendix of the updated paper, and on the general response above.

The experimental design does not effectively differentiate the proposed method from established sequence modeling algorithms.

Our method does not introduce any new "sequence modeling algorithm". As opposed to that, the rationale behind our work is to marry together well-established deep learning architectures and training techniques with the core RL machinery of policy gradients. The result is a framework that, in its simplicity, both provides new theoretical understanding and natural ways to get improved credit assignment. We believe our work is a first stepping stone exploiting advanced sequence models in policy gradient methods, and we hope future work will put engineering efforts to scale our methodology to more complex tasks.

A more robust comparison to state-of-the-art sequence modeling techniques, while not SAC and One-Step model, is necessary to validate the claims of the paper.

We are not aware of any state-of-the-art sequence modeling RL algorithms immediately appropriate for our problem setting, which is concerned with long-term credit assignment in deterministic MDPs for online RL. The closest method highlighted in this review is the online decision transformer from [1]. We are in the process of running experiments using online decision transformers applied completely online, and will provide an update as soon as the results are out (See edit, the experiments we're added into Appendix F.3).

The inclusion of common continuous control benchmark tasks, such as Hopper, HalfCheetah, Walker and Antmaze, is essential for a more comprehensive evaluation.

Our method is empirically evaluated for its ability to solve difficult temporal credit assignment tasks. As shown by recent work [2], the mentioned robotics locomotion benchmarks are not interesting from a temporal credit assignment perspective, and we are not aware of any common continuous temporal credit assignment MDPs beyond the environments we used. Also, note that, despite being generally low-dimensional, tasks from Myriad are based on existing dynamical systems and realistic applications such as cancer treatment, and could be considered by some readers as compelling as robotics locomotion, if not more.

[1] Zheng, Qinqing, Amy Zhang, and Aditya Grover. "Online decision transformer." international conference on machine learning. PMLR, 2022.

[2] Ni, Tianwei, et al. "When do transformers shine in rl? decoupling memory from credit assignment." arXiv preprint arXiv:2307.03864 (2023).

Edit 20/11/2023: We have included additional results on the online decision transformer on the Myriad environments. Please see Appendix F.3 and the general response. We hope that these results, along with our extended related work, clears up any confusion that our work does not differ enough from the line of work on decision transformers.

you said “Corollary 1.1 explains both the difficulties … and the limitations …”. What exactly are the difficulties and limitations here?

Prior works in model-based policy gradients have noted the difficulties of unrolling models for long horizons, and Corollary 1.1 justifies this difficulty due to the exploding gradients.

Why does your upper bound result indeed explain them (rather than just coincide with them)?

These exploding gradients are subsequently shown in Figure 5 to result in poor policies in Figure 4.

In your experiments, what's the difference between the "ASM(RNN)" model and the "One-step Model"?

Fundamentally, we note in Section 3.1 that they differ only in their training objectives: "When training an ASM as an RNN with teacher forcing, we are essentially training its recurrent cell as a one-step model". Experimentally, the ASM(RNN) is detailed in Appendix E: it is modeled as a 2-layer RNN with ReLU activations. Notably, we have added the architecture details of the one-step model, which also uses a 2-layer MLP with non-linear ReLU activations.

did you also try to upgrade the one-step model from a simple linear model to more sophisticated ones

Yes, all experiments use a non-linear 2 layer MLP for the one-step model. No experiment uses one-step linear models.

In Section 4.3, [...] are the environments here partially observable? [...] one-step Markovian models are just not appropriate for POMDP environments.

No, as specified in our background section, all the environments we employ in our paper are fully-observable and Markovian.

I am not sure that the capability to see the additional state info for "History Transformer" can really account for its bad performance, given that the attention modules in Transformer can be trained to simply ignore the state info if they are not helpful.

The poor performance of the History Transformer does not come from the inability of a transformer to ignore uninformative state information. Instead, for most environments, a transformer conditioned on the entire history will find the state information particularly useful for predicting the next states. As a notable example, in a Markovian case, the transformer can actually learn to just use the last state and action to predict the next state: in this case, unrolling the transformer will create a computational graph that would be similar to the one of an unrolled Markovian model, and thus potentially lead to ill-behaved gradients. For its particular structure, an Action-Sequence Model is instead constrained to directly attend to the actions at its input, thus directly propagating gradients from rewards to actions without unnecessarily chaining next-state predictions.

[1] Metz, Luke, et al. "Gradients are not all you need." arXiv preprint arXiv:2111.05803 (2021).

[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] Ghugare, Raj, et al. "Simplifying model-based rl: learning representations, latent-space models, and policies with one objective." arXiv preprint arXiv:2209.08466 (2022).

[5] Pascanu, Razvan, Tomas Mikolov, and Yoshua Bengio. "On the difficulty of training recurrent neural networks." International conference on machine learning. Pmlr, 2013.

[6] Heess, Nicolas, et al. "Learning continuous control policies by stochastic value gradients." Advances in neural information processing systems 28 (2015).

Thank you for the detailed feedback, we appreciate the thoroughness of your review!

in practice the estimated gradient is often normalized so the norm is less important, in my impression.

Clipping exploding gradients often results in biased gradients and does not always provide a benefit [1]. In fact, we tried to normalize gradients for all of our experiments, for one-step and true models, via both clipping and normalization, with no signficant gain in performance.

On the other hand, the accuracy of the gradient direction may be a more important factor, I suspect, but is not analyzed at all in the theory part.

Indeed, for some environments, accuracy can play an important role, but it is difficult to analyze in theory, since it would require accurately characterizing the training dynamics and generalization of a neural network. Instead, we provide empirical evidence to complement and confirm our theory: experiments in Section 4.1 show improved gradient accuracy for well-behaved systems, and the double-pendulum experiment in Section 4.2 demonstrates a case where using the gradient coming from a smooth but possibly inaccurate model leads to better policy optimization compared to using the real policy gradient.

This highlights a highly desirable property of using architectures such as Transformers as Action-Sequence Models. On the one hand, they are able to provide an accurate gradient when the dynamics of the environment is smooth enough; on the other hand, when the dynamics is non-smooth and thus not amenable to policy optimization, they will tend to approximate it with a smooth version of it, possibly surpassing the original dynamics in terms of ease of optimization via backpropagation through time.

It would be more interesting if the paper can discuss more complex situations, such as those that require close-loop control.

Note that our current analysis and experiments employ a closed-loop policy, but simply use an open-loop policy gradient, identical to the gradient used in state of the art model-based methods such as [2], [3] and [4]. The open-loop policy gradient does not necessarily induce an open-loop policy, and, despite its effect is yet to be well understood in the literature, it has been shown to lead to successful policies in complex tasks.

By “ill-behaved” do you mean the exponential upper bound in Corollary 1.1?

Yes, Corollary 1.1 formally establishes the possibility of exploding gradients under a Markovian model.

But Corollary 1.1 applies only to a special Markovian model, the model with linear units. What about other Markovian models?

In reality, the analysis provided by Corollary 1.1 considers a non-linear network $x_t=\sigma(W_x x_{t-1}) + W_a a_{t-1} + b$. The linearity with respect to the action quickly becomes non-linear for all future state predictions due to the recursive nature of $x_t$. The analysis can thus be extended to the more traditional non-linear formulation $x_t = \sigma(W_x x_{t-1} + W_a a_{t-1} + b)$, just as done in the seminal work by Pascanu [5].

Without tightness, an upper bound like Corollary 1.1 is not enough to support your claim that the gradient of RNN models will "explode exponentially fast".

In our paper, we only claim "the policy gradient [...] with an RNN [...] can explode exponentially fast [...]".

We agree that a tight bound is important in order for our statement to be relevant: for this reason, we have added a justification in Appendix A.1, explaining why the bound is indeed tight. The justification is similar to the ones used to explain exploding gradients in RNNs in supervised learning [1] [5].

Even though it really could be proved that the gradient of Markovian model with true transition kernel grows exponentially with the horizon length, [...] this means the unbiased policy gradient optimization is unstable and we perhaps should not use the model-based policy gradient method at all in this case.

This is a good observation. However, the fact that the policy gradient under the true transition kernel does not lead to successful policy optimization does not imply, surprisingly, that we should not use model-based policy gradients at all. Our experiments on the double-pendulum environment in Section 4.2 exactly demonstrate this point: the learned gradient from a better sequence model is highly preferable to the true gradient, and amenable to policy optimization.

how do we know that the gradient from transformer is different from an arbitrary small-but-biased gradient, in terms of its effectiveness to power the policy optimization

While we cannot prove this theoretically, we show empirically in Section 4.2 that the learned well-behaved gradients from a Transformer are indeed effective for policy optimization, leading to high-perfoming policies compared to other baselines.

Based on suggestions and feedback, we have also included additional results of an online decision transformer on the Myriad environments. We include these results in Appendix F.3. Notably, decision transformers are not suitable for a purely online setting, and thus exhibits worse performance in every task than any of our action-sequence models. In fact, a simple one-step model also outperforms the online decision transformer. Please refer to the amended Appendix for additional details.

We thank the reviewers for their initial feedback. Given that the deadline for the rebuttal period is today, we would love to hear if there is any additional answer we could provide to help the review process. Otherwise, we're hopeful that our answers addressed the main concerns from the reviewers, and hope reviewers could raise their score accordingly.

Best,

The Authors