Decision S4: Efficient Sequence-Based RL via State Spaces Layers
Authors: Shmuel Bar David
Itamar Zimerman
Eliya Nachmani
Lior Wolf
Venue: ICLR 2023 poster (2023)
Keywords: Sequential RL
S4
Decision transformers

Abstract: Recently, sequence learning methods have been applied to the problem of off-policy
Reinforcement Learning, including the seminal work on Decision Transformers,
which employs transformers for this task. Since transformers are parameter-heavy,
cannot benefit from history longer than a fixed window size, and are not computed
using recurrence, we set out to investigate the suitability of the S4 family of
models, which are based on state-space layers and have been shown to outperform
transformers, especially in modeling long-range dependencies. In this work, we
present two main algorithms: (i) an off-policy training procedure that works with
trajectories, while still maintaining the training efficiency of the S4 model. (ii) An
on-policy training procedure that is trained in a recurrent manner, benefits from
long-range dependencies, and is based on a novel stable actor-critic mechanism.
Our results indicate that our method outperforms multiple variants of decision
transformers, as well as the other baseline methods on most tasks, while reducing
the latency, number of parameters, and training time by several orders of magnitude,
making our approach more suitable for real-world RL

Published as a conference paper at ICLR 2023
DECISION S4: EFFICIENT SEQUENCE-BASED RL VIA STATE SPACE LAYERS
Shmuel Bar-David *t, Itamar Zimerman *t, Eliya Nachmani 1 & Lior Wolf {shmuel 1b1, b1, z me rman1] @mmail. iaaiilaau... www.ta.................
(enk100, liorwolf flf é gma com
ABSTRACT
Recently, sequence learning methods have been applied the problem off-policy Reinforcement Learning. including the seminal work Decision Transformers, which employs transformers for this task. Since transformers are parameter-heat cannot benefit from history longer than fixed window size, and not computed using recurrence, we set out to investigate the suitability of the S4 family of models, which are based on state-space layers and have been shown outperform transformers, especially in modeling long-range dependencies. In this work we present main algorithms: trajectories, while still maintaining the training efficiency of the S4 model. (ii) An on on-policy training procedure that trained recurrent manner, benefits from long-range dependencies, and is based on novel stable actor-critic mechanism. Our results indicate that our method outperforms multiple variants of decision transformers, well as the other baseline methods on most tasks, while reducing the latency, number parameters, and training time by several orders of magnitude, making approach more suitable for real-world RL.
(i) an off-policy training procedure that works with
INTRODUCTION
Robots are naturally described as being an observable state, having a multi-dimensional action space and striving to achieve a measurable goal. The complexity of these three elements, and the often diffeffeeentiable links between them, such as the shift between the states given the action and the shift between the states and the reward (with the latter computed based on additional entities), make the use of Reinforcement Learning (RL) natural, see also (Kober et al. 2013 [barz et al. 2021).
Off-policy RL has preferable sample complexity and is widely used in robotics research, e.g., (Haarnoja et al. (2018. Gu et al.) 2017). However, with the advent of accessible physical simulations for generating data, learning complex tasks without successful sample model is read- ily approached by on-policy methods (Sickmann et al. (202I) and the same holds for the task of adversarial imitation learning [Png al. (2021 2022).
The decision transformer of Chen et al. (001)) is a sequence- off- ool-cy RL method that considers sequences of tuples of the form (reward, state, action). Using the auto-regressive capability of transformers, it generates the next action given the desired reward and the current state.
The major disadvantages of the decision transformer are the size of the architecture, which is known limitation in these models, the inference runtime, which stems from the inability to compute the transformer recursively, and the fixed window size, which eliminates long-range dependencies. In this work, we propose a novel, sequence-based RL method that is far more efficient than the decision transformer and more suitable for capturing long-range effects.
The method is based on the $4 sequence model, which was designed by Gu et al. (202 la). While the original S4 method is not amendable for on-policy applications due to the fact that it designed to train on sequences rather than individual elements, we suggest new learning method that combines
"Theee authors contributed equally to this work.
Tel Aviv University
Meta AI Research

---

Published as a conference paper at ICLR 2023
off off-policy training with on-policy fine-tuning, This scheme allows to run on on-policy algorithms, while exploiting the advantages of S4. In the beginning, we trained the model in an off off-policy manner on sequences, via the convolutional view. This process exploits the ability of S4 to operate extremely fast on sequences, thanks to the fact that computations can be performed with FFT instead of several recurrent operations. Later, at the fine-tuning stage, we used an on-poliey algorithm. While pure on-policy training is a difficult task due to the instability and randomness that arise at the beginning of training, our method starts the on-poliey training a more stable point.
From the technical perspective, our method applies recurrence during the training of S4 model. As far as we can ascertain, such a capability has not been demonstrated for S4, although was part of the advantages of the earlier HIPPO Gu et al. (2020) model, which has fixed (unlearned) recurrent matrices and different parameterization and is outperformed by $4. Furthermore, in Appendix E we show that the recurrent view of the diagonal state space layer is unstable from both theoretical and empirical perspective, and we propose a method to mitigate this problem in on-policy RL. This observation provides a further theoretical explanation for why state-space layers empirically outperform RNNs. Moreover, we present a novel transfer learning technique that involves training both the recurrent and convolutional views of S4 and show its applicability for RL.
We conduct experiments on multiple Mujoco (Todorov et al. 2012) benchmarks and show the advantage of our method over existing off- aff-policy methods, including the decision transformer, and over similar on-policy methods.
2 RELATED WORK
Classic RL methods, such as dynamic programming Veinott 1966 Blackwell 1962 and Q-learning variants Schwwarz [1993]; Hasselt (2010); Rummery & Niranjan 1994) are often outperformed by
deep RL methods, starting with the seminal deep Q-learning method (Mnih et 2015) and followed by thousands of follow-up contributions. Some of the most prominent methods are AlphaGo (Silver
et al. 2016), AlphaZero (Silver et al. 2018), and Pluribus (Brown & Sandholm 2019). which outperform humans in chess, go and shogi, and poker, respectively.
Sequence Models in RL There are many RL methods that employ recurrent neural networks (RNNs), such as vanilla RNNs (Schäfer, 2008: [i et al.) 2015) or LSTMs Bakker (2001) 2007). Recurrent models are suitable for RL tasks for two reasons. First, these models are fast in inference, which is necessary for a system that operates and responds to the environment in real- salltimm. Second, since the agent should make decisions recursively based on the decisions made in the past, RL tasks are recursive in nature.
These models often suffer from lack of stability and struggle capture long-term dependencies. The latter problem stems from two main reasons: (i) propagating gradients over long trajectories is an extensive computation, and (ii) this process encourages gradients explode or vanish, which impairs the quality of the learning process. In this work, we tackle these two problems via the recent S4 layer Gu al. a02 and stable implementation of the actor-critic mechanism.
Decision transformers (Chen et al. 2021) (DT) consider RL as sequence modeling problem. Using transformers as the underlying models, state-of-the-e results are obtained on multiple tasks. DTs have drawn considerable attention, and several improvements have been proposed: Furuta et al. (2021) propose data-efficient algorithms that generalize the DT method and try to maximize the information gained from each trajectory to improve learning efficiency. Meng et al. (202T) explored the zero-shot and few-shot performance of a model that trained an offline manner on online tasks. Janner et al. (2021) employs beam search as planning algorithm to produce the most likely sequence of actions. Reid et al. (2022) investigate the performance of pre-trained transformers on RL tasks and propose several techniques for applying transfer learning in this domain. Other contributions design a generalist agent, via a scaled transformer or by applying modern training procedures (Lee et al.
2022 Reed et al. 2022 Wen et al. 2022).
The most relevant DT variant to our work is a recent contribution that applies DT in an on-policy manner, by fine-tuning. lee-med DT that was trained in off off-policy manner (Zheng et al. 2022).
Actor-Critic methods Learning off off-policy algorithms over high-dimensiona. complex data is a
central goal in RL. One of the main challenges of this problem is the instability of convergence, as
2

---

Published as a conference paper at ICLR 2023
shown in Maci al. 2009 Awate 2009 Benhamou, 2019). which is addressed by the Actor- Criti framework (Fujimoto et al. 2018: Peters & Schaal 2008 Lillicrap et al. 2015 Wang et al. 2016
Silver et al. 2014). It contains two neural components: (i) the critic that parameterizes the value function, and (ii) the actor that learns the policy. At each step, the actor tries to choose the optimal
action and the critic estimates its suitability. Similarly GANs/Goodfflllows et al. (2014), each neural component maximizes a different objective, and each provides a training signal to the other.
Over the years, several major improvements were proposed: (i) Haarnoja et al. (2018) designed relatively stable and "soft" version of the mechanism based on entropy maximization, (ii) Mnih et al. (2016) propose an asynchronous variation that achieves impressive results on the Atari benchmark, (ii) Nachum et al.] (2017) proposed the Path Consistency Learning (PCL) method, which is this work, we tailored the Actor-Critic-aasee based DPGG algorithm Lillicrap et al. (2015) to implement stable on-policy training scheme that collects data at the exploration stage and stores it in the replay buffer. This data is then used in an off-policy manner.
generalization the vanilla actor-critic system.
The S4 layer S4 (Gu et al. 202 is a variation of the time-invariant linear state-space layers (LSSL) (Gu et al. 202 which was shown capture long-range dependencies on several bench- marks such as images and videos (Islam & Bertasius 2022 Gu et al. (2022) speech (Goel et al. 2022), text Mchta et al. (2022) and more. Recently, efficient diagonal simplification was presented Gupta 2022).
3. BACKGROUND AND TERMINOLOGG
The Decision Transformer formulate offline reinforcement learning sequence modeling problem. employs causal transformers in an autoregressive manner to predict actions given previous actions, states, and desired rewards. Our method adopts this formulation, in which trajectory is viewed as a sequence of states, actions, and rewards. At each time step i, the state, action, and reward denoted by Si, di, ri, and trajectory of length L is defined by t == (so; ag Yo, -.. 81, ar, rr.)- Since the reward at some time step i can be highly dependent on the selected action a,, but not necessarily correlative futurre rewards. As result using this representation for learning conditional actions can be problematic, as planning according to past rewards does not imply high total rewards. Therefore, DT makes a critical design choice to model trajectories based on future desired rewards (returns- (eeurns--t-ggo) instead of past rewards. The returns-to- to-go at step i are defined by R, := : Ei Tv. and trajectories as the desired performance. One central advantage of this formulation is that the model is able to better plan and prioritize long-term and short-term goals based on the desired return each stage of the episode. For example, the return-to-go is highest at the beginning of the training, and the model can understand that it should plan for long-range. From an empirical point of view, the results of DT showed that this technique is effective for the autoregressive modeling of trajectories.
represented as (Ro, Sp, ag .... R,, Sr., ar.) At test time, returns-to- are fed
The S4 Layer is described in detail in Appendix A Below we supply the essential background. Given input scalar function u(t), the continuous time-invarinn state-space model (SSM) is defined by the following first-order differential equation: à Ax(() + Bu((), y(1) Cx((t) + Du(()
The SSM maps the input stream utt) to yt). It was shown that initializing A by the HIPPO matrix (Gu al. 2020) grants the state-space model (SSM) the ability to capture long-range dependencies. Similarly to previous work (Gu et al. 2021a Gupta, 2022). we interpret D as parameter-base based skip-connection. Hence, we will omit D from the SSM, by assuming
The SSM operates on continuous sequences, and it must be discretized by step size A to operate on discrete sequences. For example, the original $4 used the bilinear method to obtained a discrete approximation of the continuous SSM. Let the discretization matrices be A. B, C.
À = - - (1 (1+
B = - A1 -)- AAB
C-C
(1)
These matrices allow us to rewrite Equation 4 and obtain the recurrent view of the SSM:
3k A + BuR, yk C.r
(2)
3

---

Published as a conference paper at ICLR 2023
Decision 54
Input Progecton
54 Glock
Block
54 Brock
R4-4
Activaace
Desove
Figure 1: The architecture of the actor network X.
Since Eq. 2 is linear, it can be computed a closed form:
Ut
K K;= ČA y= K
(3)
And this is the convolution view of the SSM, when denote non-circular convolution.
METHOD
Below, we describe both an off-policy method, which motivated by DT and replaces the transformers with S4, and an on-policy method, which applies a tailored actor-critic scheme suitable for S4 layers.
4.1 OFF-POLICY
The DT formulation generates actions based on future desired returns. For a trajectory of length L, the desired returns ("returns-to-go" are denoted as R; = E Fi, where r; is the obtained reward on the step i. Formally, DT observes trajectory as ($1, (1, R1...., $L, 0L, R2, sT) sequence of length that contains states, actions, and desired total rewards at the moment. Since the transformer has quadratic space and time complexity. running is unfeasible, and trajectories were truncated to a length of 20 for most of the environments. While transformer-based variants have achieved excellent results on RL benchmarks, the $4 model is far more suitable for RL tasks for several reasons: (i) The S4 model was shown to be a long-context learner (Goel al. 2022, Islam & Bertasius 2022 Gu al. 2021a). which is relevant to RL (Wang et al. 2020) (ii) RL systems require low latency. and the S4 model is one or two orders of magnitude less computationally demanding during inference than Transformer with similar capacity. For example, on CIFAR-10 density estimation, the S4 outperforms the transformers and is 65 times faster(see Appendix Tab, 4 (iii The number of parameters is much smaller in S4 in comparison to transformers of similar level of accuracy. These are discussed in great detail in Appendix
complete trajectories
4.1.1 NETWORK ARCHITECTURE
In the offline setting, single network X is employed. It consists of three components, as shown in Fig. [ (i) The component contains three stacked S4 blocks. The internal S4 architecture follows the work of Gu et al. (2021a) and contains linear layer, batch normalization layer, GELU activations, dropout layer and residual connections. (ii) An encoding layer that encodes the current state, the previous action, and the previous reward. Each encoder is one fully-connected layer with ReLU activations. (iii) The input and output projection layers, each which contains a single fully-connected layer and ReLu activations. The encoder-decoder components, such the number of S4 blocks follows the DT architecture.
We used the S4 diagonal parameterization GGupta 2022 and set the state size N = 64, similarly to (Gupta 2022 Gu et al. 2021a), and the number of channels as H = 64, half of the DT embedding dimension, since some $4 parameters complex and consume twice as much memory.
The training scheme of deep reinforcement learning via sequence modeling was extensively ex- plored (Zheng al. 2022 Kostrikov et al. 2021 Furuta al. 2021 Chen et al. 2021). Unfor-
4

---

Published as a conference paper at ICLR 2023
tunately, training S4 via this scheme is inefficient, because it involves the S4 recurrent view instead the convolutional view (3) Hence, we modify the training scheme to operate on complete trajectories instead of steps.
4.1.2 TRAINING SCHEME
A sketch of our training scheme is shown in Alg. |
For simplicity, the sketch ignores batches.
Algorithm Off-yy Training Scheme
1: X (8;, R2-1,0 01-1 0x) Init() D Init the network X according the HIPPO initialization
2: for trajectory dataloader
3: for i E (0..... T.] do
b Assuming R_1 = Ry 0-1 0
4: á; = X (5;, Ri-1. a;-1)
5: loss a. a,l"
6: Take gradient step on 0x according to the loss
7: end for
8:
end for
At each training step batch of trajectories is randomly drawn from the dataset. The trajectories pad with zeros reach the same length. Given the current state 8;, the previous return-to-ge Rt-1 and the previous action 0;-1, the network X predicts the next action d; The loss is calculated by applying the L2 loss over batch of true actions ag and the predicted ones di.
In comparison to the DT, our scheme operates on complete instead of truncated trajectories. These lead to a much more data-efficient approach that can handle long-context, which, as shown by our experiments, is important for optimal performance. This technique exploits the strengths of S4, and relevant to any off-pollit policy task. For each task, we select the initial return-to-go R_ from the values 0.5, 0.75, 1.0, choosing the one that results in the best performance.
4.2 ON-POLICY FINE-TUNING
For on-policy fine-tuning, we present modified variant of the DDPG actor-critic algorithm (Lillicrap al. 2015) that includes networks: The actor target X, which has the S4-based architecture mentioned above, (ii) the actor copy X" (iii) the critic target C, and (iv) the critic copy C². These actor and critic target networks, using an update
copy networks are used to achieve soft updates factor that determines how fast models updated.
The updated mechanism can be interpreted as using an average of polices. This approach is taken in order to deal with the instability of the actor-critic mechanism, which is well-known phenomenon, see Sec. 2 Overall, our variant of DDPG has the following properties.
(i) Training The Critic To boost stability at the beginning of training. first freeze the actor and train the critic network alone on the recorded data. Since the purpose of the critic to estimate the Q-function
the policies generated by the actor, train the critic using the recorded data and not the original dataset. Additionally, we observe that training the actor and critic together from the start increases instability. Both the training of the actor and critic use the replay-buffer to sample batches of transitions. There were instabilities because the actor was changing faster compared the critic, and also because the critic is not accurate enough before sufficient training iterations when producing gradients for the actor. Those problems that appear when training both models simultaneously required to change the training method. Hence, we trained the actor and critic intermittently, similarly to GANS. In addition, we used a relatively lower learning rate for the actor networks, to prevent situations where the critic falls behind the actor and generates irrelevant reward estimates.
(ii) Stable Exploration Another challenge we encountered was maintaining stability while per- forming efficient exploration. Similarly to other contributions, we limit exploration during training. After the initialization the critic, the noise is sampled from a normal distribution N (0, depisode). where Cepisode linearly decays over the training from Additionally, relatively low learning rates are used for all networks: specifically, the critic network used a learning rate of ac, and ax for the actor.
5

---

Published as a conference paper at ICLR 2023
(iii) Freezing The S4 Kernel On-policy training requires additional regularization, and we reg- ularize the S4 core by freezing A. Moreover, in the initial off-po training, the S4 kernel learns to handle long history via the conventional view, whereas this ability can be lost when training the model via the recurrent view. See Appendixes. and for further discussion of the role these two views play.
Setting Return-to-Go At the exploration stage, the trained model produces a replay buffer that contains the explored trajectories. We set the initial return-to-go R_ value of 10% higher than the model's current highest total rewards. The initial highest result is estimated by running the trained model with a normalized target of 1.0. This technique allows us adjust the initial return-to- go-go to the recorded data without additional Syype--arameters.
Our on-policy fine-tuning scheme is summarized in Alg. is based on the Deep Deterministic Policy Gradient (DDPG) algorithm (Lillicrap et al. 2015 As on-policy algorithm, it contains exploration and exploitation stages. At exploration, the actor collects trajectories and stores them in replay buffer. We apply the actor-critic framework to these trajectories to narrow the training instability.
Algorithm 2 On-Policy Training Scheme
Copy the model weights to the actor target network X'(8s; R2-1, 0:-1 0x:)
2: Initialize C(s;, 4;-1 0c) and copy C'(s;, 0:-1 (c,)
3: Train the critic target network (s,, a;-1 0c) with respect to X(s;, Ri 11-1.00-1 0x)
4: Apply our on-policy tailored DDPG algorithm in iterations:
5: for episode = 1, .M do
6: for each step in the environment i do
7: Sample noise wi (O, (episode) where Repisode
M Meeepiodee
8: Define action (₁ == X ((s, R4-1, a;-1 0xx Wt
9:
Play a, and observe the current reward r; and the next state 9;+1
10:
Store ((s;, R.-1)) ri, ($s+1, 4, R;)) in the replay buffer D.
Every K₃ steps:
12:
Train the critic network C' with relatively high learning rate.
[(), a,- R;-1), ($+41, a, R&)} C Calculate Bellman target estimated returns r; + ('i11 ('(+11,, Ri, ai))
Sample batch from replay-buffer
14:
Update gradient descent for critic: Oc + Oc ac oz (C(si, (());;a)) y)²
16:
If actor freezing is over, update gradient ascent for actor:
17:
0x + 0x + ax Vax
R,,,,,,1,,411 a.-))
18:
Every K₃ steps:
19:
Soft update the critic target C' from C, and the actor target X' from X using 7
20:
Update the critic target 8c (1 7)6 70c
21:
Update the actor target 0'x + (1 7)0 + 50x
22:
end for
23:
end for
Critic Architecture
The architecture of the actor networks X, X' is the same as the one depicted Fig. Three fully-connected layers with ReLU activations are used for the critic networks C, C³.
5 EXPERIMENTS
We evaluate our method on data from the Mujoco physical simulator (Todorov et al.) 2012) and AntMaze-v2 (Fu et al. 2020). Additional details about the benchmarks are presented in Appendix C.I To train the model, we used D4RL (Fu et al. (2020) datasets of recorded episodes the environments. We trained these datasets in the form of off-policy, prquence-to-squueee
supervised learning.
After training the model, we move on to on-policy training, by running the agent on the environment, while creating and storing in the replay buffer its states, actions, rewards, returns-to and also the state of the S4 modules, including those of the next step. A single input for the model can be denoted 2; (8i, Oj, Rz), where includes both the state provided by the environment, and the states of

---

Published as a conference paper at ICLR 2023
Table 1: Normalized Reward obtained for various methods on the Mujoco benchmark. The first segment of the table contains off-line results for non-DT methods, the second segment for DT methods and ours. The bottom part shows results for on-policy fine-tuning. The rows with the present the difference from the corresponding row in the off-line results. Bold denotes the model that empirically best for that benchmark.
Dataset:
Hopper
HalfCheetah
Walker2D
Mean
Method
Medium Replay
Expert
Medium Replay
Expert
Medium Replay
Expert
Rank
IQL
Kostrikov et al.
2021
63.81
92.13
91.5
47.37
44.1
86.7
79.89
73.67
109.6
3.66
CQL
(Kumar el al.
2020)
58.0
48.6
111.0
44.0
46.2
62.4
79.2
26.7
98.7
5.22
BEAR
(Kumar
2019
52.1
33.7
96.3
41.7
38.6
53.4
59.1
19.2
40.1
8.22
BRAC
Wu
0.8
47.7
41.9
81.1
0.9
81.6
6.88
Transformer- based methods:
DT
Chen et al.
2021)
67.6
82.7
107.6
42,6
36.6
86.8
74.0
66.6
108.1
5.55
CDT
Funtta al.
2021
46.6
80.7
52.6
87.6
71.0
106.7
5.66
ODT
Zheng et al. 2022
66.95
86.64
42.72
39.99
72.19
68.92
6.16
Sample-base based methods:
TT. m
Janner et al.
2021
67.4
99.4
106.0
44.0
44.1
40.8
81.3
79.4
91.0
4.33
TT,
Janner al.
2021
61.1
91.5
41.9
95.0
79.0
82.6
101.9
3.88
DS4 (ours)
DS4
89.47
87.67
110.52
47.31
43.81
94.75
81.36
80.26
109.56
2.44
On-policy fine-tuning:
IQL (Kumar 2020
66.77 79
96.23
47.1
44.14
80.33
70.55
for IQL
2.98
4.1
0.04
0.04
0.44
-3.12
ODT (Zheng et al.
2022
97.54
88.89
42.72
39.99
72.19
68.92
2.66
8 for ODT
30.59
2.25
0.43
-0.54
4.60
7.94
DS4
111.3
88.42
47.32
45.17
81.71
76.97
1.33
DS4
0.03
-0.18
0.26
0.35
-4.44
the S4 kernel. A single replay buffer object can be expressed as (2i, Ti, 2141). detailed description of our experimental setup can be found in Appendix C.2
Results On Mujoco Tab. lists the results Mujoco obtained for our method and for baseline methods, both D'T-based and other state-of-the-arr methods. As can be seen, the our method out- perform most other methods most benchmarks, obtaining the lowest average rank. Compared to DT, which employs sequences in a similar way, we reduce the model size by 84%, while obtaining consistently better performance.
Furthermore, our method often outperforms TT, and TT, which are cumple-based methods that employ beam search optimize the selected actions over past and future decisions. During inference. these methods run the actor network PW times, where P is the planning horizon and W is the Beam Width, and choose the action that results the optimal sequence of actions. The results we cite employ P = 15, W = 256 and the run run-time is 17sec in comparison to 2.1ms (see Tab. 4 of our method (4 orders of magnitude).
On-policy results are given for medium and medium-replay. We observe sizable improvement in hopper-medium, and slight improvement in the other environments, This can be a result of the diminishing returns effect, i.c., the performance of the hopper-medium model was close to the maximal oeturn-to-go already, and the model generated trajectories that are close to this expected reward. In contrast, for the half-cheetah and walker2d, the model finishes the off-policy training phase with rewards significantly lower than out target.
An interesting phenomenon is that for the off-policy performance comparisons, algorithm appears to be the best when used for the "expert" datasets, rather than "medium" or "medium replay". We hypothesize that one of the main differences between "expert" and "medium" players the ability
7

---

Published as a conference paper at ICLR 2023
0.6
1.0
Rof-cheetah
HaPtteetan eqptt
110
13
1.0
00
maduumnsplyy
0.7
0.6
0.5
05
0.4
24
0.
Walser
02
02
Happer mediu
Whiter medium
81
0.
Maaper mediumorppyy
Walter med reedum-nptyy
0.0
0.0 01
02
03
04
05 06 57 05 100
0.5 0. 02 53 D4 05 06 07 ba 05 1.0
0. 01 6. 0.3 0.4 26 DJ7 0.8 Roadie longhh from bojectoy
Beaahee mostmmm langhn Dugpdaryy
Resative maximum longth takon Dayedory
(a) HalfCheetah
(b) Hopper
(c) Walker2D
Figure 2: The effect of limiting the number of recurrent steps while training model. Maximum mean rewards achieved are presented per environment: (a) HalfCheetah, (b) Hopper (c) Walker2D
to plan for long-range goals. Due to this, the S4 model can fulfill its potential much better in the "expert" case. This hypothesis is supported by Fig. 2 which shows that when the model is enabled to use only a short context (0.1 of the maximal context length), all the levels achieved roughly the same performance.
Results On AntMaze Results on AntMaze shown in Tab. Evidently, in both of the tested environments, DS4 outperforms the other offline method by a large margin. even outperforms methods that perform online-finetuning (which did not run yet for DS4 this benchmark).
Table 2: Normalized Reward obtained for various methods on the AntMaze-v2 benchmark. The baseline results are from (Zheng et al. 2022). Methods that use online- fiine-finetunning are denoted by star (+).
Method:
Transformer-Basse
IQL.
Ours
Dataset
DT ODT
ODT*
IQL
IQL*
DS4
AntMaze-umaze (Kostrikov et al.
2021
53.3
53.1
88.5
87.1
89.5
90.0
AntMaze-diverse (Kumar et al. 2020
52.5 50.2
56.0
64.4
56.8
79.1
Long-Range Dependencies and The Importance History We also study the importance of S4's ability to capture long-term dependencies. For this, allow the gradients propagate only k steps backward. Put differently, the gradients are truncated after recursive steps.
Fig. presents the normalized reward obtained for various values of k. given percentage of the maximal trajectory length. Evidently, in all tasks and all off- off-policy training datasets, as the maximal number of recursive steps drops, performance also drops. This phenomenon is also valid for DT (section 5.3 of (Chen et al.) 2021 ),, supporting the notion that shorter contexts affect the results negatively. Since these benchmarks are fully observable,
might wonder what is the benefit of using any history. The authors of DT hypothesize that the model policy consists of distribution of policies and the context allows the transformer determine which policy results in better dynamics.
We would like to add two complementary explanations: (i) The function of recurrent memory is simply memorize history. It also allows the model to make future decisions, For example, assume that during evolution a cheetah robot is overturned in way that has never been seen before on the training data. smart policy would be to try several different methods, one by one, until the robot overturned again. Namely, recurrent memory can extend the hypothesis class from single to sequential decisions, and the result can improve the robustness of the model. (ii) Even in cases where history does not provide any advantage in terms of expressiveness, is not clear that the same is true in terms of optimization. Even though all relevant information is included in the state s;, this does not mean that the model can extract all this information. Recurrent models process the previous states in way that may make this information more explicit.
Comparisons To RNN At test time, as a special linear RNN with a transition matrix A. To measure the impact of the critical S4 components, compare model to RNN and practically ablate various components, such as
the recurrent view of S4, and the model functions
8

---

Published as a conference paper at ICLR 2023
2.2
Nomalldd toratt
04
Hod
Expe1t
Figure 3: A comparison between S4 and RNN on the HalfCheetah dataset.
HIРРО initialization and the convolutional view (see Sec. B for discussion of these components). For fair comparison, only the S4 core model is replaced with RNN, maintaining the DT scheme (for example, using return-to-go), following the technique of training on complete trajectories, and using all non-S4 layers in the internal S4 block (Normalization layers, skip-contentions, and more).
We run these experiments on all the levels of HalfCheetah. The results are depicted in Fig. Evidently, S4 outperforms RNN in all expertise levels, especially expert datasets. Learning from an expert dataset allows the S4 model to fulfill its potential and learn long-range planning, a skill that cannot be learned at the same level by RNN.
Low-Budget To examine the performance of our model under low-budget scenarios we conducted several additional experiments on all levels of Halfcheetah. While our regular version used 64 channels (H) and an S4 state size of 64 (N), we analyzed how decreasing N and H affect the results.
Envirooments
(N=64.H=64)
(N=32,H=64)
(N=16,H=64)
(N=64,H=32)
(N=64,H=16)
(N=32,H==2)
DT
Medium
47.31
47.01
47.42
47.26
47.15
47.12
42.6
Replay
43.80
43.74
43.61
40.71
25.97
39.83
36.6
Export
91.75
91.03
91.41
91.27
93. 12.12
94.48
86,8
Parameters (Ours)
% Parameters (DT)
0.15
0.13
0.12
0.07
0.04
0.06
1.0
Table 3: Results of smaller models HalfCheetah. Each of the smaller models is denoted by (i) N the S4 state size, and (ii) the number channels.
Tab. depicts the results for low-budget models. be seen that for all levels, even small models (H 32) outperform DT, despite using 93% fewer parameters. In addition, our method doesn't seem sensitive hyperparameter choices on these tasks. The only drastic decrease in performance occurs when using a tiny model (N 64, H 16), which uses around 27K parameters. Additionally, it can be seen that decreasing the S4 state size from 64 to 32 or does not affect results, whereas decreasing H from 64 to 32 16 negatively affects results in some cases (for example, in the replay environment).
CONCLUSION
We present evidence that S4 outperforms transformers as the underlying architecture for the (re- ward,statte,action sequence-basee based RL framework. S4 has less parameters than the comparable transformer, is low-latency, and can be applied efficiently trajectories. We also extend the DDPG framework to employ S4 in on- on-policy RL. While we were unable to train on-policy from scratch, by finetuning offline policies, S4 outperforms the analogous extension of decision transformers.

---

Published as a conference paper at ICLR 2023
ACKNOWLEDGMENTS
This project has received funding from the European Research Council (ERC) under the European Unions Horizon 2020 research, innovation program (grant ERC CoG 725974). This work was further supported by grant from the Tel Aviv University Center for Al and Data Science (TAD). The contribution of IZ is part of a Ph.D. thesis research conducted at Tel Aviv University.
REFERENCES
Ugurkan Ates. Long-term planning with deep reinforcement learning on autonomous drones. In 2020 Innovations in Intelligent Systems and Applications Conference (ASYU), 1-6. IEEE, 2020.
Yogesh Awate. Algorithms for variance reduction in policy-gradies based actor-critic framework. In 2009 IEFE Symposium on Adaptive Dynamic Programming and Reinforcement Learning, pp- 130-136. IEEE, 2009.
Bram Bakker. Reinforcement learning with long short-term memory. Advances in neural information processing systems, 14, 2001.
Bram Bakker. Reinforcement learning by backpropagation through an Istm model/critic. In 2007 IEEE International Symposium on Approximate Dynamic Programming and Reinforcement Learning, pp. 127-134. IEFE, 2007.
Eric Benhamou. Variance reduction in actor critic methods (acm). arXiv preprint 17X1v:190..09765 2019.
David Blackwell. Discrete dynamic programming. The Annals of Mathematical Statistics, pp. 719-726, 1962.
Noam Brown and Tuomas Sandholm. Superhuman ai for multiplayer poker. Science, 365(6456): 885-890, 2019.
Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. Advances in neural information processing systems, 34, 2021.
Aidan Curtis, Minjian Xin, Dilip Arumugam, Kevin Feigelis, and Daniel Yamins. Flexible and
efficient long-range planning through curious exploration. In International Conference on Machine Learning, pp. 2238-2249, PMLR, 2020.
Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4rl: Datasets for deep data-driven reinforcement learning. arXiv preprint arXiv:2004.07219, 2020.
Scott Fujimoto, Herke Hoof, and David Meger. Addressing function approximation error in actor-
critic methods. In International conference on machine learning, pp. 1587-1596. PMLR, 2018.
Hiroki Furuta, Yutaka Matsuo, and Shixiang Shane Gu. Generalized decision transformer for offline hindsight information matching. arXiv preprint arXiv:2111.10364, 2021.
Karan Goel, Albert Gu, Chris Donahue, and Christopher Ré. It's raw! audio generation with state-space models. arXiv preprint arXiv:2202.09722 2022.
David Goldberg. What every computer scientist should know about floating-point arithmetic. ACM computing surveys (CSUR), 23(1):5-48, 1991.
Ian Goodfellow, Jean Poudet-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets. Advances neural information processing systems, 27, 2014.
Albert Gu, Tri Dao, Stefano Ermon, Atri Rudra, and Christopher Ré. Hippo: Recurrent memory with optimal polynomial projections. Advances Neural Information Processing Systems, 33: 1474-1487, 2020.
10

---

Published as a conference paper at ICLR 2023
Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces, arXiv preprint arXiv:2111.00396, 2021a.
Albert Gu, Isys Johnson, Karan Goel, Khaled Saab, Tri Dao, Atri Rudra, and Christopher Ré. Combining recurrent, convolutional, and continuous-time models with linear state space layers. Advances in Neural Information Processing Systems, 34, 2021b.
Shixiang Gu, Ethan Holly, Timothy Lillicrap, and Sergey Levine. Deep reinforcement learning for robotic manipulation with asynchronous off-policy updates. In 2017 IEEE international conference on robotics and automation (ICRA), pp. 3389-3396. IEEE, 2017.
Ankit Gupta. Diagonal state spaces are as effective as structured state spaces. arXiv preprint orXiv:2203.14343 2022.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. In International conference on machine learning, pp. 1861-1870. PMIR, 2018.
Hado Hasselt. Double q-learning. Advances neural information processing systems, 23, 2010.
Nicholas Higham. Accuracy and stability of numerical algorithms. SIAM, 2002.
Chia-Chun Hung, Timothy Lillicrap, Josh Abramson, Yan Wu, Mehdi Mirza, Federico Carnevale, Arun Ahuja, and Greg Wayne. Optimizing agent behavior over long time scales by transporting value. Nature communications, 10(1):1-12, 2019.
Julian Ibarz, Jie Tan, Chelsea Finn, Mrinal Kalakrishnan, Peter Pastor, and Sergey Levine. How to train your robot with deep reinforcement learning: lessons we have learned. The International Journal of Robotics Research, 40(4-5):698-72)) -21, 2021.
Md Mohaiminul Islam and Gedas Bertasius. Long movie clip classification with state-space video models. arXiv preprint arXiv:2204.0166 2022.
Michael Janner, Qiyang Li, and Sergey Levine. Offline reinforcement learning as big sequence modeling problem. Advances in neural information processing systems, 34, 2021.
William Kahan. Pracniques: further remarks on reducing truncation errors. Communications of the ACM, 8(1):40, 1965.
Diederik P Kingma and Jimmy Ba. Adam: method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014.
Jens Kober, J Andrew Bagnell, and Jan Peters. Reinforcement learning in robotics: A survey, The International Journal of Robotics Research, 32(11):228811274 2013.
Ilya Kostrikov, Ashvin Nair, and Sergey Levine. Offline reinforcement learning with implicit q-learning. arXiv preprint arXiv:2110.06166 2021.
Aviral Kumar, Justin Fu, Matthew Soh, George Tucker, and Sergey Levine. Stabilizing off-policy l-leannin via bootstrapping error reduction. Advances in Neural Information Processing Systems, 32, 2019.
Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. Advances Neural Information Processing Systems, 33:1 179-1191, 2020.
Kuang-Huei Lee, Ofr Nachum, Mengjiao Yang, Lisa Lee, Daniel Freeman, Winnie Xu, Sergio Guadarrama, Ian Fischer, Eric Jang, Henryk Michalewski, et al. Multi-game decision transformers. arXiv preprint arXiv:2205.1222 2022.
Xiujun Li, Lihong Li, Jianfeng Gao, Xiaodong He, Jianshu Chen, Li Deng. and Ji He. Recurrent reinforcement learning: hybrid approach. arXiv preprint 1rXiv:150.0004: 2015.
Timothy P Lillicrap, Jonathan J Hunt, Alexander Pritzel, Nicolas Heess, Tom Erez, Yuval Tassa, David Silver, and Daan Wierstra. Continuous control with deep reinforcement learning. arXiv preprint arXiv: /509 2015.

---

Published as a conference paper at ICLR 2023
Hamid Maei, Csaba Szepesvari, Shalabh Bhatnagar, Doina Precup, David Silver, and Richard Sutton. Convergent remporal-differene learning with arbitrary smooth function approximation. Advances in neural information processing systems, 22, 2009.
Harsh Mehta, Ankit Gupta, Ashok Cutkosky, and Behnam Neyshabur. Long range language modeling via gated state spaces. arXiv preprini arXiv:2206. 13947, 2022.
Linghui Meng, Muning Wen, Yaodong Yang, Chenyang Le, Xiyun Li, Weinan Zhang, Ying Wen, Haifeng Zhang, Jun Wang. and Bo Xu. Offline pre-trained multi-agent decision transformer: One big sequence model conquers all starcraftii tasks. arXiv preprint 12Xiv:2/12002844 2021.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Mare G Bellemare, Alex Graves, Martin Riedmiller, Andreas Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning. nature, 51 518(7540):529-333 533, 2015.
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-1937. PMLR, 2016.
Ole Moller Quasi coubleepprecisonn in floating point addition. BIT Numerical Mathematics, 5(1): 37-50, 1965.
Jean-Michel Muller, Nicolas Brisebarre, Florent De Dinechin, Claude-Pierre Jeannerod, Vincent Lefevre, Guillaume Melquiond, Nathalie Revol, Damien Stehlé, Serge Torres, al. Handbook of floating-pone point arithmetic. Springer, 2018.
Ofir Nachum, Mohammad Norouzi, Kelvin Xu, and Dale Schuurmans. Bridging the gap between value and policy based reinforcement learning. Advances in neural information processing systems, 30, 2017.
J Oliver. Relative error propagation in the recursive solution of linear recurrence relations, Numerische Mathematik, (4):323- 340, 1967.
Xue Bin Peng, Ze Ma, Pieter Abbeel, Sergey Levine, and Angjoo Kanazawa. Amp: Adversarial motion priors for stylized physics- character control. ACM Transactions on Graphics (TOG), 40(4):1-20, 2021.
Xue Bin Peng, Yumrong Guo, Lina Halper, Sergey Levine, and Sanja Fidler. Ase: Large-scale reusable adversarial skill embeddings for physically simulated characters. arXiv preprint arXiv:2205.01990 19900 2022.
Jan Peters and Stefan Schaal. Reinforcement learning of motor skills with policy gradients. Neural networks, 21(4):682-61 2008.
David Raposo, Sam Ritter, Adam Santoro, Greg Wayne, Theophane Weber, Matt Botvinick, Hado van Hasselt, and Francis Song. Synthetic returns for long-term credit assignment. arXiv preprint arXiv:2102. 12425, 2021.
Scott Reed, Konrad Zolna, Emilio Parisotto, Sergio Gomez Colmenarejo. Alexander Novikov, Gabriel Barth-Maron, Mai Gimenez, Yury Sulsky, Jackie Kay, Jost Tobias Springenberg, et al. A generalist agent. arXiv preprint arXiv:2205.06175. 2022.
Machel Reid, Yutaro Yamada, and Shixiang Shane Gu. Can wikipedia help offline reinforcement learning? arXiv preprint ftfi:2201.11222.. 2022.
Gavin A Rummery and Mahesan Niranjan. On-line Q-learning using connectionist systems, vol- ume 37. Citeseer, 1994.
Anton Maximilian Schäfer. Reinforcement learning with recurrent neural networks. arXiv preprint, 2008.
Anton Schwartz. A reinforcement learning method for maximizing undiscounted rewards. In Proceedings of the tenth international conference on machine learning. volume 298, pp. 298-305, 1993.
12

---

Published as a conference paper at ICLR 2023
Jonah Siekmann, Kevin Green, John Warila, Alan Fem,, and Jonathan Hurst. Blind bipedal stair traversal via sim-to-real reinforcement learning. arXiv preprint arXiv:210S.08328, 2021.
David Silver, Guy Lever, Nicolas Heess, Thomas Degris, Daan Wierstra, and Martin Riedmiller. Deterministic policy gradient algorithms. In International conference on machine learning, pp. 387-395, PMLR, 2014.
David Silver, Aja Huang, Chris J Maddison, Arthur Guez, Laurent Sifre, George Van Den Driessche, Julian Schrittwieser, Ioannis Antonoglou, Veda Panneershelvam, Mare Lanctot, et al. Mastering the game of go with deep neural networks and tree search. nature, 529(75877:484-489, 2016.
David Silver, Thomas Hubert, Julian Schrittwieser, Ioannis Antonoglou, Matthew Lai, Arthur Guez, Mare Lanctot, Laurent Sifre, Dharshan Kumaran, Thore Graepel, et al. A general reinforcement learning algorithm that masters chess, shogi, and go through self-play. Science, 362(6419): 1140-1144 114,, 2018.
Yi Tay, Mostafa Dehghani, Samira Abnar, Yikang Shen, Dara Bahri, Philip Pham, Jinfeng Rao, Liu Yang. Sebastian Ruder, and Donald Metzler. Long range arena: benchmark for efficient transformers. arXiv preprint 1.0iv:20011000000 2020.
Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: A physics engine for model-base control. In 2012 IEEE/RSJ international conference on intelligent robots and systems, pp. 5026-5033. IEEE, 2012.
Arthur F Veinott. On finding optimal policies in discrete dynamic programming with no discounting. The Annals of Mathematical Statistics, 37(5):1284-1294 1966.
Ruosong Wang. Simon S Du, Lin F Yang, and Sham M Kakade. Is long horizon reinforcement learning more difficult than short horizon reinforcement learning? arXiv preprint urXiv:2005.005277 2020.
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 arXiv:1611.01224 2016.
Muning Wen, Jakub Grudzien Kuba, Runji Lin, Weinan Zhang, Ying Wen, Jun Wang, and Yaodong Yang. Multi-agent reinforcement learning is sequence modeling problem. arXiv preprint arXiv:2205. 14953, 2022.
Yifan Wu, George Tucker, and Ofir Nachum. Behavior regularized offline reinforcement learning. CoRR, Abs//911.1111 1361, 2019. URL
11361
Linyi Yao, Qiao Dong, Jiwang Jiang, and Fujian Ni. Deep reinforcement learning for long-term pavement maintenance planning. Computer-Aided Civil and Infrastructure Engineering, 35(11): 1240-1245, 2020.
Qinqing Zheng, Amy Zhang, and Aditya Grover. Online decision transformer. arXiv preprint arXiv:2202.05667, 2022.
13

---

Published as a conference paper at ICLR 2023
A BACKGROUND
A.1 THE S4 LAYER
The S4 layer Gu et al. 2021aa is a variation of the fime-invarian linear spate layers (LSSL) Gu
al. 2021b). Below we provide the necessary background on this technology.
The State-Space Model state-space model (SSM) defined by the following first-order differential equation:
Given a imput scalar function utt) R + R. the continuous inverinvariant
á = Ax(t) + Bu(t), y(t) = Cr(() Du(t)
(4)
The model maps the input stream u(t) to y(u). It was shown that initializing A by the HIPPO matrix (Gu et al.) (020) grants the state-space model (SSM) the ability capture long-range dependencies. Similarly to previous work (Gu et al., 202Ta Guptal 2022), we interpret D as parameter-based skip-connection. Hence, we will omit 1) from SSM, by assuming
From Continuous to Discrete SSM The SSM operates on continuous sequences, and it must be discretized by a step size to operate on discrete sequences. For example, the original S4 used the bilinear method to obtained discrete approximation of the continuous SSM. Let the discretization matrices A, B,C.
= (1 A/2
A),
B= (1 (1- AB C=C
(5)
These matrices allow us rewrite Equation
Jk A Bug,
Uk
CaR
This is the recurrent view of the SSM, and it can interpreted as a linear RNN in which the state-transition matrix, and B, C are the input and output matrices.
Convolutional View The recurrent SSM view is not practical for training over long sequences, cannot be parallelized. However, the linear time-invariant SSM can be re-written as a convolution, which can be trained much more efficiently. The $4 convolutional view is obtained follows:
Given sequence of scalars u = (un, 41, u1,... of length L, the S4 recurrent view can be unrolled the following closed form:
Vi E LL - 1] E RN, 2o = Buo,
I = ABuo + Buj
...
LL-1
1
1=0
L-1 ČA
Vi E.R. 4o CBug: CABu, ČBu,
IL-1
Bu
i=0
Where N the state size.
Since the recurrent rule is linear, it can be computed in a closed form with matrix multiplication or non-circular convolution:
CB
0
0
30
uo
ČÁB
ČB
0
0
U1
ČAB
ČB
0
(8)
CBB
CAB
CB ČAB
y1-11
CAL-IB
CAL-2B
CB
ie., y Te y for some kernel k, which is convolutional kernel.
14

---

Published as a conference paper at ICLR 2023
From Theory to Deep Neural Network The parameters of each channel of the SSM A, B, C, A, and on the forward path A, B, C are computed according to Equation | Then, one can the recurrent view or the convolutional view in the forward path. The SSM is wrapped with standard deep learning components, such as skip-connections, GeI.U activations, and nor- malization layers, and they define the S4 layer(for one channel). While the S4 layer operates on multi-dimensional sequences, the calculations in Equations 6 operate on scalars. The S4 layer handles this by stacking multiple copies of the 1-D layer and using a linear layer to mix information between channels on the same position. To calculate k efficiently, structured (Gu et al. (021a) la) and diagonal (Gupta 2022) versions of the state-space model were proposed. We use the latter, which we found to be more stable in our experiments.
B MOTIVATION
S4 has few advantages in comparison to transformers, which make especially suited for solving RL tasks that are viewed as sequences.
Long horizon Reward Planner Designing RL algorithms for long-term planning is central goal in RL, which is extensively explored (Ates, 2020 Curtis et al., 2020; Yao et al. 2020 Raposo et al.
2021 Hung al., 2019). The S4 model was shown, both empirically and theoretically, to excel as a long-context learner. Empirically, S4 achieved SOTA results on every task from the Long Range Arena (Tay et al.) 2020), which is benchmark that consists of several long-context scenarios (images, text, structured data and more). In these experiments, S4 was able to reason over sequences of up to 16, 000 elements. Theoretically. the S4 model includes two well- rell-groudded features that allow it to capture long-range dependencies: (1) Convolutional View. The model can be trained via the convolutional view instead of the recursive one. While the recursive view includes repeated multiplications by A (Fquation 2). which can result in vanishing gradients, the convolutional view avoids these caleulations by using a convolution kernel k, which can be calculated efficiently and in a stable manner. For example, the original S4 does not compute k directly; instead, it applies the inverse FFT on the spectrum of k, which is calculated via Cauchy kernel and the Woodbury Identity. This is possible due to the Normal Plus Low Rank (NPLR) parameterization of A. Furthermore, in contrast to the convolutional view, when using the recurrent view, errors can accumulate during computations, as shown by Oliver (967)). In Appendix E we prove (theoretically) and demonstrated (empirically) that this phenomenon holds for state-space layers, and provide a technique for mitigating those errors during on- on-policy fine-tuning. (2) The HIPPO Initialization. As presented in (Gu et al., 2020) the HIPPO matrix, which defines the S4 initialization and the NPLR- PPLL-parramttrization of the transition matrix A, is obtained from the optimal solution of an online function approximation problem. Due to this initialization, the recurrent state is biased toward compressing the representation of the input time series (the history) into a polynomial base. The compressed representation allows the model to handle long history, even with limited state size and relatively few parameters.
To empirically measure the importance of these features for RL tasks, we compare our model to RNN, as described in 3 This ablation study provides empirical support for our main hypothesis, that S4 can be highly relevant for RL tasks, especially in long-range planning. Note that the convolutional view is relevant only for off- off-policy reinforcement learning. while the HIPPO initialization is relevant to all RL tasks.
Complexity Real-time RL systems require low latency, and the S4 model is one or two orders of magnitude less computationally demanding during inference than a Transformer with similar capacity. For example, on CIFAR-10 density estimation, the S4 outperforms the transformers and is 65x faster. We validate that this phenomenon extends to the RL world, and observe that our method is around 5x faster than DT, as described in Tab. 4 This phenomenon is caused by the fact that the time complexity of the S4 recurrent view does not depend on the sequence length L, in contrast to the DT, which is dominated by the L factor necessary for self-attention layers. Furthermore, in terms of space complexity, S4 has a critical and promising advantage for real-time real-time and low-budget robots, especially in long-range credit assignments: Note that for these tasks, DT used the entire episode length as the context (section 3.4 in (Chen et al. 2021) ) which resulted in O(L?) space complexity, making it impractical for long episodes, in contrast to S4, which does not depend on L..
15

---

Published as a conference paper at ICLR 2023
Model
DS4
DTTkk 1)
DIC = 20)
DITK = 100)
DTV = 250)
DTÍ = 300)
TT, TT. 17K
Time (rus) 2.1
2.3
3.1
5.1
10.5
18.6
Table 4: Wall clock running time for step for each model. For DT, we measured the running time for several values of k, which is the context length the model takes into account. For TT, we measured two models TTg and TT employing the peremeerrseetess that used for Mujoco in Janner et al. (202T). As can be seen, for practical tasks that require long context, S4 can be much more efficient. All experiments were conducted on HalfCheetah environment.
Furthermore, the space and time complexity for the training step of s4 is better than transformers. This gap arises from the fact that the complexity of self-attention layers is quadratic in I,, while the S4 space complexity depends only on I. and the time complexity in log L. The difference in time complexity arises from the efficiency of the S4 convolutional view, which can be calculated with FFT, compared to the transformer, which requires quadratic complexity. Moreover, at training, the S4 does materialize the states, causing the time and space complexity to be much more efficient, and, in practice, invariant to the state size.
Parameter Efficiency The number of parameters is much smaller in our method in comparison to DT with a similar level of accuracy. For instance, on Hopper we reduce the model size by 84% in comparison to DT. We hypothesize that this phenomenon arises from the recursive structure of S4, which is much more suitable for RL tasks. Since the agent should make decisions recursively, based on the decisions made in the past, RL tasks are recursive in nature, causing non-recursive models such as DT to be inefficient, as they would have to save and handle the history of the current trajectory. Note that this history must be handled even in fully-observable tasks, such D4RL. For example, the result of DT on the Atari benchmark without history drastically decreased, as shown Section 5.3 of (Chen al. 202T).
Potentially Appropriate Inductive Bias S4 has strong inductive bias for recurrent processes (even when trained in non-recurrent manner) since Equation 2 is recurrent. Also, the HIPPO initialization is designed to implement a memory update mechanism over recursive steps. allowing the model to implement rich class of recurrent rules that can exploit the compressed history. Regarding RL tasks, we hypothesize that some of the improvement of $4 arises from the suitability of recursive processes for modeling MDPs. MDPs are widely used in the RL literature, both theoretical model and as a fundamental assumption employed by various algorithms.
C EXPERIMENT DETAILS
C.1 BENCHMARKS
The bulk of the experiments done using Deep-mind's [ujooo[[Toorro ct al. 2012) physical simulation of continuous tasks. The environments we tested our method on were Hopper, HalfCheetah and Walker2D. The offline data per environment
Mujoco Benchmark
taken from the (4RL[(F et al. 2020) off- off-policy benchmark. 3 types of datasets were tested, representing different levels of quality of the recorded data: expert, medium and medium-replay. The expert dataset consists of trajectories generated from trained actor using the soft actor critic method. The medium dataset is generated from a model trained with the same method, but only halfway through. The medium-replay dataset consists of the replay- replay-buffer buffer data that the medium dataset was training on the same time the medium dataset was exported.
AntMaze Benchmark The AntMaze navigation tasks benchmark Fu et al. (2020) is designed to measure the model's 's ability to deal with sparse rewards, multitask data, and long-range planning. In the task, an ant robot learns to navigate a maze. In the umaze task, the ant starts from fixed point, while in the umaze-dive diverse it start for random one. In both tasks, the robot should plan how reach the target. The reward for this task is binary: as the ant reaches its goal, the reward arises to 1, and otherwise O. Thus, this task is considered sparse reward task.
16

---

Published as a conference paper at ICLR 2023
C.2 EXPERIMENTAL SETUP
For off-policy training we used batches of 32 trajectories, using the maximum trajectory length in that batch the length for the entire batch, then filling shorter trajectories with blank input. The training was done with a learning rate of 10-6 and about 10000 warm-up steps, with linear incrementation until reaching the highest rate. We optimized the model using Adam Kingma Ba (2014), with a weight decay of 10 The use of the target return of the trajectory per trajectory as imput helped make use of less successful trajectories, as seen in previous works, such as the DT variations Chen et al. (02T): Janner et al. (2021). After this training, we continued with on-policy training, using the DDPG method, training the critic first for 35000 steps, and then continuing with both the actor and the critic, but prioritizing the critic in order to preserve stability. This was done by training the actor once for every 3 steps of the critic. We set the target actor and critic update to = 0.1 to maintain stability.
For on-policy learning, trajectories were generated by targeting the rewards 10% higher than the current model's highest rewards. As mentioned in 4.2 the first estimation is obtained by running the model with normalized return-to-go of R-1 = 1.0. As the training continues, we store the highest reward estimate and use it to generate the trajectories. Training for the models was done in batches of 96, with different learning rates for the critic and actor, specifically ac 10-3 and ox- = 10-5 The training occurred every Ki = 200 steps of the environment, and the target models were updated every K₂ 300 steps. To isolate the effects of each on-policy fine-tuning method, we measured the pure contribution of the method to the result. We report these findings in Tab. T and denote it by 6. All experiments take less than eight hours on single NVIDIA RTX 2080Ti GPU, which is similar to DT.
D VISUALIZATION OF THE LEARNED DEPENDENCIES
and 3 that the S4 recurrent view is F = A + Bug yk = Czg and it has an equivalent convolution form: y/: Ei=o Kitt-is K; = CA'B, y= K
Recall from Eq. and
Recall also that diagonal state-space layers employ B = (1, I,.,,,... 1) and define A as (complex) diagonal matrices which parameterized by A and A. We denote the eigenvalues of A as X3 E C for all coordinate j E [N]] Hence, for a given time we can write each coordinate j of the recurrent state are (5) as Ei Xzur- Empirically, we found that [Yyl S at the end of the training, thus according to De Moivre's ' formula [Ayl Z | 2 2 1251-
The eigenvalues AAyl control which ranges of dependencies will be captured by xx (j).). For example, in the case that (Ajl 1, the history can potentially accumulate over time steps and long-range dependencies can be handled, whereas with |Λyl << the information about the history is decaying rapidly over time.
Fig. 4 Iillustrates how different coordinates handle varying ranges of dependencies and which types of dependencies can be handled by single coordinate. Coordinates are represented colors in Fig. a and as rows in Fig, 4b). In both cases, it can be seen that different coordinates capture different types of dependencies. For example, since in Fig. 44) the red curve decays much faster than the blue curve, it is clear that the red curve targets shorter dependencies.
While Fig. explains how several coordinates of the recurrent state 24 can learn a wide spectrum of dependencies, it ignores the fact that the final output of the SSM depends only on the projection of z; via C. Thus, a more accurate view of the learned dependencies must take K into account, which is parameterized by C.
Recall that in the convolutional view (Fq.3 the output of the SSM is y = K * u where K; = CA'. Therefore, for a given time t, the scalar K, in the kernel controls how the SSM takes the input u into account when calculating y/2. Fig. 5 visualizes these kernels and illustrates that they can capture complex and wide dependencies.
Fig. 5[a) a) depicts how complex dependencies can be expressed by a single kernel. In general, the early part of the blue curve appears to be the most dominant one, which allows the model to handle short dependencies. On the other hand, the last elements of the curve have non-negligible values, thus at least some of the very early history is taken into account. Fig. 5)bb depicts multiple kernels as rows
17

---

Published as a conference paper at ICLR 2023
-055
-25
-05
620
425
2.6
075
213
-.72
1.00
(a)
(b)
Figure 4: Various ranges of dependencies captured by different coordinates of the recurrent state. The x-axis represents the position E [4]. (a) Values of Aÿ are represented on y-axis and colors represent different coordinates. Each row of pixels corresponds specific position the recurrent state 2, (j). Each row represents the values of Xÿ along different positions.
1.0
1852
(a)
(b)
Figure 5: Various ranges of dependencies are captured by different kernels. The x-axis represents the position E [L] in the kernel. (c) Values of K; are represented on the y-axis. (d) Each row of pixels corresponds to kernel from separate channels. Each row include a heat map of K; along i, at log-scale.
18

---

Published as a conference paper at ICLR 2023
of matrix. Evidently, the various kernels capture wide spectrum of dependencies, and each seems to ignore different parts of the sequence, and target different ranges and frequencies.
E NUMERICAL INSTABILITY OF THE S4 RECURRENT VIEW
In this section, we prove that the recurrent view of the state-space layer is unstable in some sense. Specifically, the worst-case numerical forward error of single channel grows proportional in Lemachine where L is the sequence length and N is the state size (See E.I). We observe that empirically, under the settings in Sec. [ which contains a sequence of length L 1000, these numerical errors are negligible when the recurrent view is used for inference. In contrast, when the recurrent view used for training, these errors negatively impact the result. The combination of exposure bias and relatively high numerical errors may explain this phenomenon. Note that when using the convolutional view, the summation can be calculated in a much more stable manner than recursive sum. Specifically, a stable summation can be computed using well-know methods such as Kahan summation algorithm (Kahan 1965) or 2Sum (Møller 1965). As far as we ascertain, we are the first to present this theoretical advantage of state-space layers over vanilla RNNs.
E.1 IDENTIFY THE INSTABILITY
For simplicity, we demonstrate the instability of a real-diagonal state-space layer. Under the assump- tions that diag (A,,.., Av) and B (1,..., 1), Eq. can re-written as:
V1 Car, zz(i) = AyTz-1(j) + we
(9)
Since the obtained computations along the coordinates 2, are independent for all j E [1,...,N]
Now observe the instability of Eq 9 The intuition is that when Ay (1) and ue not the same scale, this computation can result in loss of significance, which a well-known phenomenon (Higham 2002 Goldberg [199T.. Furthermore, by setting XÚ 1 in Eq the S4 recurrent view becomes a naive recursive summation algorithm, which is unstable (Muller et al. (018).
Theorem below shows that the worst-case forward error grows proportional at least N NLC L'mhhinee- where L is the sequence length and N is the state size.
We start by standard notation: Evaluation of expression in floating point is denoted fl: B B. and we assume that any basic arithmetic operations op (for example" '+) satisfy:
fl(a op b) = (1 5 (a
op b).
1800 S
machhne
(10)
Theorem 1. Denote the forward error of the recurrent view Eq. by A 1L 2
As, I NtCmuchine Σ Eisgrusl
Proof of Theorem
Proof. We start by unrolling under the assumption in Eq. 10 and denote the value of 2; by sum;- We denote the actual value of the state under floating-point arithmetic as sum, == fl(x.)
For simplicity, and is crucial for the proof, assume that only the '4' operator accumulated errors. Assuming sum-1 0.
sum = (1 + 8+ 8ttumn
(11)
sum (1 ft (uz + Ay 8t Sttum )uo)
(12)
sum2 (1 (1+ 8 3umg) U2 + Ny (1 Otumy (us + 2y (1 + 8 Stumo )uo)]
(13)
19

---

Published as a conference paper at ICLR 2023
And in general:
SWIT
Σ
2 up 11
)oum))
(14)
1=0
Recall that |6 Sumel
Emachines
therefore we can omit factor of machine
(1 5 Sume) )= 1
St sum,
+ +
machine
(15)
km
By plugging it in Eq.
sum
U1
1 +
5
Samma + +(P22 machine)
(16)
1=0
k=
And the obtained forward error is:
A sum. =
|sum, - ssumil |
Σ
sumng + +fp Offee machine) )
S O machine
1=0
k=
1=0
(17)
E.2 EMPIRICAL EVALUATION OF THE NUMERICAL ERRORS
Fig. 6 examines the forward error across several values of Xy, where the input u, is drawn from a normal standard distribution. The z axis depicts the step number i, and the axis the absolute forward error this step, averaged over 100 repeats. We used different scales because the rate which errors accumulate highly depends on Ay. as can be expected from theorem [ We approximate the true solution sum, by computing Eq. 9 under arbitrary precision floating-point arithmetic (2K
digits).
Prower Brion 2m
Parmard Bieer
5430
6031
51400
Figure 6: Forward Error for different values X E (0.8,, 1.0, 1.005)
As be from Fig. 6 when Xy the forward error relatively smaller. However, this
the important case for two reasons: (i) Where λS s 1, long-range dependencies can be handled, since the value of the state is > = X and when Aÿ the influence of uo doesn't 'tt change over time, allowing the model to copy history between time-steps. (ii) it also the most frequent case. As explain in Appendix D when A is complex, its magnitude !Ayl control the trade-off of which type of dependencies are handled (according to De Moivre's formula). Fig. Villustrates that at least at initialization, all the eigenvalues closed to the unit disk.
20

---

Published as a conference paper at ICLR 2023
1.00
0.75
0.50
0.25
0.00
-0.25
-0.50
-0.75
-1.00
-0.75
-0.50
-0.25
0.00
0.25
0.50
0.75
1.00
Figure 7: A visualization of the eigenvalues of À on initialization, which taken form diagonal state-space layer that include H = 10 channels and state of size N = 64. The red circle is the unit disk, and the black points are eigenvalues. Eigenvalues are consistently found in the left part of the plane within and close to the edges of the unit disk.
E.3 EMPIRICAL EVALUATION ON MUJOCO
We observe that at least on Mujoco benchmark, the forward numerical errors are negligible on inference, and they did not negatively impact the results. When the recurrent view used for training this is not the case. First, for off-policy learning, we ablate the usage in the convolutional view and train the model via the recurrent view. We found that performance degraded by 0.2% 30% in all experiments. To validate that the degradation is not result of less amount of tuning, conduct several experiments for each environment over several learning rates, types of optimizer, amount of dropout, and normalization layers. We observe that the gap between using the recurrent convolutional view is the highest in the medium-replay environments (degradation of 15% 60%).
Second, for on-policy fine-tuning, show that freezing the S4 parameters (A and A) mitigates this problem in on-policy RL. The motivation is (i) adding regularization, and (ii) to validate that long-range dependencies captured during the off-policy stage will diminish. We observe that freezing those parameters (i) increases the stability of the training and decreases the hyper-parametters sensitively. (ii) Accelerate the converges by least 2x (iii) improve the on policy results by around 5%.
21


Reviews:
[{"review_id": "ShKJJeOCqA", "reviewer": "Reviewer_YrJg", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "This paper presents an offline reinforcement learning approach which is able to capture longer range dependencies than a traditional sequence modeling approaches such as a Decision Transformer (DT). Similar to DT, Decision S4 views RL as a sequence modeling problem, but using the implicit S4 model instead of attention blocks. The offline approach is trained by taking individual transitions and rewards to go, and predicting the actions. The paper also introduces an offline-to-online training approach by first freezing the actor and then the critic, as well as freezing the S4 kernel. Results show that a much smaller model can achieve a similar performance to DT on offline RL tasks. ", "summary_of_the_review": "This paper provides an interesting use of the S4 parameterization and shows that the method works well compared to alternatives (such as Decision Transformer). Details and explanations about what the S4 blocks actually capture in the context of RL are missing which would make the claims in the paper stronger. ", "strength_and_weaknesses": "Strengths: \n\n- To my knowledge this paper presents a novel approach to tackle offline RL\n- Applications of the S4/LSSM architectures to control and RL are under-explored and could be beneficial for the community \n- S4 does provide almost similar performance to DT with fewer parameters \n- Offline-to-online finetuning approach works quite well \n- paper is well written \n\nWeaknesses: \n\n- I don't particularly understand why the offline training is capturing long range dependencies, if it is predicting a single action at a time conditioned on the current state and reward to go? Where is this captured? It's important to explain this well. \n- How long does it take to train the model vs DT and what is the inference time vs DT? \n\nI would be willing to increase my rating if the questions above are answered. \n\n\nI would be willing to incre", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "8K8BFtXWnM", "reviewer": "Reviewer_15g8", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "The paper presents a method for the recent decision-making-as-sequence-modeling paradigm. It follows the recent trend of treating RL as a sequence modeling problem: given the state-action-reward history, the current state, and the desired reward-to-go, predict the next action. The proposed method is to replace the commonly used transformer model backbone with S4 (Gu et al., 2021) for better long-range sequence modeling capability and reduced model size. The paper also proposes a offline-online learning scheme where the actor is trained with offline data using the decision-transformer-like objective and finetuned online using an actor-critic algorithm.", "summary_of_the_review": "Overall I think the paper is a solid execution of a conceptually simple idea, although some part of the method is not strictly necessary and may muddle the key message of the paper.", "strength_and_weaknesses": "Strengths: Overall I like the motivation of the paper: a lean model that can capture long-range dependency is desirable for any RL agent, and S4 seems to be the perfect backbone for the job. The method proposes a relatively straightforward drop-in replacement of recurrent S4 architecture for the more heavy-duty transformer backbone and shows that the new architecture performs favorably both in task and computational performance. \n\nWeaknesses: \n- To me, the offline-online learning scheme is an awkward orthogonal component to the work that’s neither novel nor tightly connected with the core DS4 method. It appears that S4 architecture is not uniquely attuned to online finetuning (compared to transformer / regular RL agents)? And the proposed online finetuning scheme is not tailored to the DS4 architecture either. And the online finetuning is not mentioned in other aspects of evaluation, e.g., long-range dependency and parameter reduction. In the future iteration of this paper, I’d suggest the author remove this component and focus on testing S4 against more challenging domains such as domains that require strong long-term memory.\n- The paper claims runtime latency to be a major advantage of S4-based architecture. I'd like to see some empirical results that substantiate this point.\n- Minor: the writing needs improvement. Most sentences are way too long for the message they carry. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "7NQizCCNph", "reviewer": "Reviewer_1zRF", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "This paper proposes replacing Learning (RL) methods based on Transformers with the S4 family of models. Two algorithms are proposed for both the on-policy and off-policy setting, showing strong empirical results.\n", "summary_of_the_review": "This work applies a recent method (S4) to a new domain (RL), showing promising results in performance and efficiency.\n\n--------\nPost-rebuttal: The authors added substantial new content in Appendices explaining several design choices and contrasted the tunability of the model compared to baselines. More improvements have been applied, increasing the performance of the method. All my concerns are addressed and I have raised my score.", "strength_and_weaknesses": "**Strengths:**\n- The method is well-motivated. RL can require both efficient batch training as well as efficient step-by-step unrolling, and makes sense as an application of state-space models (SSMs) which can leverage both an efficient parallel training (convolutional mode) and recurrent inference (recurrent mode) of SSMs.\n- The on-policy algorithm involves a methodological innovation that involves training with the recurrent view of S4, which was previously only used during inference time.\n- The empirical evaluation is comprehensive and the model performs well on standard benchmarks with many established baselines.\n- There are extensive empirical ablations, such as on the importance of long-term dependencies, the comparison against more basic RNNs, and on the model size.\n\n**Weaknesses:**\n- The methodological innovation is limited. However, the paper is clearly positioned as an empirical work that applies a recent model to a new domain, so this is not a significant weakness.\n- It isn't clear how much tuning is required for the proposed method over DT baselines. Also, it isn't clear how much of the proposed algorithms is tailored to DS4; it could be even more compelling if the same algorithms worked for DT but improved with DS4 as a drop-in replacement.\n\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster