---

Mastering Memory Tasks with World Models

Samsami M R, Zholus A, Rajendran J, et al. Mastering Memory Tasks with World Models. ICLR, 2024.

https://openreview.net/forum?id=1vDArHJ68h

https://recall2imagine.github.io/

https://github.com/chandar-lab/Recall2Imagine

RLWorld ModelMemoryICLR 2024 OralCCF-AMila 魁北克人工智能研究所蒙特利尔大学蒙特利尔理工学院CIFAR AI 讲席

---

使用世界模型掌握记忆任务

Abstract

Current model-based reinforcement learning (MBRL) agents struggle with long-term dependencies. This limits their ability to effectively solve tasks involving extended time gaps between actions and outcomes, or tasks demanding the recalling of distant observations to inform current actions. To improve temporal coherence, we integrate a new family of state space models (SSMs) in world models of MBRL agents to present a new method, Recall to Imagine (R2I). This integration aims to enhance both long-term memory and long-horizon credit assignment.

当前的基于模型的强化学习（MBRL）智能体难以处理长期依赖关系。 这限制了它们有效解决动作与结果之间存在较长时间间隔的任务，也限制了它们通过回忆久远观察来指导当前动作的能力。 为了增强时间连贯性，我们将一类新的状态空间模型（SSM）集成到 MBRL 智能体的世界模型中，并提出一种新方法 Recall to Imagine（R2I）。 这种集成旨在同时增强长期记忆和长程信用分配能力。

Through a diverse set of illustrative tasks, we systematically demonstrate that R2I not only establishes a new state-of-the-art for challenging memory and credit assignment RL tasks, such as BSuite and POPGym, but also showcases superhuman performance in the complex memory domain of Memory Maze. At the same time, it upholds comparable performance in classic RL tasks, such as Atari and DMC, suggesting the generality of our method. We also show that R2I is faster than the state-of-the-art MBRL method, DreamerV3, resulting in faster wall-time convergence.

通过一组多样且具有代表性的任务，我们系统地证明，R2I 不仅在 BSuite 和 POPGym 等具有挑战性的记忆与信用分配强化学习任务上建立了新的最先进水平，还在复杂记忆领域 Memory Maze 中展现出超越人类的性能。 与此同时，它在 Atari 和 DMC 等经典强化学习任务上保持了相当的性能，体现了我们方法的通用性。 我们还表明，R2I 比最先进的 MBRL 方法 DreamerV3 更快，从而能在更短的实际时间内收敛。

1. Introduction

In reinforcement learning (RL), world models, which capture the dynamics of the environment, have emerged as a powerful paradigm for integrating agents with the ability to perceive, simulate, and plan within the learned dynamics. In current model-based reinforcement learning (MBRL), the agent learns the world model from past experiences, enabling it to “imagine” the consequences of its actions (such as the future environment rewards and observations) and make informed decisions.

在强化学习（RL）中，能够捕捉环境动态的世界模型已经成为一种强大范式，使智能体能够在学习到的动态中进行感知、模拟和规划。 在当前的基于模型的强化学习（MBRL）中，智能体从过去的经验中学习世界模型，从而能够“想象”其动作的后果，例如未来的环境奖励和观察，并据此作出明智决策。

MBRL necessitates learning a world model that accurately simulates the environment's evolution and future rewards, integrating the agent's actions over long horizons. This task is compounded by the credit assignment (CA) problem, where an action's impact on future rewards must be evaluated. The agent also may need to memorize and recall past experiences to infer optimal actions. The challenge of long-term memory and CA frequently arises as a result of inadequate learning of long-range dependencies, due to constraints in world models' backbone network architecture.

MBRL 要求学习一个能够准确模拟环境演化与未来奖励的世界模型，并将智能体在长时间范围内的动作纳入其中。 信用分配（CA）问题进一步加剧了这项任务，因为必须评估一个动作对未来奖励的影响。 智能体还可能需要记住并回忆过去的经验，以推断最优动作。 由于世界模型骨干网络架构的限制，长期依赖关系学习不足时，长期记忆与信用分配的挑战便会频繁出现。

More specifically, Recurrent Neural Networks (RNNs) are employed in most MBRL methods as the world models' backbone architecture because of their ability to handle sequential data. However, their efficacy is hindered by the vanishing gradients. Alternately, due to the remarkable achievements of Transformers in language modeling tasks, they have been recently adopted to build world models. Nonetheless, the computational complexity of Transformers is quadratic in its input sequence length. Even the optimized Transformers become unstable during training on long sequences. This prohibits Transformers-based world models from scaling to long input sequence lengths that might be required in certain RL tasks.

更具体地说，由于循环神经网络（RNN）能够处理序列数据，大多数 MBRL 方法都将其用作世界模型的骨干架构。 然而，梯度消失问题会削弱其有效性。 另一方面，Transformer 在语言建模任务中取得了显著成就，因此近来它们也被用于构建世界模型。 不过，Transformer 的计算复杂度与输入序列长度呈二次关系。 即使是经过优化的 Transformer，在长序列训练期间也会变得不稳定。 这使基于 Transformer 的世界模型无法扩展到某些强化学习任务可能需要的长输入序列。

Recent studies have revealed that state space models (SSMs) can effectively capture dependencies in tremendously long sequences for supervised learning (SL) and self-supervised learning (SSL) tasks. More specifically, the S4 model redefined the long-range sequence modeling research landscape by mastering highly difficult benchmarks.

近期研究表明，状态空间模型（SSM）能够在监督学习（SL）和自监督学习（SSL）任务中有效捕捉极长序列里的依赖关系。 更具体地说，S4 模型通过攻克高难度基准，重新塑造了长程序列建模的研究格局。

The S4 model is derived from a time-invariant linear dynamical system where state matrices are learned. In SL and SSL tasks, it exhibits a remarkable capability to capture dependencies extending up to 16K in length, surpassing the limitations of all prior methods. Given these achievements and MBRL methods' limitations in solving memory and CA tasks, the adoption of S4 or a modified version of it is a logical decision. In this paper, we introduce a novel method termed Recall to Imagine (R2I), which is the first MBRL approach utilizing a variant of S4.

S4 模型源自一个学习状态矩阵的时不变线性动力系统。 在监督学习和自监督学习任务中，它展现出捕捉长度高达 16K 的依赖关系的卓越能力，突破了以往所有方法的限制。 鉴于这些成果以及 MBRL 方法在解决记忆与信用分配任务时的局限，采用 S4 或其改进版本是合乎逻辑的选择。 在本文中，我们提出一种名为 Recall to Imagine（R2I）的新方法，它是首个采用 S4 变体的 MBRL 方法。

This method empowers agents with long-term memory. R2I emerges as a general and computationally efficient approach, demonstrating state-of-the-art (SOTA) performance in a range of memory domains. Through rigorous experiments, we demonstrate that R2I not only surpasses the best-performing baselines but also exceeds human performance in tasks requiring long-term memory or credit assignment, all while maintaining commendable performance across various other benchmarks. Our contributions can be summarized as follows:

该方法赋予智能体长期记忆能力。 R2I 是一种通用且计算高效的方法，在一系列记忆领域中展现出最先进的性能。 通过严谨实验，我们证明 R2I 不仅超过了表现最好的基线，还在需要长期记忆或信用分配的任务中超越了人类，同时在其他各种基准上保持了出色表现。 我们的贡献可概括如下：

