Original: Dynamics-Augmented Decision Transformer for Offline Dynamics Generalization
Revised: Dynamics-Augmented Decision Transformer for Offline Dynamics Generalization

Original: Anonymous Author(s)
Revised: Anonymous Author(s)

Original: AffiliationAddress
Revised: AffiliationAddress

Original: Abstract
Revised: Abstract

Original: Recent progress in offline reinforcement learning (RL) has shown that it is oftenpossible to train strong agents without potentially unsafe or impractical onlineinteraction.
Revised: Recent progress in offline reinforcement learning (RL) has shown that it is oftenpossible to train strong agents without potentially unsafe or impractical onlineinteraction.

Original: However, in real-world settings, agents may encounter unseen environments with different dynamics, and generalization ability is required.
Revised: However, in real-world settings, agents may encounter unseen environments with different dynamics, and generalization ability is required.

Original: Thiswork presents Dynamics-Augmented Decision Transformer (DADT), a simpleyet efficient method to train generalizable agents from offline datasets; on top ofreturn-conditioned policy using the transformer architecture, we improve generalization capabilities by using representation learning based on next state prediction.
Revised: Thiswork presents Dynamics-Augmented Decision Transformer (DADT), a simpleyet efficient method to train generalizable agents from offline datasets; on top ofreturn-conditioned policy using the transformer architecture, we improve generalization capabilities by using representation learning based on next state prediction.

Original: Our experimental results demonstrate that DADT outperforms prior state-of-theart methods for offline dynamics generalization.
Revised: Our experimental results demonstrate that DADT outperforms prior state-of-theart methods for offline dynamics generalization.

Original: Intriguingly, DADT withoutfine-tuning even outperforms fine-tuned baselines.
Revised: Intriguingly, DADT withoutfine-tuning even outperforms fine-tuned baselines.

Original:  Introduction
Revised:  Introduction

Original: Offline reinforcement learning (RL) provides a framework to train robotic agents from a loggeddataset without online interaction (Lange et al., 2012; Levine et al., 2020).
Revised: Offline reinforcement learning (RL) provides a framework to train robotic agents from a loggeddataset without online interaction (Lange et al., 2012; Levine et al., 2020).

Original: The benefit of offline
Revised: The benefit of offline

Original: RL enables us to utilize RL for complex robotics domains where environment interactions can beexpensive or dangerous (Dulac-Arnold et al., 2019).
Revised: RL enables us to utilize RL for complex robotics domains where environment interactions can beexpensive or dangerous (Dulac-Arnold et al., 2019).

Original: However, for practical applications, besides theability to learn without interaction, the ability to generalize to unseen environments with (slightly)different dynamics is required.
Revised: However, for practical applications, besides theability to learn without interaction, the ability to generalize to unseen environments with (slightly)different dynamics is required.

Original: Unfortunately, it has been evidenced that offline RL often struggles togeneralize to unseen dynamics (Li et al., 2020; Lin et al., 2022).
Revised: Unfortunately, it has been evidenced that offline RL often struggles togeneralize to unseen dynamics (Li et al., 2020; Lin et al., 2022).

Original: Several approaches have been proposed recently to address dynamics generalization in offline RL,including representation learning for dynamics inference (Li et al., 2020) and meta-learning (Linet al., 2022).
Revised: Several approaches have been proposed recently to address dynamics generalization in offline RL,including representation learning for dynamics inference (Li et al., 2020) and meta-learning (Linet al., 2022).

Original: Specifically, Li et al.
Revised: Specifically, Li et al.

Original: (2020) trains a context encoder for efficient dynamics inferenceon top of offline RL, and Lin et al.
Revised: (2020) trains a context encoder for efficient dynamics inferenceon top of offline RL, and Lin et al.

Original: (2022) learns meta-dynamics model and meta-policy that adaptto unseen dynamics in a few gradient steps.
Revised: (2022) learns meta-dynamics model and meta-policy that adaptto unseen dynamics in a few gradient steps.

Original: However, whether a simple gradient update or contextlatent variable from near past experiences could capture subtle changes in differing dynamics isquestionable, which is essential for solving dynamics generalization problems.
Revised: However, whether a simple gradient update or contextlatent variable from near past experiences could capture subtle changes in differing dynamics isquestionable, which is essential for solving dynamics generalization problems.

Original: Turning to vision and language domains, transformer-based models (Vaswani et al., 2017) haveexhibited impressive results in solving various generalization problems.
Revised: Turning to vision and language domains, transformer-based models (Vaswani et al., 2017) haveexhibited impressive results in solving various generalization problems.

Original: As a representative example,
Revised: As a representative example,

Original: Brown et al.
Revised: Brown et al.

Original: (2020), Rae et al.
Revised: (2020), Rae et al.

Original: (2021), and Chowdhery et al.
Revised: (2021), and Chowdhery et al.

Original: (2022) show that large language modelusing transformer architecture can solve various types of unseen tasks such as translation, questionanswering, even without any demonstration.
Revised: (2022) show that large language modelusing transformer architecture can solve various types of unseen tasks such as translation, questionanswering, even without any demonstration.

Original: Ramesh et al.
Revised: Ramesh et al.

Original: (2021, 2022) demonstrate that transformerbased models jointly trained with the text and image tokens can generate photorealistic images fromunseen short-length captions.
Revised: (2021, 2022) demonstrate that transformerbased models jointly trained with the text and image tokens can generate photorealistic images fromunseen short-length captions.

Original: Recently, Reed et al.
Revised: Recently, Reed et al.

Original: (2022) show that transformers can be a generalistagent that can solve a range of control tasks using inputs from different modalities with a single setof weights. 
Revised: (2022) show that transformers can be a generalistagent that can solve a range of control tasks using inputs from different modalities with a single setof weights. [Figure] [Figure]

Original: Contribution.
Revised: Contribution.

Original: Inspired by the generalization ability of transformer-based models in vision (Rameshet al., 2022;
Revised: Inspired by the generalization ability of transformer-based models in vision (Rameshet al., 2022;

Original: Dosovitskiy et al., 2020) and language domains (Brown et al., 2020;
Revised: Dosovitskiy et al., 2020) and language domains (Brown et al., 2020;

Original: Devlin et al., 2018),we present DADT:
Revised: Devlin et al., 2018),we present DADT:

Original: D ynamics-
Revised: D ynamics-

Original: A ugmented D ecision T ransformer, a  transformer-based method foroffline dynamics generalization.
Revised: A ugmented D ecision T ransformer, a new transformer-based methodfor offline dynamics generalization.

Revised: Analogously to pre-training scheme in transformer-based modelsin vision and language domains, we train DADT via a supervised learning paradigm with totallyaggregated sequences (i.e., trajectories) without distinction.

Revised: Specifically, DADT aims to optimize theobjective to find the action that achieves desired return (Chen et al., 2021) while minimizing errorin next state prediction.

Revised: We find that DADT can better generalize to unseen dynamics by explicitlyencouraging the model to learn dynamics information.

Revised: For evaluation, we consider two scenariosdepending on different accessibility to trajectories from unseen dynamics: (1) zero-shot, wheredemonstrations are not given, and (2) adaptation, where model is fine-tuned with small amount ofsamples.

Revised: We provide an overview of setup and architecture in Figure 1 and Figure 2, respectively.

Original: We show that DADT outperforms prior state-of-the-art baselinesfor offline dynamics generalization in both performance and time efficiency on various continuouscontrol tasks based on MuJoCo simulator (Todorov et al., 2012)
Revised: We show that DADT outperforms prior state-of-the-art baselines for offline dynamics generalizationin both performance and time efficiency on various continuous control tasks based on MuJoCosimulator (Todorov et al., 2012), which is widely used in the literature (Li et al., 2020; Lin et al.,2022).

Original: DADT: Dynamics-Augmented Decision Transformer
Revised: Intriguingly, DADT without fine-tuning even outperforms fine-tuned baselines. Moreover,reduced next state prediction error in unseen dynamics further supports our model’s capability fordynamics generalization.  Related Work 2.1 Dynamics generalization and adaptation

Revised: Recent works handling dynamics generalization focus on encoding inductive bias using prior knowledge (Zambaldi et al., 2018) or learning contextual latent encoder that captures local dynamics(Lee et al., 2020; Seo et al., 2020).

Revised: Apart from this, several meta-learning approaches to dynamicsgeneralization have been proposed for quickly adapting to new tasks (Gupta et al., 2018; Rakellyet al., 2019; Zintgraf et al., 2019).

Revised: Rakelly et al.

Revised: (2019) introduced probabilistic context variableaccumulating past experiences for efficient exploration and fast task inference.

Revised: Recently, Melo (2022)uses transformer architecture for capturing contextual information instead of recurrent neural network

Revised: Duan et al.

Revised: (2016) in online dynamics generalization.

Revised: It differs from our work in that (i) we usetransformer architecture not only for capturing contextual information but also for optimizing policyconditioned on desired return, and (ii) our model utilizes offline datasets.

Revised: Along with increasing attention toward offline RL (Levine et al., 2020) of which importance ishighlighted especially when environmental interaction is expensive or risky, dynamics generalizationand adaptation in the offline setting have been studied recently (Li et al., 2020;

Revised: Lin et al., 2022;

Revised: Ball et al., 2021; Cang et al., 2021).

Revised: They aim to learn a policy or dynamics model that can quickly adapt to unseen dynamics along with mitigating value overestimation issues in offline RL (Fujimotoet al., 2019).

Revised: Specifically, FOCAL (Li et al., 2020) trains a deterministic context encoder for efficientdynamics inference using contrastive loss and use it for behavior regularized policy learning (Wuet al., 2019).

Revised: MerPO (Lin et al., 2022) takes a gradient-based meta-learning approach for learningmeta-dynamics model and meta-policy that aim to adapt to unseen dynamics quickly; the metadynamics model is used to generate synthetic rollouts to boost policy learning.

Revised: Ball et al.

Revised: (2021)augment a learned dynamics model to improve the zero-shot generalization in model-based offline

Revised: RL.

Revised: 2.2 RL via sequence modelling

Original: {+Motivated by remarkable success in large-scale language modeling (Devlin et al., 
+}
Revised: {+Motivated by remarkable success in large-scale language modeling (Devlin et al., 
+}

Revised: ; Brown et al.

Revised: ,2020), it has been proposed to extend such success into RL domain by formulating RL as a sequencemodeling problem (Chen et al., 2021; Janner et al., 2021).

Revised: These works build on the reinforcementlearning as supervised learning paradigm (Schmidhuber, 2019;

Revised: Srivastava et al., 2019; Emmons et al.

Revised: ,2021) that focuses on predictive modeling of sequence of states, actions, and reward-based returns.

Revised: Specifically, Chen et al.

Revised: (2021) proposed Decision Transformer (DT), which is trained to generatethe optimal actions given desired return and history of states and actions.

Revised: Instead of conditioningdesired return, Trajectory Transformer (TT) (Janner et al., 2021) is trained to predict not only actions,but also entire elements of trajectories such as states and rewards and finds optimal actions by usingplanning (e.g., beam search) in deployment time.

Revised: Recently, Lee et al.

Revised: (2022) have demonstrated that

Revised: DT can play multiple Atari games with a single set of weights and can be efficiently fine-tuned forsolving unseen games.

Revised: Though DT and TT have shown promising results in solving various control tasks, their ability togeneralize to unseen dynamics has not investigated yet.

Revised: In this paper, we demonstrate that DT trainedon aggregated trajectories collected from different dynamics can generalize to unseen dynamics, andintroducing an auxiliary next state prediction loss can further improve the performance.

Revised:  Preliminaries

Revised: 3.1 Problem statement

Revised: We consider the standard RL framework where an agent interacts with its environment in discretetime.

Revised: Formally, we formulate our problem as a Markov decision process (MDP; Sutton & Barto(2018)), which is defined as a tuple ( S , A , p, r, γ, ρ 0 ) .

Revised: Here, S is the state space,

Revised: A is the action space,p ( s ′ | s, a ) is the transition dynamics, r ( s, a ) is the reward function, ρ 0 is the initial state distribution,and γ ∈ [0 , 1] is the discount factor.

Revised: A trajectory is made up of a sequence of states, actions, andrewards: τ = ( s 0 , a 0 , r 0 , s 1 , a 1 , r 1 , . . .

Revised: , s T , a T , r T ) , where T is the (pre-defined) horizon.

Revised: The returnof a trajectory, R ( τ )

Revised: =

Original: {+(cid:80) Tt =
+}
Revised: {+(cid:80) Tt =
+}

Revised: γ t r t , is the sum of discounted future rewards.

Revised: The goal of RL is to finda policy that maximizes the expected return E [ R ( τ )] .

Revised: In our work, we focus on offline RL settingwhere we only have access to some fixed dataset consisting of trajectory rollouts of arbitrary policiesinstead of obtaining data via environment interaction.

Revised: In order to address the problem of generalization, we further consider the distribution of MDPs, wherethe transition dynamics p c ( s ′ | s, a ) varies according to a context c .

Revised: For example, transition dynamicscan vary across different terrains for a walking robot due to different amounts of friction.

Revised: We aim tolearn a single policy π that maximizes the expected return robustly to such dynamics changes.

Revised: To bespecific, given trajectories from training dynamics with contexts sampled from p train ( c ) , we aim atlearning a policy π that can maximize expected return for test dynamics with unseen (but related)contexts sampled from p test ( c ) .

Revised: 3.2 Transformer

Revised: Transformer (Vaswani et al., 2017) is an architecture to efficiently model sequential data by usingstacked self-attention module with residual connection.

Revised: Each self-attention module takes n embeddings { x i } ni =1 and outputs another

Revised: n embeddings { z i } ni =1 with the same dimensions.

Revised: Each token x iis projected to a key k i , query q i , and value v i using linear transformation.

Revised: Then, the output z i iscomputed by weighting values v j by the normalized inner product between the query q i and other

Revised: [Equation]

Revised: s

Revised: [Equation]

Revised: a t linear emb.

Revised: + pos.

Revised: enc.

Revised: 𝐑# keys

Revised: [Equation]

Revised: The self-attention modules work by performing the inner product between the current embedding x iwith the query matrix Q to give q i and with the key matrix K to give the key k i .

Revised: The model takes theinner product between the key and query for j ∈ [1 , n ] , yielding n logits, over which the model takes the softmax and yields a probability distribution.

Revised: The final result is then a convex combination of thevalue vectors v i , with the weights dictated by the probability distribution.

Revised: More concisely, the i -thcomponent of the output z i is given by

Revised: [Equation]

Revised:  DADT: Dynamics-Augmented Decision Transformer

Original: Following Chen et al.
Revised: Following Chen et al.

Original: (2021), we consider the problem of offline reinforcement learning as a sequencemodeling problem; we model the probability of the next sequence token τ i conditioned on all tokensprior to it: P θ ( τ i | τ <i ) .
Revised: (2021), we consider the problem of offline reinforcement learning as a sequencemodeling problem; we model the probability of the next sequence token τ i conditioned on all tokensprior to it: P θ ( τ i | τ <i ) .

Original: The sequences we consider have the form:
Revised: The sequences we consider have the form:

Original: [Equation]
Revised: [Equation]

Original: where T denotes the max timestep, and ˆ R t is the return-to-go , which is the sum of return for the restof the sequence.
Revised: where T denotes the max timestep, and ˆ R t is the return-to-go , which is the sum of return for the restof the sequence.

Original: Architecture.
Revised: Architecture.

Original: we largely follow the architecture from Chen et al.
Revised: we largely follow the architecture from Chen et al.

Original: and then append nextstate prediction head, which will be explained below.
Revised: and then append nextstate prediction head, which will be explained below.

Original: Specifically, our transformer takes the last Ktimesteps as an input; namely, 3 K tokens consisting of return-to-go, state, and action are fed.
Revised: Specifically, our transformer takes the last K timesteps as an input; namely, 3 K tokens consisting of return-to-go, state, and action are fed.

Original: Rawinputs are projected into the embedding dimension by a linear layer for each modality followed bylayer normalization (Ba et al., 2016).
Revised: Raw inputs are projected into the embedding dimension by a linear layer for each modality followedby layer normalization (Ba et al., 2016).

Revised: Additionally, we learn positional embedding per timestepand add it to each token; one timestep corresponds to three tokens.

Original: The tokens are then processed by a causal transformer, followedby action prediction head; future action tokens are predicted via autoregressive modeling along withmasking future tokens in an input sequence.
Revised: The tokens are then processedby a causal transformer, followed by action prediction head; future action tokens are predicted viaautoregressive modeling along with masking future tokens in an input sequence.

Original: Analogously to the action prediction head, we stack alinear decoder for next state prediction on top of the transformer. 
Revised: Analogously to theaction prediction head, we stack a linear decoder for next state prediction on top of the transformer. We provide an illustration of our architecture in Figure 2.

Original: Training with next state prediction.
Revised: Training with next state prediction.

Original: Along with the action prediction proposed by Chen et al.(2021), we train our DADT using next state prediction as an auxiliary training task to boost dynamicsgeneralization.
Revised: Along with the action prediction proposed by Chen et al.(2021), we train our DADT using next state prediction as an auxiliary training task to boost dynamicsgeneralization.

Original: First, given a dataset of offline trajectories from varying dynamics, we sample minibatches of sequence length K from the dataset.
Revised: First, given a dataset of offline trajectories from varying dynamics, we sample minibatches of sequence length K from the dataset.

Original: Note that we do not include any dynamics-id ininput sequences; this can be considered as a more practical setting because annotating such ids wouldrequire additional labeling costs in practice.
Revised: Note that we do not include any dynamics-id ininput sequences; this can be considered as a more practical setting because annotating such ids wouldrequire additional labeling costs in practice.

Original: For the total loss, we add action prediction loss andauxiliary next state prediction loss, multiplied by balancing coefficient λ .
Revised: For the total loss, we add action prediction loss and [Table] auxiliary next state prediction loss, multiplied by balancing coefficient λ .

Original: This next state predictionloss enriches the training signal, which can help capture subtle changes across varying dynamics.
Revised: This next state predictionloss enriches the training signal, which can help capture subtle changes across varying dynamics.

Revised: While our next state prediction loss has a connection with that of Trajectory Transformer (Janneret al., 2021) in that both of them predict next state given past transitions, we remark that DADT doesnot use discretization, which might be a negative factor in dynamics generalization.

Revised: Specifically,approximating a set of values from different dynamics using the exact discrete quantities could hurtcapturing subtle changes in distinguishing the difference between dynamics.

Original: Evaluation.
Revised: Evaluation.

Original: When we encounter unseen dynamics, we evaluate both zero-shot generalization and fewshot adaptation ability of DADT.
Revised: When we encounter unseen dynamics, we evaluate both zero-shot generalization and fewshot adaptation ability of DADT.

Original: For zero-shot, we do not update our DADT for the new dynamicsassuming that trajectories from the new environments are unavailable.
Revised: For zero-shot, we do not update our DADT for the new dynamicsassuming that trajectories from the new environments are unavailable.

Original: For few-shot adaptation, wefine-tune DADT in a supervised manner given a fixed number of samples from the new environments,which is the same as in the training phase. 
Revised: For few-shot adaptation, wefine-tune DADT in a supervised manner given a fixed number of samples from the new environments,which is the same as in the training phase. For deployment, we specify the desired performance byconditioning return-to-go ˆ R 0 along with the starting state s 0 . Then, we generate action for the currentstate (in a deterministic manner). After feeding the generated action, we decrement the target returnby

Revised: the achieved reward and repeat until episode termination.

Original:  Experiments3.1 Setups
Revised:  Experiments

Revised: 5.1 Setups

Original: Environments.
Revised: Environments.

Original: We demonstrate the effectiveness of our model in various simulated robotic controlenvironments (i.e., Walker, Hopper, HalfCheeatah, SlimHumanoid, Ant) based on the OpenAI
Revised: We demonstrate the effectiveness of our model in various simulated robotic controlenvironments (i.e., Walker, Hopper, HalfCheeatah, SlimHumanoid, Ant) based on the OpenAI

Original: Gym (Brockman et al., 2016).
Revised: Gym(Brockman et al., 2016).

Original: In each environment, we consider 20 training dynamics and 5test dynamics with a different set of dynamics parameters (i.e., mass, inertia, damping, friction)following setups of Finn et al.
Revised: In each environment, we consider 20 training dynamics and 5 test dynamicswith a different set of dynamics parameters (mass, inertia, damping, friction) following setups of Finnet al.

Original: (2017) and Rakelly et al.
Revised: (2017) and Rakelly et al.

Original: Note that test dynamics have differentdynamics parameters from training dynamics.
Revised: Note that test dynamics have different dynamics parametersfrom training dynamics.

Original: We provide detailed descriptions of our environmentsin Appendix B.
Revised: We provide detailed descriptions of our environments in Appendix A.

Original: Offline data collection.
Revised: Offline data collection.

Original: We follow the same procedure for offline data collection in Li et al.
Revised: We follow the same procedure for offline data collection in Li et al.

Original: Specifically, we train policies using soft actor-critic (Haarnoja et al., 2018) for every single dynamicsand collect trajectories every 50K timesteps.
Revised: Specifically, we train policies using soft actor-critic (Haarnoja et al., 2018) for every single dynamicsand collect trajectories every 50K timesteps.

Original: Collected trajectories have diverse qualities; trajectoriesfrom early timestep are likely to have lower average return, while later ones have higher return.
Revised: Collected trajectories have diverse qualities; trajectoriesfrom early timestep are likely to have lower average return, while later ones have higher return.

Original: Forcollecting each trajectory, we set the max horizon T as 200 following Li et al.
Revised: Forcollecting each trajectory, we set the max horizon T as 200 following Li et al.

Original: and
Revised: and

Original: Lin et al.(2022)
Revised: Lin et al.(2022)

Original: We provide detailed statistics of collected offline datasets in Appendix C.
Revised: We provide detailed statistics of collected offline datasets in Appendix B.

Original: [Table]

Original: Implementation details.
Revised: Implementation details.

Original: We largely follow the implementation details of the original Decision
Revised: We largely follow the implementation details of the original Decision

Original: Transformer (Chen et al., 2021) and we determined some key distinct hyperparameters for trainingand adaptation.
Revised: Transformer (Chen et al., 2021). Specifically, we train DADT for 1M steps using AdamW optimizer(Loshchilov & Hutter, 2017) with a learning rate of 3 × 10 − 5 containing linear warmup steps of 10K,weight decay of 10 − 4 , gradient clip of 1.0, and batch size of 64. For fine-tuning DADT on unseen [Table]

Revised: dynamics, we train DADT with a learning rate of 10 − 5 , weight decay of 10 − 2 , and batch size of 32.

Original: We provide further implementation details in Appendix D.
Revised: We provide further implementation details in Appendix C.

Original: Evaluation protocol.
Revised: Evaluation protocol.

Original: We train DADT and baselines with 5 independent runs and evaluate themacross 5 (unseen) test dynamics.
Revised: We train DADT and baselines with 5 independent runs and evaluate themacross 5 (unseen) test dynamics.

Original: To assess performance in the test dynamics, we measure returns of10 rollouts whose max timestep is 200.
Revised: To assess performance in the test dynamics, we measure returnsof 10 rollouts whose max timestep is 200.

Original: We provide detailed explanation of evaluating schemes in 
Revised: Then, to facilitate comparison across different environments, we normalize returns for each dynamics to the range between 0 to 100, by computing Normalized _ return = 100 × [Table] following setups in Chen et al.

Original: Appendix D.
Revised: (2021)and Fu et al. For random _ return and expert _ return , we use minimum/maximum returnof collected trajectories, respectively, which are collected periodically while training a soft actor-critic(Haarnoja et al., 2018) for each test dynamics. Finally, we average out 5 (runs) × 5 (test dynamics)× 10 (rollouts) number of returns to measure performance. For DADT and baselines, we measure thetest performance periodically during training and report the best return.

Original: Baselines.
Revised: Baselines.

Original: To evaluate the performance of our method, we consider the following state-of-the-artmethods for offline dynamics generalization:
Revised: To evaluate the performance of our method, we consider the following state-of-the-artmethods for offline dynamics generalization:

Original: • FOCAL (Li et al., 2020): an offline model-free actor-critic meta-RL method with (i) a determin- istic context encoder for efficient task inference and (ii) behavior regularization (Wu et al., 2019)
Revised: • FOCAL (Li et al., 2020): an offline model-free actor-critic meta-RL method with (i) a deterministic context encoder for efficient task inference and (ii) behavior regularization (Wu et al.

Original: to mitigate bootstrapping errors from offline dataset.
Revised: , 2019)to mitigate bootstrapping errors from offline dataset.

Original: We only evaluate FOCAL in a zero-shot setup because FOCAL is designed to be fully offline without adaptation to unseen dynamics.
Revised: We only evaluate FOCAL in a zero-shotsetup because FOCAL is designed to be fully offline without adaptation to unseen dynamics.

Original: In particular, context encoder in FOCAL is trained with distance metric loss (Sohn, 2016) that pushes away samples from other dynamics in the embedding space.
Revised: In particular, context encoder in FOCAL is trained with distance metric loss (Sohn, 2016) thatpushes away samples from other dynamics in the embedding space.

Original: This makes it non-trivial to fine-tune the model using samples from a single unseen dynamics.
Revised: This makes it non-trivialto fine-tune the model using samples from a single unseen dynamics.

Original: • MerPO (Lin et al., 2022): an offline model-based meta-RL method in which proximal meta- learning (Zhou et al., 2019) is used for training both task-specific dynamics model and taskspecific actor-critic networks.
Revised: A comparison with thismethod shows that DADT could capture the context of different dynamics without contrastivelearning.• MerPO (Lin et al., 2022): an offline model-based meta-RL method in which proximal metalearning (Zhou et al., 2019) is used for training both task-specific dynamics model and taskspecific actor-critic networks.

Original: We report evaluation results (i) using meta-policy and metadynamics model without adaptation (zero-shot) and (ii) using task-specific policy and dynamicsmodel after adaptation to unseen dynamics (adaptation).
Revised: We report evaluation results (i) using meta-policy and metadynamics model without adaptation (zero-shot) and (ii) using task-specific policy and dynamicsmodel after adaptation to unseen dynamics (adaptation).

Revised: A comparison with this method showsthat simple supervised learning scheme with transformer-based model could perform comparablyor better than complex meta-learning approaches (Finn et al., 2017; Stadie et al., 2018;

Revised: Rothfusset al., 2018; Mitchell et al., 2021).

Original: 3.2 How does DADT perform compared to prior state-of-the-art methods?
Revised: 5.2 How does DADT perform compared to prior state-of-the-art methods?

Original: We compare DADT to the prior state-of-the-art methods in two scenarios: zero-shot and adaptation.
Revised: We compare DADT to the prior state-of-the-art methods in two scenarios: zero-shot and adaptation.

Original: As shown in Table 1, DADT surpasses the prior arts across all the environments and evaluationscenarios.
Revised: As shown in Table 1, DADT surpasses the prior arts across all the environments and evaluationscenarios.

Original: We also find that DADT without any adaptation outperforms adapted MerPO in 4 out of 5environments, which further highlights the dynamics generalization capability of DADT.
Revised: We also find that DADT without any adaptation outperforms adapted MerPO in 4 out of 5environments, which further highlights the dynamics generalization capability of DADT.

Original: We alsoemphasize that DADT does not use any dynamics annotation in training, while other baselines aretrained with trajectories annotated with dynamics-id.
Revised: We alsoemphasize that DADT does not use any dynamics annotation in training, while other baselines aretrained with trajectories annotated with dynamics-id.

Original: 3.3 Jointly optimizing state and action prediction boosts dynamics generalization
Revised: 5.3 How does DADT perform compared to DT and BC?

Original: Tofurther investigate whether including the proposed auxiliary  state prediction loss is helpful fordynamics generalization, we compare our method and alternative training scheme, which also learnsstate prediction but in a different manner, with respect to next state prediction error in unseendynamics.
Revised: To investigate the effect of the proposed auxiliary next state prediction for improving dynamicsgeneralization, we compare DADT with DT in Table 2 in a zero-shot setup, where the same networkarchitecture and hyperparameters are used for both DADT and DT. We find that DADT outperforms [Table]

Original: In this context, we design DT+, a variant of DT; it is trained using only action prediction
Revised: DT, demonstrating effectiveness of the proposed prediction loss. also compare DTand DADT with Transformer-BC, which does not condition on return-to-go and is trained in a manner of behavior cloning (BC).

Original: [Table]
Revised: We find that both DT and DADT significantly outperform Transformer-BC. These results exhibit that return-to-go conditioned training is more capable of processing offline dataset, which contains suboptimal behaviors, more effectively than BC in dynamics generalization. 5.4 Jointly optimizing state and action prediction boosts dynamics generalization To further investigate whether including the proposed auxiliary state prediction loss is helpful fordynamics generalization, we compare our method and alternative training scheme, which also learnsstate prediction but in a different manner, with respect to next state prediction error in unseendynamics.

Original: loss (as the original DT does),  then the additional linear layer for state prediction is stacked and trained on top of the (frozen) trained transformer
Revised: In this context, we design DT+, a variant of DT; it is trained using only action predictionloss (as the original DT does), and then the additional linear layer for state prediction is stacked andtrained on top of the (frozen) trained transformer.

Original: Namely, a comparison between our DADT and
Revised: Namely, a comparison between our DADT and

Original: DT+ could support how jointly optimizing state prediction loss can help transformer’s capabilities of dynamics generalization.
Revised: DT+ could support how jointly optimizing state prediction loss can help transformer’s capabilities ofdynamics generalization.

Original: As shown in Table  exhibits across the  even if DT+ uses the same amount of  and same 
Revised: As shown in Table 3, DADT exhibits far lower next state prediction erroracross the environments even if DT+ uses the same amount of data and same objective function.

Revised: Thisresult further supports that learning next state prediction in joint with action conditioned on desiredreturn is effective in dynamics generalization.

Original: 3.4 How is DADT compute-efficient?
Revised: 5.5 How is DADT compute-efficient?

Original: [Figure]
Revised: [Figure]

Original: Figure 1 shows the training throughput of DADT and baselines on Ant environment.
Revised: Figure 3 shows the training throughput of DADT and baselines on Ant environment.

Original: The throughput is measuredas the number of updated gradient steps per second ona single GeForce RTX 2080 Ti GPU and an Intel Xeon
Revised: The throughput is measuredas the number of updated gradient steps per second ona single GeForce RTX 2080 Ti GPU and an Intel Xeon

Original: CPU E5-2630 @ 2.80GHz.
Revised: CPU E5-2630 @ 2.80GHz.

Original: We find that DADT requiresabout 6x less time than other baselines.
Revised: We find that DADT requiresabout 6x less time than other baselines.

Original: This is because
Revised: This is because

Original: FOCAL has quadratic time complexity in distance metriclearning, and MerPO trains separated networks for dynamics model and policy in order, while DADT jointly trainsboth return-conditioned policy and dynamics informationin linear time complexity with single transformer.
Revised: FOCAL has quadratic time complexity in distance metriclearning, and MerPO trains separated networks for dynamics model and policy in order, while DADT jointly trainsboth return-conditioned policy and dynamics informationin linear time complexity with single transformer.

Original: 3.5 Importance of transformer architecture in next state prediction
Revised: 5.6 Importance of transformer architecture in next state prediction

Original: [Figure]
Revised: [Figure]

Original: To further investigate the source of capabilities for dynamics generalization of DADT, we compare next stateprediction error of DADT and alternative architectures.
Revised: To further investigate the source of capabilities fordynamics generalization of DADT, we comparenext state prediction error of DADT and alternative architectures.

Original: For the comparison, we consider three types of architectures: Multi-Layer Perceptron (MLP), LSTM (Hochreiter & Schmidhuber, 1997), and Context-aware Dynamics Model (CaDM) (Lee et al., 2020), which extracts contextfrom history of state-action pairs to predict next state.
Revised: For the comparison, we considerthree types of architectures: Multi-Layer Perceptron(MLP), LSTM (Hochreiter & Schmidhuber, 1997),and Context-aware Dynamics Model (CaDM) (Leeet al., 2020), which extracts context from history ofstate-action pairs to predict next state.

Original: Weremark that MLP only consumes the current state without considering past transitions, while LSTM and CaDMwatch a fixed number of past transitions along with recurrent model and ensemble of MLPs, respectively.
Revised: We remarkthat MLP only consumes the current state withoutconsidering past transitions, while LSTM and CaDMwatch a fixed number of past transitions along withrecurrent model and ensemble of MLPs, respectively.

Original: Asshown in Figure 2, we observe that DADT shows a far lowprediction error overall.
Revised: As shown in Figure 2, we observe that DADT shows a far low prediction error overall.

Original: This result demonstrates that the generalization ability of DADT comes fromthe high state prediction ability of transformer architecture. 
Revised: This resultdemonstrates that the generalization ability of DADT comes from the high state prediction ability oftransformer architecture. [Figure]

Revised: 5.7 How does DADT perform in low data-quality regime compared to baselines?

Revised: To investigate how DADT performs in a low data-quality regime, we compare DADT and baselineswith a subset of offline dataset , which is consisted of suboptimal trajectories.

Original: Conclusion
Revised: Specifically, we builda new offline dataset by selecting trajectories of which the average return is in the bottom 30% ofthe original offline dataset; note that the dataset for adaptation is not re-built. As shown in Figure 5, DADT shows the best performance both in zero-shot and adaptation in the low data-quality regime. 5.8 How does balancing coefficient affect performance?

Original: In this work, we proposed Dynamics-Augmented Decision Transformer (DADT), a simple yeteffective transformer-based model for offline dynamics generalization with auxiliary state predictionloss.
Revised: of DADT with varying balancing coeffi- [Table] Finally, we investigate how the the balancing coefficient λ ,which determines how state prediction loss occupies totalloss, affects the performance of DADT.

Original: We have shown that our model outperforms state-of-the-art algorithms in various control tasks. Remarkably, DADT without additional adaptation surpasses fine-tuned baselines in unseen dynamics.
Revised: Figure 4 shows theoverall performance in unseen dynamics by aggregatingthe average return across all environments with different λ . One can observe that the difference in performance by λ isnot significant.

Original: We believe that  DADT will guide new interesting directions in offline dynamics generalization.
Revised: This result indicates that our performanceof DADT is robust to the choice of hyperparameter λ .

Original:  A Related WorkA.1 Dynamics generalization and adaptation
Revised:  Conclusion

Original: Recent works handling dynamics generalization focus on encoding inductive bias using prior knowledge (Zambaldi et al., 2018) or learning contextual latent encoder that captures local dynamics(Lee et al., 2020; Seo et al., 2020). Apart from this, several meta-learning approaches to dynamicsgeneralization have been proposed for quickly adapting to new tasks (Gupta et al., 2018; Rakellyet al., 2019; Zintgraf et al., 2019). Rakelly et al.

Original: (2019) introduced probabilistic context variableaccumulating past experiences for efficient exploration and fast task inference.
Revised: In this work, we proposed Dynamics-Augmented Decision

Original: Recently, Melo (2022)uses transformer architecture for capturing contextual information instead of recurrent neural network Duan et al. (2016) in online dynamics generalization.
Revised: Transformer (DADT), a simple yet effective transformer-based model for offline dynamics generalization with auxiliary state prediction loss.

Original: It differs from our work in that(i) we usetransformer architecture not only for capturing contextual information but also for optimizing policyconditioned on desired return, and (ii) our model utilizes offline datasets. Along with increasing attention toward offline RL (Levine et al., 2020) of which importance ishighlighted especially when environmental interaction is expensive or risky, dynamics generalizationand adaptation in the offline setting have been studied recently (Li et al., 2020; Lin et al., 2022; Ball et al., 2021; Cang et al., 2021).
Revised: We have shown that our model outperforms state-of-the-artalgorithms in various control tasks.

Original: They aim to learn a policy or dynamics model that can quicklyadapt to unseen dynamics along with mitigating value overestimation issues inoffline RL (Fujimotoet al., 2019). Specifically, FOCAL (Li et al., 2020) trains a deterministic context encoder for efficientdynamics inference using contrastive loss and use it for behavior regularized policy learning (Wuet al., 2019). MerPO (Lin et al., 2022) takes a gradient-based meta-learning approach for learningmeta-dynamics model and meta-policy that aim to adapt to unseen dynamics quickly; the metadynamics model is used to generate synthetic rollouts to boost policy learning. Ball et al.
Revised: Remarkably, DADT without additional adaptation surpassesfine-tuned baselines in unseen dynamics.

Original: (2021)augment a learned dynamics model to improve the zero-shot generalization in model-based offline 
Revised: We believe that DADT will guide new interesting directionsin offline dynamics generalization.

Original: RL.
Revised: Limitation.

Original: Motivated by remarkable success in [-large-scale language modeling (Devlin et al., 
 ; Brown et al.-]
Revised: While our work shows impressive results in [-large-scale language modeling (Devlin et al., 
 ; Brown et al.-]offline dynamics generalization with variouscontinuous control tasks, we have considered the limited context of varying dynamics (i.e., mass,friction, and inertia).

Original: ,2020), it has been proposed to extend such success into RL domain by formulating RL as a sequencemodeling problem (Chen et al., 2021; Janner et al., 2021). These works build on the reinforcementlearning as supervised learning paradigm (Schmidhuber, 2019; Srivastava et al., 2019; Emmons et al. ,2021) that focuses on predictive modeling of sequence of states, actions, and reward-based returns.
Revised: However, in the real world, agents often encounter a more diverse and difficultsituation of dynamics generalization rather than varying dynamics coefficient, i.e., navigating variousshapes of mazes.

Original: Specifically, Chen et al. (2021) proposed Decision Transformer (DT), which is trained to generatethe optimal actions given desired return and history of states and actions. Instead of conditioningdesired return, Trajectory Transformer (TT) (Janner et al., 2021) is trained to predict not only actions,but also entire elements of trajectories such as states and rewards and finds optimal actions by usingplanning (e.g., beam search) in deployment time. Recently, Lee et al.
Revised: Therefore, it would be an interesting future direction to extend our work into amore complex setting with more practical experimental setups like autonomous driving (i.e., CARLA(Dosovitskiy et al., 2017)).

Original: (2022) have demonstrated that DT can play multiple Atari games with a single set of weights and can be efficiently fine-tuned forsolving unseen games. Though DT and TT have shown promising results in solving various control tasks, their ability togeneralize to unseen dynamics has not investigated yet. In this paper, we demonstrate that DT trainedon aggregated trajectories collected from different dynamics can generalize to unseen dynamics, andintroducing an auxiliary next state prediction loss can further improve the performance.
Revised: Extending our work into setups where reward function also changes alongwith varying dynamics would be an interesting future step.

Original: B Environment details
Revised: A Environment details

Original: We use 5 simulated robotic control environments from OpenAI Gym (Brockman et al., 2016).
Revised: We use 5 simulated robotic control environments from OpenAI Gym (Brockman et al., 2016).

Original: Welargely follow the implementation details of the Gym 1 for resetting agent state and shaping rewardfunction of Walker, Hopper, HalfCheetah, and Ant.
Revised: Welargely follow the implementation details of the Gym 1 for resetting agent state and shaping rewardfunction of Walker, Hopper, HalfCheetah, and Ant.

Original: For SlimHumanoid, we change the environmentcode of Humanoid in Gym, following the description in (Wang et al., 2019).
Revised: For SlimHumanoid, we change the environmentcode of Humanoid in Gym, following the description in (Wang et al., 2019).

Original: Walker.
Revised: Walker.

Original: Walker 
Revised: Walker is

Original: is a two-dimensional two-legged figure with 7 rigid links 
Revised: a two -dimensional two-legged figure with 7 rigid links including a torso

Original: , including a torso and
Revised: and

Original: 2 legs.
Revised: 2 legs.

Original: The goal is to move forward as fast as possible while maintaining the standing height and
Revised: The goal is to move forward as fast as possible while maintaining the standing height andconsuming as small control input as possible

Original: consuming

Original: as small control input as possible

Original: • Observation.
Revised: • Observation.

Original: Observation is a 17-dimensional vector that includes 1) root joint’s position, 2)angular position and velocities.• Action.
Revised: Observation is a 17-dimensional vector that includes 1) root joint’s position, 2)angular position and velocities. • Action.

