Paper: A Differentiable Sequence Model Perspective on Policy Gradients (ICLR-2024)

Reviewer comment (Experiments):
There is some confusion about the experiments, especially on the comparison between action-sequence models and history-sequence models (which have states as conditions). The authors provide some explanations in the final paragraph of Page 8 and Page 9 but I don't think I am convinced. To me, state conditioning is necessary to predict the next states. Only conditioning on the actions does not provide complete information about the environment.

Author response:
> 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].

Reviewer comment (Novelty):
The connection between policy gradients and RNNs/sequence models seems obvious in the literature. RL policies interact with environments in a recurrent function application manner, which corresponds to RNNs/sequence models. I don't quite see what brings the novel insights from the proposed understanding.

Author response:
> 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.

Reviewer comment (Novelty):
Originality is a major concern for this submission. 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]. This paper does not clearly establish its unique contributions to the field, and the related work section is insufficiently detailed, lacking a critical analysis of how this work diverges from existing methodologies.

Author response:
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.

Reviewer comment (Experiments):
The experimental design does not effectively differentiate the proposed method from established sequence modeling algorithms. 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. Additionally, the benchmarks chosen for testing the methodology do not cover the breadth of scenarios needed to substantiate the authors' assertions. The inclusion of common continuous control benchmark tasks, such as Hopper, HalfCheetah, Walker and Antmaze, is essential for a more comprehensive evaluation.

Author response:
Thank you for your feedback! 

> 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.

Reviewer comment (Theory):
I am not sure that the theory part (Section 3) well support the claim that action-sequence models are better choice than one-step Markovian models. Only the norm of the gradients induced by two special cases of these two classes of models are compared, but in practice the estimated gradient is often normalized so the norm is less important, in my impression. 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. Also see my Question 1~4 below for several soundness concerns about the theory part.

Author response:
> 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.

Reviewer comment (Experiments):
I am not sure if the baselines in the experiment part (Section 4) are strong enough to establish the advantage of ASM over one-step models. In particular, the one-step model tested in the experiment seems to be a very simple one. See my Question 5~6 below for the detailed concerns.

Author response:
> 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.

Reviewer comment (Experiments):
The current results of this paper seem to have limited applicability scope: it seems to mainly applicable to environment that is deterministic, differentiable, with fixed episode length, where open-loop policies are sufficient for the environment. In this special case, the RL problem degenerates to a simple black-box optimization problem where we maximize an unknown but deterministic objective function over an action-sequence space. It would be more interesting if the paper can discuss more complex situations, such as those that require close-loop control.

Author response:
>  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.

Reviewer comment (Theory):
Page 4, in the paragraph below Proposition 1, you said it's a "fundamental fact" that PG with Markovian model is "fundamentally ill-behaved", and you said this fact is "analyzed in-depth in this section". I don't quite understand this sentence and am not sure about the analysis either. By "ill-behaved" do you mean the exponential upper bound in Corollary 1.1? But Corollary 1.1 applies only to a special Markovian model, the model with linear units. What about other Markovian models? In what sense can we conclude that Corollary 1.1 is due to the Markovian property, instead of to the linearity or other limits of the model under consideration?

Author response:
> 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].

Reviewer comment (Theory):
How do we know that the upper bound in Corollary 1.1 is tight? 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". To support such a claim, we typically need a *lower bound* result.

Author response:
> 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].

Reviewer comment (Theory):
Even though it really could be proved that the gradient of Markovian model with true transition kernel grows exponentially with the horizon length -- even though we suppose this were true in this question -- 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, isn't it? Importantly, although the gradient with Transformer is better bounded in this case, since we know that it's *not* the true policy gradient (because the true policy gradient has larger norm), 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?

Author response:
> 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.

Reviewer comment (Theory):
Page 5, in the paragraph below Corollary 1.1, you said "Corollary 1.1 explains both the difficulties ... and the limitations ...". What are exactly the difficulties and limitations here? Why does your upper bound result indeed *explain* them (rather than just coincident with them)? The argument here is not self-contained so it's hard to evaluate its soundness.

Author response:
> 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.

Reviewer comment (Experiments):
In Section 4.3, are the environments here partially observable? Does a state info \(s_t\) give a Markovian state or only the partial observation of the full state? 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. On the other hand, one-step Markovian models are just not appropriate for POMDP environments, so I'm not sure they should be included as baselines in this experiment.

Author response:
> 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.