• We introduce R2I, a memory-enhanced MBRL agent built upon DreamerV3 that uses a modification of S4 to handle temporal dependencies. R2I inherits the generality of DreamerV3, operating with fixed world model hyperparameters on every domain, while also offering an improvement in computational speed of up to 9 times.
• We demonstrate SOTA performance of the R2I agent in a diverse set of memory domains: POPGym, Behavior Suite (BSuite), and Memory Maze. Notably, in the Memory Maze, which is a challenging 3D domain with extremely long-term memory needed to be solved, R2I outperforms human.
• We investigate R2I's performance in established RL benchmarks, namely Atari and DMC. We show that R2I's improved memory does not compromise performance across different types of control tasks, highlighting its generality.
• We conduct ablation experiments to show the impact of the design decisions made for R2I.

• 我们提出 R2I，这是一种建立在 DreamerV3 之上的记忆增强型 MBRL 智能体，使用 S4 的改进版本来处理时间依赖关系。R2I 继承了 DreamerV3 的通用性，在所有领域中使用固定的世界模型超参数，同时将计算速度最高提升 9 倍。
• 我们证明 R2I 智能体在 POPGym、Behavior Suite（BSuite）和 Memory Maze 等多种记忆领域中取得了最先进的性能。尤其是在需要极长期记忆才能解决、具有挑战性的三维领域 Memory Maze 中，R2I 超越了人类。
• 我们研究了 R2I 在 Atari 和 DMC 等成熟强化学习基准上的表现。我们表明，R2I 增强的记忆能力不会损害其在不同类型控制任务中的性能，凸显了其通用性。
• 我们开展消融实验，以展示 R2I 各项设计决策的影响。

2. Background

2.1 State Space Models

A recent work has introduced a novel Structured State Space Sequence model (S4). This model has shown superior performance in SL and SSL tasks, compared to common deep sequence models, including RNNs, convolutional neural networks (CNNs), and Transformers. It outperforms them in terms of both computational efficiency and the ability to model extremely long-range dependencies. S4 is a specific instance of state space models (SSMs), which can be efficiently trained by using specialized parameterization.

近期工作提出了一种新的结构化状态空间序列模型（S4）。 与 RNN、卷积神经网络（CNN）和 Transformer 等常见深度序列模型相比，该模型在监督学习和自监督学习任务中表现更优。 它在计算效率和建模极长程依赖关系的能力方面都优于这些模型。 S4 是状态空间模型（SSM）的一个特例，可以通过专门的参数化方式进行高效训练。

SSMs are derived from a linear dynamical system with control variable $u(t)\in \mathbb{R}$ and observation variable $y(t)\in \mathbb{R}$, utilizing state variables $x(t)\in {\mathbb{C}}^N$ for a state size $N$. The system is represented by the state matrix $\mathbf{A}\in {\mathbb{C}}^{N\times N}$ and other matrices $\mathbf{B}\in {\mathbb{C}}^{N\times 1}$, $\mathbf{C}\in {\mathbb{C}}^{1\times N}$, and $\mathbf{D}\in {\mathbb{R}}^{1\times 1}$:

SSM 源自一个线性动力系统，其中控制变量为 $u(t)\in \mathbb{R}$，观察变量为 $y(t)\in \mathbb{R}$，并使用状态大小为 $N$ 的状态变量 $x(t)\in {\mathbb{C}}^N$。 该系统由状态矩阵 $\mathbf{A}\in {\mathbb{C}}^{N\times N}$ 以及其他矩阵 $\mathbf{B}\in {\mathbb{C}}^{N\times 1}$、$\mathbf{C}\in {\mathbb{C}}^{1\times N}$ 和 $\mathbf{D}\in {\mathbb{R}}^{1\times 1}$ 表示：

$$\begin{array}{rl}{x}^{\prime }(t)& =\mathbf{A}x(t)+\mathbf{B}u(t),\\ y(t)& =\mathbf{C}x(t)+\mathbf{D}u(t).\end{array}$$

Note that these SSMs function on continuous sequences. They can be discretized by a step size $\mathrm{\Delta }$ to allow discrete recurrent representation:

请注意，这些 SSM 作用于连续序列。 可以使用步长 $\mathrm{\Delta }$ 将其离散化，以获得离散的循环表示：

$$\begin{array}{rl}{x}_n& =\overline{\mathbf{A}}{x}_{n-1}+\overline{\mathbf{B}}{u}_n,\\ y_n& =\overline{\mathbf{C}}{x}_n+\overline{\mathbf{D}}{u}_n.\end{array}$$

Here, $\overline{\mathbf{A}}$, $\overline{\mathbf{B}}$, $\overline{\mathbf{C}}$, and $\overline{\mathbf{D}}$ are discrete-time parameters obtained from the continuous-time parameters and $\mathrm{\Delta }$ using methods like zero-order hold and bilinear technique. These representations are incorporated as a neural network layer, and each SSM is used to process a single dimension of the input sequence and map it to a single output dimension. This means that there are separate linear transformations for each input dimension, which are followed by a nonlinearity. This allows working with discrete sequence tasks, such as language modeling, speech classification, and pixel-level 1D image classification.

其中，$\overline{\mathbf{A}}$、$\overline{\mathbf{B}}$、$\overline{\mathbf{C}}$ 和 $\overline{\mathbf{D}}$ 是通过零阶保持、双线性变换等方法，根据连续时间参数和 $\mathrm{\Delta }$ 获得的离散时间参数。 这些表示被整合为神经网络层，每个 SSM 处理输入序列的一个维度，并将其映射到一个输出维度。 这意味着每个输入维度都有单独的线性变换，之后再接一个非线性函数。 因此，它可以处理语言建模、语音分类和像素级一维图像分类等离散序列任务。

S4 model characterizes $\mathbf{A}$ as a matrix with a diagonal plus low-rank (DPLR) structure. One benefit of this “structured” representation is that it helps preserve the sequence history; S4 employs HiPPO framework to initialize the matrix $\mathbf{A}$ with special DPLR matrices. This initialization grants the SSMs the ability to decompose $u(t)$ into a set of infinitely long basis functions, enabling the SSMs to capture long-range dependencies.

S4 模型将 $\mathbf{A}$ 表示为具有对角加低秩（DPLR）结构的矩阵。 这种“结构化”表示的一个优点是有助于保留序列历史；S4 使用 HiPPO 框架，以特殊的 DPLR 矩阵初始化 $\mathbf{A}$。 这种初始化使 SSM 能够将 $u(t)$ 分解为一组无限长的基函数，从而捕捉长程依赖关系。

Further, to make S4 more practical on modern hardware, the mapping $u_{1:T},x_0\to y_{1:T},x_T$ has been reparameterized as a global convolution, referred to as the convolution mode, thereby avoiding sequential training (as in RNNs). This modification has made S4 faster to train, and S4 models can be thought of as a fusion of CNNs, RNNs, and classical SSMs. Parallel scan can compute $u_{1:T},x_0\to y_{1:T},x_{1:T}$ as efficiently as convolution mode.

此外，为了让 S4 更适合现代硬件，映射 $u_{1:T},x_0\to y_{1:T},x_T$ 被重新参数化为全局卷积，称为卷积模式，从而避免像 RNN 那样进行序列式训练。 这一改动使 S4 的训练更快，也可以将 S4 模型视为 CNN、RNN 与经典 SSM 的融合。 并行扫描能够以与卷积模式同等的效率计算 $u_{1:T},x_0\to y_{1:T},x_{1:T}$。

S4 has demonstrated impressive empirical results on various established SL and SSL benchmarks involving long dependencies, and it outperforms Transformers in terms of inference speed and memory consumption due to its recurrent inference mode. Moreover, some recent works have focused on understanding S4 models, as well as refining them and augmenting their capabilities. For the sake of simplicity in this study, we will be referring to all the S4 model variations as “SSMs”. It is worth highlighting that a few recent methods optimize the performance of SSMs by integrating them with Transformers. This enhances the SSMs by adding a powerful local attention-based inductive bias.

S4 已经在多个涉及长程依赖的成熟监督学习和自监督学习基准上取得了令人瞩目的实证结果，并且凭借循环推理模式，在推理速度和内存消耗方面优于 Transformer。 此外，近期一些工作致力于理解 S4 模型，并对其进行改进和能力增强。 为简化表述，本文将所有 S4 模型变体统称为“SSM”。 值得强调的是，一些近期方法通过将 SSM 与 Transformer 集成来优化其性能。 这种做法为 SSM 增加了强大的局部注意力归纳偏置。

2.2 From Imagination To Action

We frame a sequential decision-making problem as a partially observable Markov decision process (POMDP) with observations $o_t$, scalar rewards $r_t$, agent's actions $a_t$, episode continuation flag $c_t$, and discount factor $\gamma \in (0,1)$, all following dynamics $o_t,r_t,c_t\sim p(o_t,r_t,c_t\mid o_{<t},a_{<t})$. The goal of RL is to train a policy $\pi$ that maximizes the expected value of the discounted return ${\mathbb{E}}_{\pi }[\sum _{t\ge 0}{\gamma }^t{r}_t]$.

我们将序列决策问题表示为部分可观测马尔可夫决策过程（POMDP），其中包含观察 $o_t$、标量奖励 $r_t$、智能体动作 $a_t$、回合继续标志 $c_t$ 和折扣因子 $\gamma \in (0,1)$，它们都遵循动态 $o_t,r_t,c_t\sim p(o_t,r_t,c_t\mid o_{<t},a_{<t})$。 强化学习的目标是训练策略 $\pi$，使折扣回报的期望值 ${\mathbb{E}}_{\pi }[\sum _{t\ge 0}{\gamma }^t{r}_t]$ 最大化。

In MBRL, the agent learns a model of the environment's dynamics (i.e., the world model), through an iterative process of collecting data using a policy, training the world model on the accumulated data, and optimizing the policy through the world model. The Dreamer agent and its subsequent versions have been impactful MBRL systems that learn the environment dynamics in a compact latent space and learn the policy entirely within that latent space. Dreamer agents consist of three primary components: the world model, which predicts the future outcomes of potential actions, the critic, which estimates the value of each state, and the actor, which learns to take optimal actions.

在 MBRL 中，智能体通过一个迭代过程学习环境动态模型，也就是世界模型：使用策略收集数据，在累积数据上训练世界模型，再通过世界模型优化策略。 Dreamer 智能体及其后续版本是具有重要影响的 MBRL 系统，它们在紧凑的潜在空间中学习环境动态，并完全在该潜在空间中学习策略。 Dreamer 智能体包含三个主要组件：预测潜在动作未来结果的世界模型、估计各状态价值的评论器，以及学习采取最优动作的行动器。

In Dreamer, an RNN-based architecture called Recurrent State-Space Model (RSSM) serves as the core of the world model, and it can be described as follows. For every time step $t$, it represents the latent state through the concatenation of deterministic state $h_t$ and stochastic state $z_t$. Here, $h_t$ is updated using a Gated Recurrent Unit (GRU), and then is utilized to compute $z_t$, which incorporates information about the current observation $o_t$ and is subsequently referred to as the posterior state. Additionally, the prior state ${\hat{z}}_t$ which predicts $z_t$ without access to $o_t$ is computed using $h_t$. By leveraging the latent state $(z_t,h_t)$, we can reconstruct various quantities such as $o_t$, $r_t$, and $c_t$.

在 Dreamer 中，一种基于 RNN、称为循环状态空间模型（RSSM）的架构构成世界模型的核心，其描述如下。 对于每个时间步 $t$，它通过拼接确定性状态 $h_t$ 和随机状态 $z_t$ 来表示潜在状态。 其中，$h_t$ 由门控循环单元（GRU）更新，随后用于计算 $z_t$；$z_t$ 融合当前观察 $o_t$ 的信息，因此被称为后验状态。 此外，还使用 $h_t$ 计算先验状态 ${\hat{z}}_t$，它在无法访问 $o_t$ 的情况下预测 $z_t$。 利用潜在状态 $(z_t,h_t)$，我们可以重构 $o_t$、$r_t$ 和 $c_t$ 等各种量。

The RSSM comprises three components: a sequence model ($h_t=f_{\theta }(h_{t-1},z_{t-1},a_{t-1})$), a representation model ($z_t\sim q_{\theta }(z_t\mid h_t,o_t)$), and a dynamics model (${\hat{z}}_t\sim p_{\theta }({\hat{z}}_t\mid h_t)$), where $a_{t-1}$ is the action at time step $t-1$, and $\theta$ denotes the combined parameter vector of all components. In addition to the RSSM, the world model has separate prediction heads for $o_t$, $r_t$, $c_t$. Within the imagination phase, it harnesses the RSSM to simulate trajectories. This is performed through an iterative computation of states ${\hat{z}}_t,h_t$ and actions ${\hat{a}}_t\sim \pi ({\hat{a}}_t\mid {\hat{z}}_t,h_t)$ without the need for observations (except in the initial step). The sequences of ${\hat{z}}_{1:T},h_{1:T},{\hat{a}}_{1:T}$ are used to train the actor and the critic.

RSSM 包含三个组件：序列模型（$h_t=f_{\theta }(h_{t-1},z_{t-1},a_{t-1})$）、表示模型（$z_t\sim q_{\theta }(z_t\mid h_t,o_t)$）和动态模型（${\hat{z}}_t\sim p_{\theta }({\hat{z}}_t\mid h_t)$）；其中 $a_{t-1}$ 是时间步 $t-1$ 的动作，$\theta$ 表示所有组件的联合参数向量。 除 RSSM 外，世界模型还为 $o_t$、$r_t$ 和 $c_t$ 配置了独立的预测头。 在想象阶段，它使用 RSSM 模拟轨迹。 除初始步骤外，该过程无需观察，而是迭代计算状态 ${\hat{z}}_t,h_t$ 和动作 ${\hat{a}}_t\sim \pi ({\hat{a}}_t\mid {\hat{z}}_t,h_t)$。 随后使用 ${\hat{z}}_{1:T},h_{1:T},{\hat{a}}_{1:T}$ 序列训练行动器和评论器。

3. Methodology

We introduce R2I (Recall to Imagine), which integrates SSMs in DreamerV3's world model, giving rise to what we term the Structured State-Space Model (S3M). The design of the S3M aims to achieve two primary objectives: capturing long-range relations in trajectories and ensuring fast computational performance in MBRL. S3M achieves the desired speed through parallel computation during training and recurrent mode in inference time, which enables quick generation of imagined trajectories. In Figure 1, a visual representation of R2I is provided, and we will now proceed to describe its design.

我们提出 R2I（Recall to Imagine），它将 SSM 集成到 DreamerV3 的世界模型中，形成我们所称的结构化状态空间模型（S3M）。 S3M 的设计旨在实现两个主要目标：捕捉轨迹中的长程关系，并确保 MBRL 具备较快的计算性能。 S3M 通过训练期间的并行计算和推理期间的循环模式获得所需速度，从而能够快速生成想象轨迹。 图1给出了 R2I 的可视化表示，下面我们将介绍其设计。

图1：R2I 的图形化表示。左：世界模型编码过去经验，将观察和动作转换为紧凑的潜在状态；重构轨迹为塑造这些潜在状态提供学习信号。右：策略从世界模型想象的、基于潜在状态的轨迹中学习。图中表示对应完整状态策略，为简化图示省略了评论器。

3.1 World Model Details