Original: a ∈ [ − 1 . 0 , 1 . 0] 6 represents torques applied at six joints connecting the six body parts.• Reward.
Revised: a ∈ [ − 1 . 0 , 1 . 0] 6 represents torques applied at six joints connecting the six body parts. • Reward.

Original: r t = ˙ x torso,t − 0 .
Revised: r t = ˙ x torso,t − 0 .

Original: 001 ∥ a t ∥ 2 where ˙ x torso,t represents forward velocity of the torso.
Revised: 001 ∥ a t ∥ 2 where ˙ x torso,t represents forward velocity of the torso.

Original: Wealso add an alive bonus of 1 to the agents at every time-step.
Revised: We also add an alive bonus of 1 to the agents at every time-step.

Original: Hopper.
Revised: Hopper.

Original: Hopper is a two-dimensional one-legged figure consisting of 4 rigid links, including a torso,thigh, leg, and foot.
Revised: Hopper is a two-dimensional one-legged figure consisting of 4 rigid links, including a torso,thigh, leg, and foot.

Original: The goal is to move forward as fast as possible while maintaining the standingheight and consuming as small control input as possible.• Observation.
Revised: The goal is to move forward as fast as possible while maintaining the standingheight and consuming as small control input as possible. • Observation.

Original: Observation is a 11-dimensional vector that includes the angular position andvelocity of all the joints, except for the x position of the root joint.• Action.
Revised: Observation is a 11-dimensional vector that includes the angular position andvelocity of all the joints, except for the x position of the root joint.• Action.

