SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Lucy Xiaoyang Shi $^{1*}$ , Yunfan Jiang $^{1*}$ , Jake Grigsby $^{2}$ , Linxi “Jim” Fan $^{3\dagger}$ , Yuke Zhu $^{23\dagger}$ $^{1}$ Stanford University  $^{2}$ The University of Texas at Austin  $^{3}$ NVIDIA Research  $^{*}$ Equal contribution  $^{+}$ Equal advising
We present a new algorithm, Cross-Episodic Curriculum (CEC), to boost the learning efficiency and generalization of Transformer agents. Central to CEC is the placement of cross-episodic experiences into a Transformer's context, which forms the basis of a curriculum. By sequentially structuring online learning trials and mixed-quality demonstrations, CEC constructs curricula that encapsulate learning progression and proficiency increase across episodes. Such synergy combined with the potent pattern recognition capabilities of Transformer models delivers a powerful cross-episodic attention mechanism. The effectiveness of CEC is demonstrated under two representative scenarios: one involving multi-task reinforcement learning with discrete control, such as in DeepMind Lab, where the curriculum captures the learning progression in both individual and progressively complex settings, and the other involving imitation learning with mixed-quality data for continuous control, as seen in RoboMimic, where the curriculum captures the improvement in demonstrators' expertise. In all instances, policies resulting from CEC exhibit superior performance and strong generalization. Code is open-sourced on the project website cec-agent.github.io to facilitate research on Transformer agent learning.
The paradigm shift driven by foundation models [8] is revolutionizing the communities who study sequential decision-making problems [80], with innovations focusing on control [2, 45, 38, 9], planning [76, 32, 33, 78, 17], pre-trained visual representation [57, 50, 67, 51, 82], among others. Despite the progress, the data-hungry nature makes the application of Transformer [75] agents extremely challenging in data-scarce domains like robotics [52, 53, 19, 38, 9]. This leads us to the question: Can we maximize the utilization of limited data, regardless of their optimality and construction, to foster more efficient learning?
To this end, this paper introduces a novel algorithm named Cross-Episodic Curriculum (CEC), a method that explicitly harnesses the shifting distributions of multiple experiences when organized into a curriculum. The key insight is that sequential cross-episodic data manifest useful learning signals that do not easily appear in any separated training episodes. As illustrated in Figure 1, CEC realizes this through two stages: 1) formulating curricular sequences to capture (a) the policy improvement on single environments, (b) the learning progress on a series of progressively harder environments, or (c) the increase of demonstrators' proficiency; and 2) causally distilling policy improvements into the model weights of Transformer agents through cross-episodic attention. When a policy is trained to predict actions at current time steps, it can trace back beyond ongoing trials and internalize improved behaviors encoded in curricular data, thereby achieving efficient learning
$^{1}$ Following the canonical definition in Sutton and Barto [73], we refer to the sequences of agent-environment interaction with clearly identified initial and terminal states as "episodes". We interchangeably use "episode", "trial", and "trajectory" in this work.
Curricular Data
Policy improvement on single environments:
Learning progress on a series of progressively harder environments:
Increase of demonstrator's proficiency:
Figure 1: Cross-episodic curriculum for Transformer agents. CEC involves two major steps: 1) Preparation of curricular data. We order multiple experiences such that they explicitly capture curricular patterns. For instance, they can be policy improvement in single environments, learning progress in a series of progressively harder environments, or the increase of the demonstrator's expertise. 2) Model training with cross-episodic attention. When training the model to predict actions, it can trace back beyond the current episode and internalize the policy refinement for more efficient learning. Here each  $\tau$  represents an episode (trajectory).  $\hat{a}$  refers to actions predicted by the model. Colored triangles denote causal Transformer models.
Model Training with Cross-Episodic Attention
and robust deployment when probed with visual or dynamics perturbations. Contrary to prior works like Algorithm Distillation (AD, Laskin et al. [42]) which, at test time, samples and retains a single task configuration across episodes for in-context refinement, our method, CEC, prioritizes zero-shot generalization across a distribution of test configurations. With CEC, agents are evaluated on a new task configuration in each episode, emphasizing adaptability to diverse tasks.
We investigate the effectiveness of CEC in enhancing sample efficiency and generalization with two representative case studies. They are: 1) Reinforcement Learning (RL) on DeepMind Lab (DMLab) [5], a 3D simulation encompassing visually diverse worlds, complicated environment dynamics, ego-centric pixel inputs, and joystick control; and 2) Imitation Learning (IL) from mixed-quality human demonstrations on RoboMimic [53], a framework designed to study robotic manipulation with proprioceptive and external camera observations and continuous control. Despite RL episodes being characterized by state-action-reward tuples and IL trajectories by state-action pairs, our method exclusively employs state-action pairs in its approach.
In challenging embodied navigation tasks, despite significant generalization gaps (Table 1), our method surpasses concurrent and competitive method Agentic Transformer (AT, Liu and Abbeel [47]). It also significantly outperforms popular offline RL methods such as Decision Transformer (DT, Chen et al. [13]) and baselines trained on expert data, with the same amount of parameters, architecture, and data size. It even exceeds RL oracles directly trained on test task distributions by  $50\%$  in a zero-shot manner. CEC also yields robust embodied policies that are up to  $1.6\times$  better than RL oracles when zero-shot probed with unseen environment dynamics. When learning continuous robotic control, CEC successfully solves two simulated manipulation tasks, matching and outperforming previous well-established baselines [53, 25, 41]. Further ablation reveals that CEC with cross-episodic attention is a generally effective recipe for learning Transformer agents, especially in applications where sequential data exhibit moderate and smooth progression.
2 Cross-Episodic Curriculum: Formalism and Implementations
In this section, we establish the foundation for our cross-episodic curriculum method by first reviewing the preliminaries underlying our case studies, which encompass two representative scenarios in sequential decision-making. Subsequently, we formally introduce the assembly of curricular data and the specifics of model optimization utilizing cross-episodic attention. Lastly, we delve into the practical implementation of CEC in the context of these two scenarios.
Reinforcement learning. We consider the setting where source agents learn through trial and error in partially observable environments. Denoting states  $s \in S$  and actions  $a \in \mathcal{A}$ , an agent interacts in a Partially Observable Markov Decision Process (POMDP) with the transition function  $p(s_{t+1}|s_t, a_t) : S \times \mathcal{A} \to S$ . It observes  $o \in \mathcal{O}$  emitted from observation function  $\Omega(o_t|s_t, a_{t-1}) : S \times \mathcal{A} \to \mathcal{O}$  and receives scalar reward  $r$  from  $R(s, a) : S \times \mathcal{A} \to \mathbb{R}$ . Under the episodic task setting, RL seeks to learn a parameterized policy  $\pi_\theta(\cdot|s)$  that maximizes the return over a fixed length  $T$  of interaction steps:  $\pi_\theta = \arg \max_{\theta \in \Theta} \sum_{t=0}^{T-1} \gamma^t r_t$ , where  $\gamma \in [0,1)$  is a discount factor. Here we follow the canonical definition of an episode  $\tau$  as a series of environment-agent interactions with length  $T$ ,  $\tau := (s_0, a_0, r_0, \ldots, s_{T-1}, a_{T-1}, r_{T-1}, s_T)$ , where initial states  $s_0$  are sampled from initial state distribution  $s_0 \sim \rho_0(s)$  and terminal states  $s_T$  are reached once the elapsed timestep exceeds  $T$ . Additionally, we view all RL tasks considered in this work as goal-reaching problems [39, 26] and constrain all episodes to terminate upon task completion. It is worth noting that similar to previous work [42], training data are collected by source RL agents during their online learning. Nevertheless, once the dataset is obtained, our method is trained offline in a purely supervised manner.
Imitation learning. We consider IL settings with existing trajectories composed only of state-action pairs. Furthermore, we relax the assumption on demonstration optimality and allow them to be crowdsourced [10, 12, 11]. Data collected by operators with varying expertise are therefore unavoidable. Formally, we assume the access to a dataset  $\mathcal{D}^N \coloneqq \{\tau_1, \dots, \tau_N\}$  consisting of  $N$  demonstrations, with each demonstrated trajectory  $\tau_i \coloneqq (s_0, a_0, \dots, s_{T-1}, a_{T-1})$  naturally identified as an episode. The goal of IL, specifically of behavior cloning (BC), is to learn a policy  $\pi_\theta$  that accurately models the distribution of behaviors. When viewed as goal-reaching problems, BC policies can be evaluated by measuring the success ratio in completing tasks [26].
2.2 Curricular Data Assembly and Model Optimization
Meaningful learning signals emerge when multiple trajectories are organized and examined cross-episodically along a curriculum axis. This valuable information, which is not easily discernible in individual training episodes, may encompass aspects such as the improvement of an RL agent's navigation policy or the generally effective manipulation skills exhibited by operators with diverse proficiency levels. With a powerful model architecture such as Transformer [75, 16], such emergent and valuable learning signals can be baked into policy weights, thereby boosting performance in embodied tasks.
For a given embodied task  $\mathcal{M}$ , we define its curriculum  $\mathcal{C}_{\mathcal{M}}$  as a collection of trajectories  $\tau$  consisting of state-action pairs. A series of ordered levels  $[\mathcal{L}_1, \ldots, \mathcal{L}_L]$  partitions this collection such that  $\bigcup_{l \in \{1, \ldots, L\}} \mathcal{L}_l = \mathcal{C}_{\mathcal{M}}$  and  $\bigcap_{\forall i, j \in \{1, \ldots, L\}, i \neq j} \mathcal{L}_{\{i, j\}} = \emptyset$ . More importantly, these ordered levels characterize a curriculum by encoding, for example, learning progress in single environments, learning progress in a series of progressively harder environments, or the increase of the demonstrator's expertise.
With a curriculum  $\mathcal{C}_{\mathcal{M}} := \{\tau_i\}_{i=1}^N$  and its characteristics  $[\mathcal{L}_1, \ldots, \mathcal{L}_L]$ , we construct a curricular sequence  $\mathcal{T}$  that spans multiple episodes and captures the essence of gradual improvement in the following way:
$$
\mathcal {T} := \bigoplus_ {l \in \{1, \dots , L \}} \left[ \tau^ {(1)}, \dots , \tau^ {(C)} \right], \quad \text {w h e r e} \quad C \sim \mathcal {U} \left(\llbracket \mathcal {L} _ {l} \rrbracket\right) \quad \text {a n d} \quad \tau^ {(c)} \sim \mathcal {L} _ {l}. \tag {1}
$$
The symbol  $\oplus$  denotes the concatenation operation.  $\mathcal{U}([K])$  denotes a uniform distribution over the discrete set  $\{k\in \mathbb{N},k\leq K\}$ . In practice, we use values smaller than  $|\mathcal{L}_l|$  considering the memory consumption.
(a) Goal Maze
(b) Watermaze
(c) Irreversible Path
Figure 2: We evaluate our method on five tasks that cover challenges such as exploration and planning over long horizons in RL settings, as well as object manipulation and continuous control in IL settings. Figures are from Beattie et al. [5] and Mandlekar et al. [53].
(d) Lift
(e) Can
We subsequently learn a causal policy that only depends on cross-episodic historical observations  $\pi_{\theta}(\cdot |o_{\leq t}^{(\leq n)})$ . Note that this modeling strategy differs from previous work that views sequential decision-making as a big sequence-modeling problem [13, 37, 42, 38]. It instead resembles the causal policy in Baker et al. [4]. Nevertheless, we still follow the best practice [36, 60, 22] to provide previous action as an extra modality of observations in POMDP RL tasks.
We leverage the powerful attention mechanism of Transformer [75] to enable cross-episodic attention. Given observation series  $O_{t}^{(n)} \coloneqq \{o_{0}^{(1)},\ldots ,o_{\leq t}^{(\leq n)}\}$  (shorthand as  $O$  hereafter for brevity), Transformer projects it into query  $Q = f_{Q}(O)$ , key  $K = f_{K}(O)$ , and value  $V = f_{V}(O)$  matrices, with each row being a  $D$ -dim vector. Attention operation is performed to aggregate information:
$$
\operatorname {A t t e n t i o n} (Q, K, V) = \operatorname {s o f t m a x} \left(\frac {Q K ^ {\intercal}}{\sqrt {D}}\right) V. \tag {2}
$$
Depending on whether the input arguments for  $f_{Q}$  and  $f_{\{K,V\}}$  are the same, attention operation can be further divided into self-attention and cross-attention. Since tasks considered in this work do not require additional conditioning for task specification, we follow previous work [4, 83] to utilize self-attention to process observation series. Nevertheless, ours can be naturally extended to handle, for example, natural language or multi-modal task prompts, following the cross-attention introduced in Jiang et al. [38].
Finally, this Transformer policy is trained by simply minimizing the negative log-likelihood objective  $\mathcal{J}_{\mathrm{NLL}}$  of labeled actions, conditioned on cross-episodic context:
$$
\mathcal {J} _ {\mathrm {N L L}} = - \log \pi_ {\theta} (\cdot | \mathcal {T}) = \frac {1}{| \mathcal {T} | \times T} \sum_ {n = 1} ^ {| \mathcal {T} |} \sum_ {t = 1} ^ {T} - \log \pi_ {\theta} \left(a _ {t} ^ {(n)} | o _ {\leq t} ^ {(\leq n)}\right). \tag {3}
$$
Regarding the specific memory architecture, we follow Baker et al. [4], Adaptive Agent Team et al. [1] to use Transformer-XL [16] as our model backbone. Thus, during deployment, we keep its hidden states propagating across test episodes to mimic the training settings.
2.3 Practical Implementations
We now discuss concrete instantiations of CEC for 1) RL with DMLab and 2) IL with RoboMimic. Detailed introductions to the benchmark and task selection are deferred to Sec. 3. We investigate the following three curricula, where the initial two pertain to RL, while the final one applies to IL:
Learning-progress-based curriculum. In the first instantiation, inspired by the literature on learning progress [54, 27, 65, 40], we view the progression of learning agents as a curriculum. Concretely, we train multi-task PPO agents [70, 63] on tasks drawn from test distributions. We record their online interactions during training, which faithfully reflect the learning progress. Finally, we form the learning-progress-based curriculum by sequentially concatenating episodes collected at different learning stages. Note that this procedure is different from Laskin et al. [42], where for each environment, the learning dynamics of multiple single-task RL agents has to be logged. In contrast, we only track a single multi-task agent per environment.
Task-difficulty-based curriculum. In the second instantiation, instead of taking snapshots of RL agents directly trained on test configurations, we collect learning progress on a series of easier
Table 1: Generalization gaps between training and testing for DMLab levels. Note that agents resulting from task-difficulty-based curricula are not trained on test configurations. Therefore, their performance should be considered as zero-shot.
Level Name	Difficulty Parameter	Test Difficulty	Training Difficulty
Ours (Learning Progress)	Ours (Task Difficulty)	BC w/ Expert Data	RL (Oracle)	Curriculum RL (Oracle)
Goal Maze	Room Numbers	20	20	5→10→15	20	20	5→10→15→20
Watermaze	Spawn Radius	580	580	150→300→450	580	580	150→300→450→580
Irreversible Path	Built-In Difficulty	.9	.9	.1→.3→.5→.7	.9	.9	.1→.3→.5→.7→.9
but progressively harder tasks. For instance, in an embodied navigation task, the test configuration includes 20 rooms. Rather than logging source agents' learning progression in the 20-room maze, we record in a series of mazes with 5, 10, and 15 rooms. We then structure stored episodes first following learning progress and then the increase of layout complexity. This practice naturally creates a task-difficulty-based curriculum, which resembles curriculum RL that is based on task difficulty [54, 58]. We find it especially helpful for hard-exploration problems where the source RL agent does not make meaningful progress.
Expertise-based curriculum. For the setting of IL from mixed-quality demonstrations, we instantiate a curriculum based on demonstrators' expertise. This design choice is motivated by literature on learning from heterogeneous demonstrators [6, 81], with the intuition that there is little to learn from novices but a lot from experts. To realize this idea, we leverage the Multi-Human dataset from RoboMimic [53]. Since it contains demonstrations collected by human demonstrators with varying proficiency, we organize offline demonstration trajectories following the increase of expertise to construct the expertise-based curriculum.
In this section, we elaborate on the experimental setup of our case studies. Our investigation spans two representative and distinct settings: 1) online reinforcement learning with 3D maze environments of DMLab [5], and 2) imitation learning from mixed-quality human demonstrations of RoboMimic [53]. For each of them, we discuss task selection, baselines, and training and evaluation protocols. Teasers of these tasks are shown in Figure 2.
3.1 Task Settings and Environments
DeepMind Lab [5] is a 3D learning environment with diverse tasks. Agents spawn in visually complex worlds, receive ego-centric (thus partially observable) RGB pixel inputs, and execute joystick actions. We consider three levels from this benchmark: Goal Maze, Watermaze [56], and Sky Maze with Irreversible Path. They challenge agents to explore, memorize, and plan over a long horizon. Their goals are similar — to navigate in complicated mazes and find a randomly spawned goal, upon which sparse rewards will be released. Episodes start with randomly spawned agents and goals and terminate once goals are reached or elapsed steps have exceeded pre-defined horizons.
RoboMimic [53] is a framework designed for studying robot manipulation and learning from demonstrations. Agents control robot arms with fixed bases, receive proprioceptive measurements and image observations from mounted cameras, and operate with continuous control. We evaluate two simulated tasks: "Lift" and "Can". In the "Lift" task, robots are tasked with picking up a small cube. In the "Can" task, robots are required to pick up a soda can from a large bin and place it into a smaller target bin. Episodes start with randomly initialized object configuration and terminate upon successfully completing the task or exceeding pre-defined horizons.
The primary goal of these case studies is to assess the effectiveness of our proposed cross-episodic curriculum in increasing the sample efficiency and boosting the generalization capability of Transformer agents. Therefore, in online RL settings, we compare against source RL agents which generate training data for our method and refer to them as oracles. These include a) PPO agents directly
Figure 3: Evaluation results on DMLab. Our CEC agents perform comparable to RL oracles and on average outperform other baseline methods. On the hardest task Irreversible Path where the RL oracle and BC baseline completely fail, our agents outperform the curriculum RL oracle by  $50\%$  even in a zero-shot manner. For our methods, DT, AT, and the BC w/ expert data baselines, we conduct 20 independent evaluation runs, each consisting of 100 episodes for Goal Maze and Watermaze and 50 episodes for Irreversible Path due to longer episode length. We test RL oracles for 100 episodes. The error bars represent the standard deviations over 20 runs.
trained on test task distributions, denoted as “RL (Oracle)” hereafter, and b) curriculum PPO agents that are gradually adapted from easier tasks to the test difficulty, which is referred to as “Curriculum RL (Oracle)”. Furthermore, we compare against one concurrent and competitive method Agen-tic Transformer [47], denoted as “AT”. It is closely related to our method, training Transformers on sequences of trajectory ascending sorted according to their rewards. We also compare against popular offline RL method Decision Transformer [13], denoted as “DT”. Additionally, we include another behavior cloning agent that has the same model architecture as ours but is trained on optimal data without cross-episodic attention. This baseline is denoted as “BC w/ Expert Data”. For the case study on IL from mixed-quality demonstrations, we adopt the most competing approach, BC-RNN, from Mandlekar et al. [53] as the main baseline. We also include comparisons against other offline RL methods [44] such as Batch-Constrained Q-learning (BCQ) [25] and Conservative Q-Learning (CQL) [41].
3.3 Training and Evaluation
We follow the best practice to train Transformer agents, including adopting AdamW optimizer [49], learning rate warm-up and cosine annealing [48], etc. Training is performed on NVIDIA V100 GPUs. During evaluation, for agents resulting from our method, each run involves several test rollouts to fill the context. We keep hidden states of Transformer-XL [16] propagating across episodes. We run other baselines and oracles for 100 episodes to estimate their performances. For our methods on RL settings, we compute the maximum success rate averaged across a sliding window over all test episodes to account for in-context improvement. The size of the sliding window equals one-quarter of the total test episodes. These values are averaged over 20 runs to constitute the final reporting metric. For our methods on the IL setting, since all training data are successful trajectories, we follow Mandlekar et al. [53] to report the maximum success rate achieved over the course of training, directly averaged over test episodes.
We aim to answer the following four research questions through comprehensive experiments.
1. To what extent can our cross-episodic curriculum increase the sample efficiency of Transformer agents and boost their generalization capability?
2. Is CEC consistently effective and generally applicable across distinct learning settings?
3. What are the major components that contribute to the effectiveness of our method?
Figure 4: Generalization results on DMLab. Top row: Evaluation results on Goal Maze with unseen maze mechanism and Irreversible Path with out-of-distribution difficulty levels. Bottom row: Evaluation results on three levels with environment dynamics differing from training ones. CEC agents display robustness and generalization across various dimensions, outperforming curriculum RL oracles by up to  $1.6 \times$ . We follow the same evaluation protocol as in Figure 3. The error bars represent the standard deviations over 20 runs.
4.1 Main Evaluations
We answer the first two questions above by comparing learned agents from our method against 1) Reinforcement Learning (RL) oracles in online RL settings and 2) well-established baselines on learning from mixed-quality demonstrations in the Imitation Learning (IL) setting.