Non-recurrent representation model. Our objective when updating the world model is to calculate S3M deterministic states $h_{1:T}$ in parallel by simultaneously feeding all actions $a_t$ and stochastic state $z_t$, where $T$ represents the length of the entire sequence. We aim to carry out this computation as $h_{1:T},x_{1:T}=f_{\theta }((a_{1:T},z_{1:T}),x_0)$ where $x_t$ is a hidden state and $f_{\theta }$ is a sequence model with a SSM network. To achieve this, prior access to all actions $a_{1:T}$ and stochastic states $z_{1:T}$ is required. However, we encounter a challenge due to the sequential nature of the relationship between the representation model $q_{\theta }(z_t\mid h_t,o_t)$ and sequence model $f_{\theta }(h_{t-1},z_{t-1},a_{t-1})$: at time step $t$, the representation model's most recent output, denoted as $z_{t-1}$, is fed into the sequence model, and the resulting output $h_t$ is then used within the representation model to generate $z_t$.

非循环表示模型。 更新世界模型时，我们的目标是同时输入所有动作 $a_t$ 和随机状态 $z_t$，并行计算 S3M 的确定性状态 $h_{1:T}$，其中 $T$ 表示整个序列的长度。 我们希望将这一计算表示为 $h_{1:T},x_{1:T}=f_{\theta }((a_{1:T},z_{1:T}),x_0)$，其中 $x_t$ 是隐状态，$f_{\theta }$ 是采用 SSM 网络的序列模型。 为此，需要预先获得所有动作 $a_{1:T}$ 和随机状态 $z_{1:T}$。 然而，表示模型 $q_{\theta }(z_t\mid h_t,o_t)$ 与序列模型 $f_{\theta }(h_{t-1},z_{t-1},a_{t-1})$ 之间具有序列依赖关系，这带来了一项挑战：在时间步 $t$，表示模型最新输出 $z_{t-1}$ 被输入序列模型，序列模型输出的 $h_t$ 又被表示模型用于生成 $z_t$。

Hence, by eliminating the dependency on $h_t$ in the representation model, we transform it to a non-recurrent representation model $q_{\theta }(z_t\mid o_t)$. This modification allows us to compute the posterior samples independently for each time step, enabling simultaneous computation for all time steps. By utilizing a parallelizable function $f_{\theta }$, we can then obtain $h_{1:T}$ in parallel. A systematic analysis in the appendix investigates how this modification impacts the performance of DreamerV3 across a diverse set of tasks. The results indicate that transforming $q_{\theta }(z_t\mid o_t,h_t)$ to $q_{\theta }(z_t\mid o_t)$ does not hurt the performance.

因此，通过消除表示模型对 $h_t$ 的依赖，我们将其转换为非循环表示模型 $q_{\theta }(z_t\mid o_t)$。 这一修改使我们能够在每个时间步独立计算后验样本，从而同时计算所有时间步。 随后，利用可并行化函数 $f_{\theta }$，我们可以并行获得 $h_{1:T}$。 附录中的系统分析研究了这一修改对 DreamerV3 在多种任务上性能的影响。 结果表明，将 $q_{\theta }(z_t\mid o_t,h_t)$ 转换为 $q_{\theta }(z_t\mid o_t)$ 不会损害性能。

Architecture details. Inspired by Dreamer, R2I's world model consists of a representation model, a dynamics model, and a sequence model (together forming S3M). In addition to that, there are three prediction heads: an observation predictor $p_{\theta }({\hat{o}}_t\mid z_t,h_t)$, a reward predictor $p_{\theta }({\hat{r}}_t\mid z_t,h_t)$, and an episode continuation predictor $p_{\theta }({\hat{c}}_t\mid z_t,h_t)$. At each time step, S3M processes a pair of $(a_t,z_t)$ to output the deterministic state $h_t$. Inside, it operates over the hidden state $x_t$, so it can be defined as $h_t,x_t=f_{\theta }((a_{t-1},z_{t-1}),x_{t-1})$.

架构细节。 受 Dreamer 启发，R2I 的世界模型由表示模型、动态模型和序列模型组成，三者共同构成 S3M。 除此之外，它还包含三个预测头：观察预测器 $p_{\theta }({\hat{o}}_t\mid z_t,h_t)$、奖励预测器 $p_{\theta }({\hat{r}}_t\mid z_t,h_t)$ 和回合继续预测器 $p_{\theta }({\hat{c}}_t\mid z_t,h_t)$。 在每个时间步，S3M 处理一对 $(a_t,z_t)$ 并输出确定性状态 $h_t$。 在其内部，模型基于隐状态 $x_t$ 运行，因此可定义为 $h_t,x_t=f_{\theta }((a_{t-1},z_{t-1}),x_{t-1})$。

Specifically, $f_{\theta }$ is composed of multiple layers of SSMs, each one calculating outputs according to the discrete recurrent representation. The outputs are then passed to GeLU, which is followed by a fully-connected GLU transformation, and finally by a LayerNorm. The deterministic state $h_t$ is the output from the final SSM layer. The set of all SSM layer hidden states is denoted $x_t$. In image-based environments, we leverage a CNN encoder for $q_{\theta }(z_t\mid o_t)$ and a CNN decoder for $p_{\theta }({\hat{o}}_t\mid z_t,h_t)$. In contrast, in tabular environments, both $q_{\theta }(z_t\mid o_t)$ and $p_{\theta }({\hat{o}}_t\mid z_t,h_t)$ are MLPs.

具体而言，$f_{\theta }$ 由多层 SSM 组成，每一层都根据离散循环表示计算输出。 这些输出随后依次经过 GeLU、全连接 GLU 变换和 LayerNorm。 确定性状态 $h_t$ 是最后一个 SSM 层的输出。 所有 SSM 层的隐状态集合记为 $x_t$。 在基于图像的环境中，我们为 $q_{\theta }(z_t\mid o_t)$ 使用 CNN 编码器，并为 $p_{\theta }({\hat{o}}_t\mid z_t,h_t)$ 使用 CNN 解码器。 相比之下，在表格型环境中，$q_{\theta }(z_t\mid o_t)$ 和 $p_{\theta }({\hat{o}}_t\mid z_t,h_t)$ 都采用 MLP。

Training details. R2I optimizes the following objective:

训练细节。 R2I 优化以下目标：

$$\begin{array}{rl}\mathcal{L}(\theta )={\mathbb{E}}_{z_{1:T}\sim q_{\theta }}\sum _{t=1}^T[& {\mathcal{L}}^{\mathrm{pred}}(\theta ,h_t,o_t,r_t,c_t,z_t)\\ & +{\mathcal{L}}^{\mathrm{rep}}(\theta ,h_t,o_t)+{\mathcal{L}}^{\mathrm{dyn}}(\theta ,h_t,o_t)].\end{array}$$$$\begin{array}{rl}{\mathcal{L}}^{\mathrm{pred}}& =-{\beta }_{\mathrm{pred}}(\ln p_{\theta }(o_t\mid z_t,h_t)+\ln p_{\theta }(r_t\mid z_t,h_t)+\ln p_{\theta }(c_t\mid z_t,h_t)),\\ {\mathcal{L}}^{\mathrm{dyn}}& ={\beta }_{\mathrm{dyn}}max(1,\mathrm{KL}[\mathrm{sg}(q_{\theta }(z_t\mid o_t))\parallel p(z_t\mid h_t)]),\\ {\mathcal{L}}^{\mathrm{rep}}& ={\beta }_{\mathrm{rep}}max(1,\mathrm{KL}[q_{\theta }(z_t\mid o_t)\parallel \mathrm{sg}(p(z_t\mid h_t))]),\\ h_{1:T},x_{1:T}& =f_{\theta }((a_{1:T},z_{1:T}),x_0).\end{array}$$