Original: a ∈ [ − 1 . 0 , 1 . 0] 3 represents torques applied at three joints connecting the four body parts.• Reward.
Revised: a ∈ [ − 1 . 0 , 1 . 0] 3 represents torques applied at three joints connecting the four body parts.• Reward.

Original: r t = ˙ x torso,t − 0 .
Revised: r t = ˙ x torso,t − 0 .

Original: 001 ∥ a t ∥ 2 , where ˙ x torso,t represents forward velocity of the torso.
Revised: 001 ∥ a t ∥ 2 , where ˙ x torso,t represents forward velocity of the torso.

Original: Wealso add an alive bonus of 1 to the agents at every time-step.
Revised: Wealso add an alive bonus of 1 to the agents at every time-step.

Original: HalfCheetah.
Revised: HalfCheetah.

Original: HalfCheetah is a two-dimensional two-legged figure.
Revised: HalfCheetah is a two-dimensional two-legged figure.

Original: It consists of7 rigid links (1for torso, 3 for forlimb„ 3  for hindlimb).
Revised: It consists of rigid links (  for torso

Original: It is connected by 6 joints
Revised: 3 for forlimb„ 3 for

Original:  to which an actuator is attachedper each joint.
Revised: hindlimb). It is connected by 6 joints, to which an actuator is attached per each joint.