Here, sg represents the stop gradient operation. This loss, resembling the objective utilized in DreamerV3, is derived from Evidence Lower Bound (ELBO), but our objective differs from ELBO in three ways. First, we clip KL-divergence when it falls below the threshold of $1$. Secondly, we use KL-balancing to prioritize the training of the S3M. Third, we use scaling coefficients ${\beta }_{\text{pred}}$, ${\beta }_{\text{rep}}$, ${\beta }_{\text{dyn}}$ to adjust the influence of each term in the loss function. Some works on SSMs recommend optimizing state matrices using a smaller learning rate; however, our experiments indicate that the most effective approach is to use the same learning rate used in the rest of the world model.

其中，sg 表示停止梯度操作。 该损失与 DreamerV3 使用的目标相似，源自证据下界（ELBO），但我们的目标与 ELBO 有三点不同。 首先，当 KL 散度低于阈值 $1$ 时，我们对其进行截断。 其次，我们使用 KL 平衡来优先训练 S3M。 第三，我们使用缩放系数 ${\beta }_{\text{pred}}$、${\beta }_{\text{rep}}$ 和 ${\beta }_{\text{dyn}}$ 调整损失函数各项的影响。 一些 SSM 工作建议使用较小的学习率优化状态矩阵；然而，我们的实验表明，最有效的做法是采用与世界模型其余部分相同的学习率。

SSMs Computational Modeling. To enable the parallelizability of world model learning, we have the option to select between two distinct approaches: convolution and parallel scan. After thorough deliberation, we opted for parallel scan due to several compelling reasons. Firstly, it is essential to pass hidden states $x_t$ to the policy in memory environments, a critical finding we empirically analyze in the appendix. Another consequence of not yielding $x_t$ via convolution mode is that it would necessitate several burn-in steps to obtain correct hidden states, resulting in quadratic imagination complexity.

SSM 的计算建模。 为了使世界模型学习可以并行化，我们可以在卷积和并行扫描两种不同方法之间选择。 经过充分权衡，我们出于几个重要原因选择了并行扫描。 首先，在记忆环境中必须将隐状态 $x_t$ 传递给策略，我们在附录中对这一关键发现进行了实证分析。 如果卷积模式不产生 $x_t$，还会导致另一个后果：为了得到正确的隐状态，需要执行若干预热步骤，从而使想象过程具有二次复杂度。

Furthermore, parallel scan enables scaling of sequence length in batch across distributed devices, a capability not supported by the convolution mode. Table 1 summarizes computational complexities associated with different types of recurrences, including RNNs, SSMs, and Attention.

此外，并行扫描允许在分布式设备上扩展批次中的序列长度，而卷积模式不支持这种能力。 表1总结了 RNN、SSM 和注意力等不同循环机制的计算复杂度。

表1：不同架构的渐近运行时间。L 为序列长度，H 为想象范围。想象过程的外层循环无法并行化。注意力和 SSM+Conv 接收由 O(L+H) 个预热步骤与想象步骤组成的完整上下文，因此注意力的单步复杂度为 O((L+H)²)，SSM+Conv 为 O(L)。SSM 将紧凑循环与并行计算结合，达到最佳渐近复杂度。

Method	Training	Inference step	Imagination step	Parallel	State Reset
Attn	O(L²)	O(L²)	O((L+H)²)	✓	✓
RNN	O(L)	O(1)	O(1)	×	✓
SSM (Conv)	O(L)	O(1)	O(L)	✓	×
SSM (Par. Scan)	O(L)	O(1)	O(1)	✓	✓

Finally, parallel scan can facilitate the resetting of hidden states. When sampling a sequence from the buffer, it may comprise of multiple episodes; thus, the hidden states coming from terminal states to the initial states in new episodes must be reset. This boosts the early training performance, when the episodes may be short. We modify the SSM inference operator to support resetting hidden states. Achieving this is not feasible with convolution mode.

最后，并行扫描还可以支持隐状态重置。 从缓冲区采样的序列可能包含多个回合，因此必须重置从终止状态传递到新回合初始状态的隐状态。 当回合可能较短时，这会改善训练早期的性能。 我们修改了 SSM 推理算子，使其支持重置隐状态。 卷积模式无法做到这一点。

3.2 Actor-Critic Details

In the design of Dreamer's world model, it is assumed that $h_t$ contains information summarizing past observations, actions, and rewards. Then, $h_t$ is leveraged in conjunction with the stochastic state $z_t$ to reconstruct or predict observations, rewards, episode continuation, actions, and values. Unlike DreamerV3, which utilizes a GRU cell wherein $h_t$ is passed both to the reconstruction heads and the next recurrent step, R2I exclusively passes $h_t$ to prediction heads, while SSM's hidden state $x_t$ is used in the next recurrent update of S3M. This implies that the information stored in $h_t$ and $x_t$ could potentially vary.

在 Dreamer 的世界模型设计中，假设 $h_t$ 包含对过去观察、动作和奖励的汇总信息。 随后，将 $h_t$ 与随机状态 $z_t$ 结合，用于重构或预测观察、奖励、回合继续状态、动作和价值。 DreamerV3 使用 GRU 单元，将 $h_t$ 同时传递给重构头和下一个循环步骤；与之不同，R2I 只将 $h_t$ 传递给预测头，而在 S3M 的下一次循环更新中使用 SSM 的隐状态 $x_t$。 这意味着 $h_t$ 和 $x_t$ 中存储的信息可能有所不同。

Empirically, we discovered that this difference can lead to the breakdown of policy learning when using $\pi ({\hat{a}}_t\mid z_t,h_t)$, but it remains intact when we use $\pi ({\hat{a}}_t\mid z_t,x_t)$ in memory-intensive environments. Surprisingly, we found that incorporating all features into the policy $\pi ({\hat{a}}_t\mid z_t,h_t,x_t)$ is not a remedy. The reason lies in the non-stationarity of these features; their empirical distribution changes over time as the world model trains, ultimately leading to instability in the policy training process. A similar phenomenon was also observed in prior work. We study the dependency of policy features on the performance in the appendix, where we cover a diverse set of environments: from non-memory vector-based ones to image-based memory environments.

实验中我们发现，在记忆密集型环境里，使用 $\pi ({\hat{a}}_t\mid z_t,h_t)$ 时，这种差异可能导致策略学习失效；而使用 $\pi ({\hat{a}}_t\mid z_t,x_t)$ 时，策略学习仍能正常进行。 令人意外的是，将所有特征都纳入策略 $\pi ({\hat{a}}_t\mid z_t,h_t,x_t)$ 并不能解决问题。 原因在于这些特征具有非平稳性：随着世界模型不断训练，它们的经验分布会随时间变化，最终导致策略训练不稳定。 以往工作也观察到了类似现象。 我们在附录中研究了策略特征对性能的影响，覆盖从无记忆的向量环境到基于图像的记忆环境等多种场景。

In different environments, we condition the policy and value function on the information from S3M in the following ways: we use the output state policy that takes $(z_t,h_t)$ as input, the hidden state policy that takes $(z_t,x_t)$ as input, and the full state policy that takes $(z_t,h_t,x_t)$ as input. To train actor-critic, we opt for the procedure proposed in DreamerV3.

在不同环境中，我们以以下方式使用 S3M 的信息来调节策略和价值函数：以 $(z_t,h_t)$ 为输入的输出状态策略、以 $(z_t,x_t)$ 为输入的隐状态策略，以及以 $(z_t,h_t,x_t)$ 为输入的完整状态策略。 我们采用 DreamerV3 提出的流程训练行动器和评论器。

4. Experiments

We conduct a comprehensive empirical study to assess the generality and memory capacity of R2I across a wide range of domains, including credit assignment, memory-intensive tasks, and non-memory tasks, all while maintaining fixed hyperparameters of the world model. We cover five RL domains: BSuite, POPGym, Atari 100K, DMC, and Memory Maze. The section is organized as follows. In Section 4.1 and Section 4.2, we evaluate R2I's performance in two distinct memory-intensive settings: simple tabular environments and complex 3D environments.

我们开展了一项全面的实证研究，在固定世界模型超参数的前提下，评估 R2I 在信用分配、记忆密集型任务和无记忆任务等广泛领域中的通用性与记忆容量。 实验覆盖五个强化学习领域：BSuite、POPGym、Atari 100K、DMC 和 Memory Maze。 本节组织如下。 在第 4.1 节和第 4.2 节中，我们在两种不同的记忆密集型设置下评估 R2I：简单表格型环境和复杂三维环境。

We show that not only does R2I achieve the SOTA performance, but it also surpasses human-level performance in the complex Memory Maze domain. In Section 4.3, we demonstrate that we do not trade the generality for improved memory capabilities. Figure 2 shows R2I's impressive computational efficiency, with a speed increase of up to 9 times compared to its predecessor, DreamerV3. Note that the image environments are representative of Memory Maze, and the vector environments represent POPGym.

我们表明，R2I 不仅取得了最先进的性能，还在复杂的 Memory Maze 领域中超越了人类水平。 在第 4.3 节中，我们证明记忆能力的增强并未以牺牲通用性为代价。 图2展示了 R2I 出色的计算效率，其速度相较前身 DreamerV3 最高提升 9 倍。 其中，图像环境代表 Memory Maze，向量环境代表 POPGym。

图2：DreamerV3 与 R2I 的计算用时，越低越好。

We reuse most of the world model hyperparameters from DreamerV3. In all environments, we use a First-in First-out (FIFO) replay buffer size of 10M steps to train R2I. We found this helps stabilize the world model and prevent overfitting on a small buffer.

我们复用了 DreamerV3 的大多数世界模型超参数。 在所有环境中，我们都使用容量为 1000 万步的先进先出（FIFO）回放缓冲区训练 R2I。 我们发现，这有助于稳定世界模型并防止其在小缓冲区上过拟合。

Also, we vary features that the policy is conditioned on (i.e., output state policy $\pi ({\hat{a}}_t\mid z_t,h_t)$, hidden state policy $\pi ({\hat{a}}_t\mid z_t,x_t)$, or full state policy $\pi ({\hat{a}}_t\mid z_t,h_t,x_t)$). Our primary takeaway is to leverage the output state policy in non-memory environments and the full state policy or hidden state policy within memory environments, as explained in Section 3.2. We also found that even in memory environments, the full state policy cannot be preferred over the hidden state policy because of the instability of features — since the world model is trained alongside the policy, the former might change the feature distribution which introduces non-stationarity for the policy.

此外，我们改变策略所依赖的特征，包括输出状态策略 $\pi ({\hat{a}}_t\mid z_t,h_t)$、隐状态策略 $\pi ({\hat{a}}_t\mid z_t,x_t)$ 和完整状态策略 $\pi ({\hat{a}}_t\mid z_t,h_t,x_t)$。 我们的主要结论是：在无记忆环境中使用输出状态策略，在记忆环境中使用完整状态策略或隐状态策略，详见第 3.2 节。 我们还发现，即使在记忆环境中，也不能认为完整状态策略优于隐状态策略，原因在于特征不稳定：世界模型与策略同步训练，前者可能改变特征分布，从而给策略引入非平稳性。

4.1 Quantifying Memory of R2I

In this section, we study the performance of R2I in challenging memory environments of BSuite and POPGym domains, which are tabular environments. Despite their simplicity, these environments pose a challenge for MBRL algorithms since the world model needs to learn causal connections over time. While SSMs have shown their ability to handle extremely long-range dependencies in SL and SSL, this capability does not necessarily translate to MBRL, even though the world model optimizes the same supervised objective. This discrepancy arises from the lifelong nature of world model training.

本节研究 R2I 在 BSuite 和 POPGym 两个表格型领域的高难度记忆环境中的表现。 尽管这些环境很简单，但它们仍对 MBRL 算法构成挑战，因为世界模型需要学习跨时间的因果联系。 SSM 已经在监督学习和自监督学习中展示出处理极长程依赖的能力，但即使世界模型优化相同的监督目标，这种能力也不一定能直接迁移到 MBRL。 这种差异源于世界模型训练的终身性质。

That is, it needs to bootstrap its performance from a very small dataset with hugely imbalanced reward “labels” (as opposed to big and well-balanced long-range datasets on which SSMs shine). Additionally, the continuously growing replay buffer imposes the need to quickly learn the newly arrived data which requires an ability for quick adaptation of the world model throughout its optimization. The section's goal is to give an insight into how extensive are R2I's memory capabilities.

也就是说，它需要从一个奖励“标签”极不平衡的小型数据集开始建立性能，这与 SSM 擅长的大型、均衡长程序列数据集不同。 此外，不断增长的回放缓冲区要求模型快速学习新到达的数据，这要求世界模型在整个优化过程中具备快速适应能力。 本节旨在揭示 R2I 的记忆能力究竟能延伸到何种程度。

图3：DreamerV3（此前的最先进方法）与 R2I 在 BSuite 环境中的成功率。横轴上的每个点都单独训练一个模型；图中绘制 10 个随机种子的中位数，并以填充区域表示第 25 至第 75 百分位数。

Behavior Suite experiments. To study the ability of the R2I model to handle longer episodes, we conduct quantitative experiments within a subset of the BSuite environments. These environments are specifically designed to evaluate an agent's memory capacity and its ability to effectively perform credit assignment. In particular, we carry out experiments within Memory Length and Discounting Chain environments. The former focuses on memory, and the latter serves as a credit assignment task.

Behavior Suite 实验。 为研究 R2I 模型处理更长回合的能力，我们在 BSuite 的部分环境中开展定量实验。 这些环境专门用于评估智能体的记忆容量以及有效执行信用分配的能力。 具体而言，我们在 Memory Length 和 Discounting Chain 环境中进行实验。 前者关注记忆，后者则是一项信用分配任务。

In Memory Length environment, the goal is to output an action which is dictated by the initial observation (the episode length i.e., the memory steps number is an environment parameter). Essentially, the agent must carry the information from the initial observation throughout the entire episode. In the Discounting Chain, the first action (which is categorical) causes a reward that is only provided after a certain number of steps, specified by the parameter reward delay.

在 Memory Length 环境中，目标是输出由初始观察决定的动作；回合长度，也就是记忆步数，是一个环境参数。 本质上，智能体必须在整个回合中始终保留初始观察的信息。 在 Discounting Chain 中，第一个类别动作会引发奖励，但该奖励只有在参数奖励延迟指定的若干步之后才会给出。

As depicted in Figure 3, the previous SOTA DreamerV3 learns the dependencies between actions and rewards in both Discounting Chain and Memory Length with reward delays of up to 30 environment steps. Note that every run either converged to a maximum reward or failed (based on the random seed). We plot the success rate as the fraction of runs that achieved success. R2I excels in both tasks, significantly outperforming in the preservation of its learning ability across a wider range of varying environment complexities. In these experiments, we leverage the output state policy (i.e., operating on latent variable $z_t$ and S3M output $h_t$).