Original: The goal is to move forward as quickly as possible while minimizing the cost forcontrolling agent as small as possible.
Revised: The goal is to move forward as quickly as possible while minimizing the cost for controlling agent as small as possible.

Original: • Observation.
Revised: • Observation.

Original: Observation is given by a 17-dimensional vector that includes 1) root joint’s position(except for the x-coordinate) and velocity, and 2) center of mass of the torso.
Revised: Observation is given by a 17-dimensional vector that includes 1) root joint’s position(except for the x-coordinate) and velocity, and 2) center of mass of the torso.

Original: • Action.
Revised: • Action.

Original: a ∈ [ − 1 . 0 , 1 . 0] 6 represents torques applied at six joints.
Revised: a ∈ [ − 1 . 0 , 1 . 0] 6 represents torques applied at six joints.

Original: • Reward.
Revised: • Reward.

Original: r t = ˙ x torso,t − 0 .
Revised: r t = ˙ x torso,t − 0 .

Original: 1 ∥ a t ∥ 2 , where ˙ x torso,t represents forward velocity of the torso.
Revised: 1 ∥ a t ∥ 2 , where ˙ x torso,t represents forward velocity of the torso.

Original: SlimHumanoid.
Revised: SlimHumanoid.

Original: Humanoid is a 3D robot for simulating a human.
Revised: Humanoid is a 3D robot for simulating a human.