如图3所示，此前最先进的 DreamerV3 能够在 Discounting Chain 和 Memory Length 中学习动作与奖励之间的依赖关系，但奖励延迟最多只能达到 30 个环境步。 请注意，每次运行要么收敛到最大奖励，要么根据随机种子而失败。 我们将成功率绘制为成功运行所占的比例。 R2I 在两项任务中都表现出色，面对范围更广、复杂度不断变化的环境时，仍能保持学习能力，显著优于基线。 在这些实验中，我们采用输出状态策略，即在潜在变量 $z_t$ 和 S3M 输出 $h_t$ 上运行。

POPGym experiments. We perform a study to assess R2I in a more challenging benchmark, namely, POPGym. This suite offers a range of RL environments designed to assess various challenges related to POMDPs, such as navigation, noise robustness, and memory. We select the three most memory-intensive environments: RepeatPrevious, Autoencode, and Concentration. These environments require an optimal policy to memorize the highest number of events (i.e., actions or observations) at each time step. Each environment in POPGym has three difficulty levels: Easy, Medium, and Hard.

POPGym 实验。 我们在更具挑战性的 POPGym 基准中评估 R2I。 该套件提供一系列强化学习环境，用于评估与 POMDP 有关的导航、抗噪性和记忆等多种挑战。 我们选择了三个记忆需求最高的环境：RepeatPrevious、Autoencode 和 Concentration。 这些环境要求最优策略在每个时间步记住最多数量的事件，也就是动作或观察。 POPGym 中的每个环境都有 Easy、Medium 和 Hard 三个难度级别。

In the memory environments of this study, the complexity is increased by the number of actions or observations that the agent should keep track of simultaneously. All environments in this study have categorical observation and action spaces. As POPGym was not included in the DreamerV3 benchmark, we performed hyperparameter tuning of both DreamerV3 and R2I, solely on adjusting the network sizes of both. This is because DreamerV3 is a generalist agent that works with a fixed set of hyperparameters and in this environment, with sizes primarily influencing its data efficiency. We observed a similar characteristic in R2I.

在本研究的记忆环境中，复杂度由智能体需要同时跟踪的动作或观察数量决定。 本研究中的所有环境都具有离散的观察空间和动作空间。 由于 DreamerV3 基准没有包含 POPGym，我们仅通过调整网络规模，对 DreamerV3 和 R2I 都进行了超参数调优。 这是因为 DreamerV3 是使用固定超参数集的通用智能体，而在该环境中，网络规模主要影响其数据效率。 我们在 R2I 中也观察到了类似特征。

For R2I, we use the hidden state policy: $\pi ({\hat{a}}_t\mid z_t,h_t)$ as we found it much more performant, especially in memory-intensive tasks. We train R2I in POPGym environments using a unified and fixed set of hyperparameters. In addition to R2I and DreamerV3, we include model-free baselines from POPGym. These include PPO model-free policy with different observation backbones, such as GRU, LSTM, MLP, and MLP with timestep number added as a feature (PosMLP). PPO with GRU is the best-performing model-free baseline of POPGym while PPO+LSTM is the second best. PPO+MLP and PPO+PosMLP are included for a sanity check — the better their performance is, the less is the memory needed in the environment.

对于 R2I，我们使用隐状态策略 $\pi ({\hat{a}}_t\mid z_t,h_t)$，因为我们发现它的表现要好得多，尤其是在记忆密集型任务中。 我们使用一套统一且固定的超参数，在 POPGym 环境中训练 R2I。 除 R2I 和 DreamerV3 外，我们还加入了 POPGym 的无模型基线。 其中包括采用不同观察骨干的 PPO 无模型策略，例如 GRU、LSTM、MLP，以及将时间步编号作为特征加入 MLP 的 PosMLP。 采用 GRU 的 PPO 是 POPGym 中表现最好的无模型基线，PPO+LSTM 位居第二。 PPO+MLP 和 PPO+PosMLP 用于合理性检查：它们的表现越好，说明该环境所需的记忆越少。

图4：R2I 在 POPGym 记忆密集型环境中的结果。我们的方法在最困难的记忆环境中建立了新的最先进水平：Autoencode 的 Easy、Medium，RepeatPrevious 的 Medium、Hard，以及 Concentration 的 Medium。Concentration 可以在一定程度上不依赖记忆而解决。

As illustrated in Figure 4, R2I demonstrates the new SOTA performance, outperforming every baseline in Autoencode, Easy and Medium tasks. Note that R2I outperforms all 13 model-free baselines of the POPGym benchmark by a huge margin. R2I also shows consistently strong performance in RepeatPrevious tasks, setting a new SOTA in both Medium and Hard. In Concentration, the model-free memory baselines fail to outperform a simple MLP policy, suggesting that they all converge to a non-memory-based suboptimal policy. R2I advances this towards a better memory policy. Its performance is roughly equal to DreamerV3 in an Easy and slightly better in the Medium task.

如图4所示，R2I 展现出新的最先进性能，在 Autoencode 的 Easy 和 Medium 任务上超过了所有基线。 R2I 还以巨大优势超过了 POPGym 基准的全部 13 个无模型基线。 R2I 在 RepeatPrevious 任务中也始终表现强劲，在 Medium 和 Hard 两个难度上都建立了新的最先进水平。 在 Concentration 中，无模型记忆基线无法超过简单的 MLP 策略，这表明它们都收敛到了不依赖记忆的次优策略。 R2I 将其推进到更好的记忆策略。 它在 Easy 任务中的表现与 DreamerV3 大致相当，在 Medium 任务中略好。

All RepeatPrevious tasks require up to $64$ memorization steps, while Autoencode Easy and Medium require up to $104$. In Concentration Easy and Medium this length is up to $208$ steps, however, since PPO+MLP shows somewhat good performance, likely less than $208$ memorization steps are required. This observation is consistent with the results of the BSuite experiments, which demonstrate that our model is capable of memorizing up to approximately 100 steps in time. To summarize, these results indicate that R2I significantly pushes the memory limits.

所有 RepeatPrevious 任务最多需要记忆 $64$ 步，而 Autoencode 的 Easy 和 Medium 最多需要 $104$ 步。 在 Concentration 的 Easy 和 Medium 中，这一长度最多为 $208$ 步；不过，由于 PPO+MLP 的表现尚可，实际需要的记忆步数可能少于 $208$。 这一观察与 BSuite 实验结果一致，后者表明我们的模型能够记住时间上相隔约 100 步的信息。 总之，这些结果表明 R2I 显著拓展了记忆极限。

4.2 Evaluating Long-term Memory In Complex 3D Tasks

Memory Maze presents randomized 3D mazes where the egocentric agent is repeatedly tasked to navigate to one of multiple objects. For optimal speed and efficiency, the agent must retain information about the locations of objects, the maze's wall layout, and its own position. Each episode can extend for up to 4K environment steps. An ideal agent equipped with long-term memory only needs to explore each maze once, a task achievable in a shorter time than the episode's duration; subsequently, it can efficiently find the shortest path to reach each requested target. This task poses a fundamental challenge for existing memory-augmented RL algorithms, which fall significantly behind human performance in these tasks.

Memory Maze 提供随机生成的三维迷宫，以第一人称视角行动的智能体需要反复导航到多个物体中的某一个。 为了达到最佳速度和效率，智能体必须保留物体位置、迷宫墙体布局以及自身位置的信息。 每个回合最多可持续 4000 个环境步。 具备长期记忆的理想智能体只需探索每个迷宫一次，而且所需时间短于回合时长；之后，它就能高效找到到达每个指定目标的最短路径。 该任务对现有记忆增强型强化学习算法构成根本挑战，这些算法在此类任务上的表现明显落后于人类。

In this benchmark, we found that DreamerV3 works equally well as DreamerV2 reported in prior work. Therefore, we use the size configuration of Dreamer outlined in that work. Note that this baseline also leverages truncated backpropagation through time (TBTT), a technique demonstrated to enhance the preservation of information over time. We use the “medium memory” size configuration of R2I in this work. We use the full state policy ($\pi ({\hat{a}}_t\mid z_t,h_t,x_t)$, i.e., conditioning on stochastic state, and S3M output, and hidden states at step $t$) in this environment.

在该基准中，我们发现 DreamerV3 的表现与以往工作报告的 DreamerV2 相当。 因此，我们采用该工作中给出的 Dreamer 规模配置。 请注意，该基线还使用了截断时间反向传播（TBTT），这种技术已被证明能够增强信息随时间保留的能力。 本文使用 R2I 的“中等记忆”规模配置。 在该环境中，我们使用完整状态策略，即 $\pi ({\hat{a}}_t\mid z_t,h_t,x_t)$，策略以随机状态、S3M 输出和时间步 $t$ 的隐状态为条件。

We trained and tested R2I and other methods on 4 existing maze sizes: 9x9, 11x11, 13x13, and 15x15. The difference between them is in the number of object rooms and the episode lengths. More difficult maze sizes have more environment steps in the episode making it more challenging to execute a successful series of object searches. R2I and other baselines are evaluated after 400M environment steps or two weeks of training. We also compare R2I with IMPALA, which is the leading model-free approach.

我们在 9x9、11x11、13x13 和 15x15 四种现有迷宫规模上训练并测试 R2I 和其他方法。 它们的差别在于物体房间数量和回合长度。 难度更高的迷宫在每个回合中包含更多环境步，因此更难成功完成一系列物体搜索。 R2I 和其他基线都在 4 亿个环境步或两周训练之后进行评估。 我们还将 R2I 与领先的无模型方法 IMPALA 进行比较。

图5：经过 4 亿个环境步后的 Memory Maze 分数。R2I 在所有难度级别上都超过基线，成为该领域新的最先进方法。由于计算效率更高，R2I 的训练天数少于 Dreamer。

As shown in Figure 5, R2I consistently outperforms baseline methods in all of these environments. In 9x9 mazes, it demonstrates performance similar to the Dreamer, while significantly outperforming IMPALA. In 11x11, 13x13, and 15x15 mazes, it has a remarkably better performance than both baselines. Moreover, it has surpassed human-level abilities in solving 9x9, 11x11, and 13x13 mazes. These results establish R2I as a SOTA in this complex 3D domain.

如图5所示，R2I 在所有这些环境中都持续优于基线方法。 在 9x9 迷宫中，它的表现与 Dreamer 相近，同时显著超过 IMPALA。 在 11x11、13x13 和 15x15 迷宫中，它的表现明显优于两个基线。 此外，它在 9x9、11x11 和 13x13 迷宫上的求解能力已经超过人类水平。 这些结果确立了 R2I 在这一复杂三维领域中的最先进地位。

4.3 Assessing the Generality of R2I in Non-Memory Domains

We conduct a sanity check by assessing R2I's performance on two widely used RL benchmarks: Atari and DMC, as parts of the DreamerV3 benchmark. Even though these tasks are nearly fully observable and do not necessitate extensive memory to solve (it is often enough to model the dynamics of only the last few steps), evaluating R2I on them is essential as we aim to ensure our agent's performance across a wide range of tasks that require different types of control: continuous control (in DMC) and discrete (in Atari).

我们在两个广泛使用的强化学习基准 Atari 和 DMC 上评估 R2I，进行合理性检查；两者都是 DreamerV3 基准的一部分。 尽管这些任务几乎完全可观测，无需大量记忆即可解决，通常只需建模最近几个步骤的动态，但在这些任务上评估 R2I 仍然至关重要，因为我们希望确保智能体能够应对需要不同控制类型的广泛任务：DMC 中的连续控制和 Atari 中的离散控制。

图6：R2I 在 Atari 和 DMC 上的平均性能。

In all the experiments conducted within Atari 100K and DMC, we fix hyperparameters of the world model. In Atari and the proprio benchmark in DMC, we utilize output state policies, as we found them more performant. In the visual benchmark in DMC, we use hidden state policy. Note that for continuous control, the policy is trained via differentiating through the learned dynamics. R2I maintains a performance similar to DreamerV3 in these domains, as demonstrated in Figure 6, implying that in the majority of standard RL tasks, R2I does not sacrifice generality for improved memory capabilities.

在 Atari 100K 和 DMC 的所有实验中，我们固定世界模型的超参数。 在 Atari 和 DMC 的本体感知基准中，我们采用输出状态策略，因为它们的表现更好。 在 DMC 的视觉基准中，我们使用隐状态策略。 请注意，对于连续控制，策略通过对学习到的动态进行微分来训练。 如图6所示，R2I 在这些领域中保持了与 DreamerV3 相近的性能，这意味着在大多数标准强化学习任务中，R2I 不会为增强记忆能力而牺牲通用性。

5. Conclusion

In this paper, we introduced R2I, a general and fast model-based approach to reinforcement learning that demonstrates superior memory capabilities. R2I integrates two strong algorithms: DreamerV3, a general-purpose MBRL algorithm, and SSMs, a family of novel parallelizable sequence models adept at handling extremely long-range dependencies. This integration helps rapid long-term memory and long-horizon credit assignment, allowing R2I to excel across a diverse set of domains, all while maintaining fixed hyperparameters across all domains. Through a systematic examination, we have demonstrated that R2I sets a new state-of-the-art in domains demanding long-term temporal reasoning: it outperforms all known baselines by a large margin on the most challenging memory and credit assignment tasks across different types of memory (long-term and short-term) and observational complexities (tabular and complex 3D).

本文提出 R2I，这是一种通用且快速的基于模型的强化学习方法，展现出卓越的记忆能力。 R2I 集成了两种强大算法：通用 MBRL 算法 DreamerV3，以及一类擅长处理极长程依赖、可并行化的新型序列模型 SSM。 这种集成有助于实现快速的长期记忆和长程信用分配，使 R2I 能够在多个领域中表现出色，同时在所有领域保持固定超参数。 通过系统检验，我们证明 R2I 在需要长期时间推理的领域中建立了新的最先进水平：面对不同记忆类型，包括长期与短期记忆，以及不同观察复杂度，包括表格型与复杂三维观察的高难度记忆和信用分配任务，它都以巨大优势超过所有已知基线。

Remarkably, it transcends human performance in complex 3D tasks. Furthermore, we have demonstrated that R2I achieves computation speeds up to 9 times faster than DreamerV3.

尤其值得注意的是，它在复杂三维任务中超越了人类表现。 此外，我们还证明 R2I 的计算速度最高可达到 DreamerV3 的 9 倍。

Our study presents the first model-based RL approach that uses SSMs. While R2I offers benefits for improving memory in RL, it also has limitations, which we leave for future research. For instance, it can be explored how R2I can be augmented with attention mechanisms, given that Transformers and SSMs exhibit complementary strengths. Hybrid architectures have been introduced in language modeling tasks. Moreover, the sequence length within the training batches for world model learning is not currently extremely long, as is the horizon (i.e., the number of steps) of imagination in actor-critic learning. Future work could focus on these aspects to further enhance memory capabilities.

我们的研究提出了首个使用 SSM 的基于模型的强化学习方法。 R2I 虽然有助于提升强化学习中的记忆能力，但也存在一些局限，我们将其留待未来研究。 例如，鉴于 Transformer 与 SSM 各有互补优势，可以探索如何使用注意力机制增强 R2I。 语言建模任务中已经出现了混合架构。 此外，目前世界模型学习的训练批次序列长度还不是特别长，行动器—评论器学习中的想象范围，也就是步数，同样如此。 未来工作可以聚焦这些方面，进一步增强记忆能力。