Original: It has a pair of legs and armsunlike other environments, which is consisted of 13 rigid links, 17 actuators.
Revised: It has a pair of legs and armsunlike other environments, which is consisted of 13 rigid links, 17 actuators.

Original: The goal is to moveforward as quickly as possible while minimizing the cost for controlling agent as small as possible.
Revised: The goal is to moveforward as quickly as possible while minimizing the cost for controlling agent as small as possible.

Original: SlimHumanoid Wang et al.
Revised: SlimHumanoid Wang et al.

Original: (2019) is a variant of Humanoid with limited observation and ease ofuse.
Revised: (2019) is a variant of Humanoid with limited observation and ease ofuse.

Original: Specifically, there are some changes in observation details (excluding the center of mass-basedquantities, external force, and actuator force) and reward function (no penalizing on external force).
Revised: Specifically, there are some changes in observation details (excluding the center of mass-basedquantities, external force, and actuator force) and reward function (no penalizing on external force).

Original: • Observation.
Revised: • Observation.

Original: Observation is a 45-dimensional vector that includes angular position and velocities.
Revised: Observation is a 45-dimensional vector that includes angular position and velocities.

Original: • Action.
Revised: • Action.

Original: a ∈ [ − 0 .
Revised: a ∈ [ − 0 .

Original: 4 , 0 . 4] 17 represents torques applied at seventeen joints.
Revised: 4 , 0 . 4] 17 represents torques applied at seventeen joints.

Original: • Reward.
Revised: • Reward.

Original: 5 / 4 × ˙ x t − 0 .
Revised: 5 / 4 × ˙ x t − 0 .

Original: 1 ∥ a t ∥ 2 2 + 5 × bool (1 . 

Revised: 1 ∥ a t ∥ 2 2 + 5 × bool (1 .

Original: < =
Revised: < =

Original: z t < = 2 . 0) , where ˙ x torso,t represents forward
Revised: z t < = 2 . 0) , where ˙ x torso,t represents forward

Original: velocity of the torso, and z t is the height of the torso.
Revised: velocity of the torso, and z t is the height of the torso.

Original: We also add an alive bonus of 5 to theagents at every time-step.
Revised: We also add an alive bonus of 5 to the agents at every time-step.

Original: Ant.
Revised: Ant.

Original: Ant is a 3D ant-like robot consisting of a torso with 4 legs, each of which is connected with2 joints.
Revised: Ant is a 3D ant-like robot consisting of a torso with 4 legs, each of which is connected with [Table] 2 joints.

Original: There are 8 actuators consisting of 13 rigid links.
Revised: There are 8 actuators consisting of 13 rigid links.

Original: The goal is to move forward as fast as
Revised: The goal is to move forward as fast as

Original: [Table]
Revised: [Table]

Original: possible while maintaining the standing height and consuming as small control input as possible.
Revised: possible while maintaining the standing height and consuming as small control input as possible.

Original: [Table]
Revised: [Table]

Original: • Observation.
Revised: • Observation.

Original: Observation is a 27-dimensional vector that includes angular position and velociyuof all 8 joints, except for the x and y positions of the root joints.
Revised: Observation is a 27-dimensional vector that includes angular position and velociyuof all 8 joints, except for the x and y positions of the root joints.

Original: [Table]
Revised: [Table]

Original: • Action.
Revised: • Action.

Original: a ∈ [ − 1 . 0 , 1 . 0] 8 represents torques applied at eight joints connecting the two links ofeach leg and torso.
Revised: a ∈ [ − 1 . 0 , 1 . 0] 8 represents torques applied at eight joints connecting the two links of

Original: [Table]
Revised: [Table]

Original: • Reward.
Revised: each leg and torso.• Reward.

Original: r t = ˙ x torso,t − 0 .
Revised: r t = ˙ x torso,t − 0 .

Original: 5 ∥ a t ∥ 22 − 0 .
Revised: 5 ∥ a t ∥ 22 − 0 .

Original: 5 × 0 .
Revised: 5 × 0 .

Original: 001 × + 1.0, where ˙ x torso,t represents forward [Table] velocity of the torso.
Revised: 001 × + 1.0, where ˙ x torso,t represents forwardvelocity of the torso.

Original: We also add an alive bonus of 1 to the agents at every time-step. 
Revised: We also add an alive bonus of 1 to the agents at every time-step. B Offline dataset details

Original: [Table]
Revised: [Table]

Original: C Offline dataset details

Original: Dataset Collection.
Revised: Dataset Collection.

Original: For each environment, we collect 20 train dynamics and 5 test dynamics bychanging the parameter of the agent, specifically body mass, body inertia, damping, and friction.
Revised: For each environment, we collect 20 train dynamics and 5 test dynamics bychanging the parameter of the agent, specifically body mass, body inertia, damping, and friction.

Original: Parameters of each dynamics are randomly chosen by multiplying coefficient (1) 1 . 5 r , r ∼ Uni ( − 3 , 3)for mass, inertia, friction and (2) 1 . 3 r , r ∼ Uni ( − 3 , 3) for damping with initial values.
Revised: Parameters of each dynamics are randomly chosen by multiplying coefficient (1) 1 . 5 r , r ∼ Uni ( − 3 , 3)for mass, inertia, friction and (2) 1 . 3 r , r ∼ Uni ( − 3 , 3) for damping with initial values.

Original: We train policies using SAC Haarnoja et al.
Revised: We train policies using SAC Haarnoja et al.

Original: (2018) for every single dynamics.
Revised: (2018) for every single dynamics.

Original: We skipped the first50K steps for warm-up.
Revised: We skipped the first50K steps for warm-up.

Original: For collecting datasets with diverse qualities, we collect 50 random rolloutsfrom trained policies every 50K timesteps.
Revised: For collecting datasets with diverse qualities, we collect 50 random rolloutsfrom trained policies every 50K timesteps.

Original: The total dataset size of each environment is from 1.3GB(Hopper) to 1.7GB (Ant).
Revised: The total dataset size of each environment is from 1.3GB(Hopper) to 1.7GB (Ant).

Original: You can download the dataset we used for experiments in the link below 2 .
Revised: You can download the dataset we used for experiments in the link below 2 .

Original: Normalization.
Revised: Normalization.

Original: Return statistics of the collected offline trajectories can be found in Table 3.
Revised: Return statistics of the collected offline trajectories can be found in Table 5.

Original: Weused the minimum and maximum value among test dynamics of each environment for normalizingthe average return in evaluation.
Revised: Weused the minimum and maximum value among test dynamics of each environment for normalizingthe average return in evaluation.

Original: [Table]
Revised: [Table]

Original: HalfCheetah
Revised: HalfCheetah

Original: Dynamics Dynamics
Revised: Dynamics Dynamics

Original: Dynamics Dynamics
Revised: Dynamics Dynamics

Original: Dynamicsmin 149.
Revised: Dynamicsmin 149.

Original: max 1681.
Revised: max 1681.

Original: mean 1222.
Revised: mean 1222.

Original: SlimHumanoid
Revised: SlimHumanoid

Original: Dynamics Dynamics
Revised: Dynamics Dynamics

Original: Dynamics Dynamics
Revised: Dynamics Dynamics

Original: Dynamicsmin 148.
Revised: Dynamicsmin 148.

Original: max 975.
Revised: max 975.

Original: mean 444.
Revised: mean 444.

Original: Ant
Revised: Ant

Original: Dynamics Dynamics
Revised: Dynamics Dynamics

Original: Dynamics Dynamics
Revised: Dynamics Dynamics

Original: Dynamicsmin -61.
Revised: Dynamicsmin -61.

Original: max 786.
Revised: max 786.

Original: mean 348.
Revised: mean 348.

Original: D Implementation details
Revised: C Implementation details

Original: In this section, we will explain the implementation details and resources used for training andevaluating DADT.
Revised: In this section, we will explain the implementation details and resources used for training andevaluating DADT.

Original: Specifically, we train DADT for 1M steps using AdamW optimizer (Loshchilov &
Revised: Specifically, we train DADT for 1M steps using AdamW optimizer (Loshchilov &

Original: Hutter, 2017) with a learning rate of 3 × 10 − 5 containing linear warmup steps of 10K, weight decayof 10 − 4 , gradient clip of 1.0, and batch size of 64.
Revised: Hutter, 2017) with a learning rate of 3 × 10 − 5 containing linear warmup steps of 10K, weight decayof 10 − 4 , gradient clip of 1.0, and batch size of 64.

Original: For fine-tuning DADT on unseen dynamics, wetrain DADT with a learning rate of 10 − 5 , weight decay of 10 − 2 , and batch size of 32.
Revised: For fine-tuning DADT on unseen dynamics, wetrain DADT with a learning rate of 10 − 5 , weight decay of 10 − 2 , and batch size of 32.

Original: We include thecode of DADT in the supplementary material.
Revised: We include thecode of DADT in the supplementary material.

Original: Evaluation Protocol To facilitate comparison across different environments, we normalize returns for each dynamics to the range between 0 to 100, by computing Normalized _ return =100 × return − random _ return expert _ return − random _ return following setups in Chen et al.
Revised: Evaluation Protocol To facilitate comparison across different environments, we normalize returns for each dynamics to the range between 0 to 100, by computing Normalized _ return =100 × return − random _ return expert _ return − random _ return following setups in Chen et al.

Original: (2021) and Fu et al.
Revised: (2021) and Fu et al.

Original: Forrandom _ return and expert _ return , we use minimum/maximum return of collected trajectories,respectively, which are collected periodically while training a soft actor-critic (Haarnoja et al.
Revised: Forrandom _ return and expert _ return , we use minimum/maximum return of collected trajectories,respectively, which are collected periodically while training a soft actor-critic (Haarnoja et al.

Original: , 2018)for each test dynamics.
Revised: , 2018)for each test dynamics.

Original: Finally, we average out 5 (runs) × 5 (test dynamics) × 10 (rollouts) numberof returns to measure performance.
Revised: Finally, we average out 5 (runs) × 5 (test dynamics) × 10 (rollouts) numberof returns to measure performance.

Original: For DADT and baselines, we measure the test performanceperiodically during training and report the best return.
Revised: For DADT and baselines, we measure the test performanceperiodically during training and report the best return.

Original: Hyperparamters.
Revised: Hyperparamters.

Original: Hyperparameters for DADT are shown in Table 4.
Revised: Hyperparameters for DADT are shown in Table 6.

Original: We choose target return-to-go,referring to expert performance for each environment, which is set to a value equivalent to or slightlyabove expert performance.
Revised: We choose target return-to-go,referring to expert performance for each environment, which is set to a value equivalent to or slightlyabove expert performance.

Original: We heuristically find that training (1) larger DADT and (2)
Revised: We heuristically find that training (1) larger DADT and (2)

Original: DADT withmore gradient steps boost the model’s performance but use the current version for faster training andevaluation.
Revised: DADT withmore gradient steps boost the model’s performance but use the current version for faster training andevaluation.

Original: [Table]
Revised: [Table]

Original: Loshchilov & Hutter, 2017)
Revised: AdamW (Loshchilov & Hutter, 2017)

Original: [Table]
Revised: [Table]

Revised: β

Original: Resources.
Revised: [Equation] = 0 . 9 , β [Equation]

Revised: =

Original:  For training and evaluating our model, we use a single NVIDIA GeForce RTX 2080 Ti
Revised: [Table] [Table] Resources. For training and evaluating our model, we use a single NVIDIA GeForce RTX 2080 Ti

Original: GPU and 40 CPU cores (Intel(R) Xeon(R) CPU E5-2630 @ 2.80GHz), taking at most 8 hours fortraining and less than 10 seconds for rollout 10 trajectories per evaluation.
Revised: GPU and 40 CPU cores (Intel(R) Xeon(R) CPU E5-2630 @ 2.80GHz), taking at most 8 hours fortraining and less than 10 seconds for rollout 10 trajectories per evaluation.

Original: Compared to the original
Revised: Compared to the original

Original: DT (Chen et al., 2021), DADT shows no significant increase in training time.E
Revised: DT (Chen et al., 2021), DADT shows no significant increase in training time.

Original: [Figure] Low data-quality regime. To investigate how DADT performs in a low data-quality regime, wecompare DADT and baselines with a subset of offline dataset , which is consisted of suboptimaltrajectories. Specifically, we build a new offline dataset by selecting trajectories of which the averagereturn is in the bottom 30% of the original offline dataset; note that the dataset for adaptation is notre-built.
Revised: D Comparison with Trajectory Transformer

Original: As shown in Figure 3, DADT shows the best performance bothin zero-shot and adaptation inthe low data-quality regime. We also provide the detailed result of performances in low data-qualityregime (Figure 3) in Table 5. DADT shows a consistently better result than baselines except in HalfCheetah with zero-shot setting.
Revised: While our next state prediction loss has a connection with that of Trajectory Transformer (Janneret al., 2021) in that both of them predict next state given past transitions, we remark that DADT doesnot use discretization, which might be a negative factor in dynamics generalization.

Original: In the adaptation setting, DADT outperforms baseline models inall environments. This result verifies the superiority of DADT to baselines in dynamics generalizationregardless of the quality of offline trajectories [Table] [Table]
Revised: Specifically,approximating a set of values from different dynamics using the exact discrete quantities could hurt

Original: [Table]
Revised: [Table]

Original: How does balancing coefficient affect performance? Finally, weinvestigate how the the balancing coefficient λ , which determineshow state prediction loss occupies total loss, affects the performanceof DADT. Figure 6 shows the overall performance in unseen dynamics by aggregating the average return across all environments withdifferent λ . One can observe that the difference in performance byλ is not significant.
Revised: capturing subtle changes in distinguishing the difference between dynamics.

Original: This result indicates that our performance of DADT is robust to the choice of hyperparameter λ . F Comparison with Trajectory Transformer While our next state prediction loss has a connection with that of Trajectory Transformer (Janneret al., 2021) in that both of them predict next state given past transitions, we remark that DADT doesnot use discretization, which might be a negative factor in dynamics generalization.
Revised: [Table]

Original: Specifically,approximating a set of values from different dynamics using the exact discrete quantities could hurtcapturing subtle changes in distinguishing the difference between dynamics.

Original: As table 7 demonstrates
Revised: As table 7 demonstrates

Original: [Table]
Revised: [Table]

Original: our DADT consistently outperforms TT.
Revised: our DADT consistently outperforms TT.

Original: We think that the underperformingof TT is caused by (1) reduced visible context due to per-dimension processing and (2) informationloss from discretization.
Revised: We think that the underperformingof TT is caused by (1) reduced visible context due to per-dimension processing and (2) informationloss from discretization.

Original: [Table]
Revised: [Table]