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

Original: Anonymous Author(s)

Original: 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 often possible to train strong agents without potentially unsafe or impractical online interaction.

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: This work presents Dynamics-Augmented Decision Transformer (DADT), a simple yet efficient method to train generalizable agents from offline datasets; on top of return-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 without fine-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 logged dataset 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 be expensive 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 the ability 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 to generalize 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 (Lin et 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 inference on 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 adapt to 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 context latent variable from near past experiences could capture subtle changes in differing dynamics is questionable, 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) have exhibited 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 model using transformer architecture can solve various types of unseen tasks such as translation, question answering, 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 from unseen 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 generalist agent that can solve a range of control tasks using inputs from different modalities with a single set of weights.

Original: [Figure] [Figure]
Revised: [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 (Ramesh et 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: D ynamics-
Revised: Devlin et al., 2018), we present DADT: D ynamics-

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

Original: 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: Analogously to pre-training scheme in transformer-based models in vision and language domains, we train DADT via a supervised learning paradigm with totally aggregated sequences (i.e., trajectories) without distinction.

Original: 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: Specifically, DADT aims to optimize the objective to find the action that achieves desired return (Chen et al., 2021) while minimizing error in next state prediction.

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

Original: 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: For evaluation, we consider two scenarios depending on different accessibility to trajectories from unseen dynamics: (1) zero-shot, where demonstrations are not given, and (2) adaptation, where model is fine-tuned with small amount of samples.

Original: We provide an overview of setup and architecture in Figure 1 and Figure 2, respectively.
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 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).
Revised: We show that DADT outperforms prior state-of-the-art baselines for offline dynamics generalization in both performance and time efficiency on various continuous control tasks based on MuJoCo simulator (Todorov et al., 2012), which is widely used in the literature (Li et al., 2020; Lin et al., 2022).

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

Original: Moreover,reduced next state prediction error in unseen dynamics further supports our model’s capability fordynamics generalization.
Revised: Moreover, reduced next state prediction error in unseen dynamics further supports our model’s capability for dynamics generalization.

Original:  Related Work 
Revised:  Related Work 2.1 Dynamics generalization and adaptation

Original: 2.1 Dynamics generalization and adaptation

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

Original: 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: Apart from this, several meta-learning approaches to dynamics generalization have been proposed for quickly adapting to new tasks (Gupta et al., 2018; Rakelly et al., 2019; Zintgraf et al., 2019).

Original: Rakelly et al.
Revised: Rakelly et al.

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

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

Original: Duan et al.
Revised: Duan et al.

Original: (2016) in online dynamics generalization.
Revised: (2016) in online dynamics generalization.

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.
Revised: It differs from our work in that (i) we use transformer architecture not only for capturing contextual information but also for optimizing policy conditioned on desired return, and (ii) our model utilizes offline datasets.

Original: 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: Along with increasing attention toward offline RL (Levine et al., 2020) of which importance is highlighted especially when environmental interaction is expensive or risky, dynamics generalization and adaptation in the offline setting have been studied recently (Li et al., 2020;

Original: Lin et al., 2022;
Revised: Lin et al., 2022;

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

Original: 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: 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 (Fujimoto et al., 2019).

Original: 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: Specifically, FOCAL (Li et al., 2020) trains a deterministic context encoder for efficient dynamics inference using contrastive loss and use it for behavior regularized policy learning (Wu et al., 2019).

Original: 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: MerPO (Lin et al., 2022) takes a gradient-based meta-learning approach for learning meta-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.

Original: Ball et al.
Revised: Ball et al.

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

Original: RL.
Revised: RL.

Original: 2.2 RL via sequence modelling
Revised: 2.2 RL via sequence modelling

Original: Motivated by remarkable success in large-scale language modeling (Devlin et al., [-
 ;-] Brown et al. ,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: Motivated by remarkable success in large-scale language modeling (Devlin et al., [-
 ;-]2018; Brown et al., 2020), it has been proposed to extend such success into RL domain by formulating RL as a sequence modeling problem (Chen et al., 2021; Janner et al., 2021).

Original: 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: These works build on the reinforcement learning 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.

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

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

Original: 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: Instead of conditioning desired 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 using planning (e.g., beam search) in deployment time.

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

Original: (2022) have demonstrated that
Revised: (2022) have demonstrated that

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

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

Original: 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: In this paper, we demonstrate that DT trained on aggregated trajectories collected from different dynamics can generalize to unseen dynamics, and introducing an auxiliary next state prediction loss can further improve the performance.

Original:  Preliminaries 
Revised:  Preliminaries 3.1 Problem statement

Original: 3.1 Problem statement

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

Original: 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: 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 ) .

Original: Here, S is the state space,
Revised: Here, S is the state space,

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

Original: 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: A trajectory is made up of a sequence of states, actions, and rewards: τ = ( s 0 , a 0 , r 0 , s 1 , a 1 , r 1 , . . .

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

Original: The returnof a trajectory, R ( τ ) =
Revised: The return of a trajectory, R ( τ ) =

Original: (cid:80) Tt =

Revised: (cid:80) Tt =

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

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

Original: 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 our work, we focus on offline RL setting where we only have access to some fixed dataset consisting of trajectory rollouts of arbitrary policies instead of obtaining data via environment interaction.

Original: 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: In order to address the problem of generalization, we further consider the distribution of MDPs, where the transition dynamics p c ( s ′ | s, a ) varies according to a context c .

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

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

Original: 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: To be specific, given trajectories from training dynamics with contexts sampled from p train ( c ) , we aim at learning a policy π that can maximize expected return for test dynamics with unseen (but related) contexts sampled from p test ( c ) .

Original: 3.2 Transformer
Revised: 3.2 Transformer

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

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

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

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

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

Original: [Equation]
Revised: [Equation]

Original: s
Revised: s

Revised: [Equation]

Revised: a

Original: [Equation]
Revised: [Equation]

Original: a t linear emb.
Revised: linear emb.

Original: + pos.
Revised: + pos.

Original: enc.
Revised: enc.

Original: 𝐑# keys 
Revised: 𝐑# keys k

Original: [Equation]
Revised: [Equation]

Original: 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 self-attention modules work by performing the inner product between the current embedding x i with the query matrix Q to give q i and with the key matrix K to give the key k i .

Original: 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 model takes the inner 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.

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

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

Original: [Equation]
Revised: [Equation]

Original:  DADT: Dynamics-Augmented Decision Transformer
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 sequence modeling problem; we model the probability of the next sequence token τ i conditioned on all tokens prior 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 rest of 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 next state prediction head, which will be explained below.

Original: Specifically, our transformer takes the last
Revised: Specifically, our transformer takes the last

Original: K timesteps as an input; namely, 3 K tokens consisting of return-to-go, state, and action are fed.
Revised: K timesteps as an input; namely, 3 K tokens consisting of return-to-go, state, and action are fed.

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

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

Original: 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.
Revised: The tokens are then processed by a causal transformer, followed by action prediction head; future action tokens are predicted via autoregressive modeling along with masking future tokens in an input sequence.

Original: Analogously to theaction prediction head, we stack a linear decoder for next state prediction on top of the transformer.
Revised: Analogously to the action prediction head, we stack a linear decoder for next state prediction on top of the transformer.

Original: We provide an illustration of our architecture in Figure 2.
Revised: 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 dynamics generalization.

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 in input sequences; this can be considered as a more practical setting because annotating such ids would require additional labeling costs in practice.

Original: For the total loss, we add action prediction loss and
Revised: For the total loss, we add action prediction loss and

Original: [Table]
Revised: [Table]

Original: auxiliary next state prediction loss, multiplied by balancing coefficient λ .
Revised: 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 prediction loss enriches the training signal, which can help capture subtle changes across varying dynamics.

Original: 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: While our next state prediction loss has a connection with that of Trajectory Transformer (Janner et al., 2021) in that both of them predict next state given past transitions, we remark that DADT does not use discretization, which might be a negative factor in dynamics generalization.

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.
Revised: Specifically, approximating a set of values from different dynamics using the exact discrete quantities could hurt capturing 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 dynamics assuming 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, we fine-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.

Original: For deployment, we specify the desired performance byconditioning return-to-go ˆ R 0 along with the starting state s 0 .
Revised: For deployment, we specify the desired performance by conditioning return-to-go ˆ R 0 along with the starting state s 0 .

Original: Then, we generate action for the currentstate (in a deterministic manner).
Revised: Then, we generate action for the current state (in a deterministic manner).

Original: After feeding the generated action, we decrement the target returnby the achieved reward and repeat until episode termination.
Revised: After feeding the generated action, we decrement the target return by the achieved reward and repeat until episode termination.

Original:  Experiments 
Revised:  Experiments 5.1 Setups

Original: 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 Gym(Brockman et al., 2016).
Revised: We demonstrate the effectiveness of our model in various simulated robotic control environments (i.e., Walker, Hopper, HalfCheeatah, SlimHumanoid, Ant) based on the OpenAI Gym (Brockman et al., 2016).

Original: 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.
Revised: In each environment, we consider 20 training dynamics and 5 test dynamics with a different set of dynamics parameters (mass, inertia, damping, friction) following setups of Finn et al.

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

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

Original: We provide detailed descriptions of our environments in Appendix A.
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 dynamics and 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; trajectories from 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: For collecting each trajectory, we set the max horizon T as 200 following Li et al.

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

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

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).
Revised: Transformer (Chen et al., 2021).

Original: 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.
Revised: 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.

Original: For fine-tuning DADT on unseen
Revised: For fine-tuning DADT on unseen

Original: [Table]
Revised: [Table]

Original: dynamics, we train DADT with a learning rate of 10 − 5 , weight decay of 10 − 2 , and batch size of 32.
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 C.
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 them across 5 (unseen) test dynamics.

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

Original: Then, to facilitate comparison across different environments, we normalize returns for each dynamics to the range between 0 to 100, by computing
Revised: Then, to facilitate comparison across different environments, we normalize returns for each dynamics to the range between 0 to 100, by computing

Original: Normalized _ return = 100 ×
Revised: Normalized _ return = 100 ×

Original: [Table]
Revised: [Table]

Original: following setups in Chenet al.
Revised: following setups in Chen

Revised: [Table]

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

Original: 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.
Revised: For random _ 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., 2018) for each test dynamics.

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

Original: For DADT and baselines, we measure thetest performance periodically during training and report the best return.
Revised: For DADT and baselines, we measure the test 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-art methods for offline dynamics generalization:

Original: • 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. , 2019)to mitigate bootstrapping errors from offline dataset.
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., 2019) to mitigate bootstrapping errors from offline dataset.

Original: We only evaluate FOCAL in a zero-shotsetup because FOCAL is designed to be fully offline without adaptation to unseen dynamics.
Revised: We only evaluate FOCAL in a zero-shot setup 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) thatpushes away samples from other dynamics in the embedding space.
Revised: 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.

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

Original: 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.
Revised: A comparison with this method shows that DADT could capture the context of different dynamics without contrastive learning. • 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 dynamics model after adaptation to unseen dynamics (adaptation).

Original: 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: A comparison with this method shows that simple supervised learning scheme with transformer-based model could perform comparably or better than complex meta-learning approaches (Finn et al., 2017; Stadie et al., 2018;

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

Original: 5.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 Table1, DADT surpasses the prior arts across all the environments and evaluationscenarios.
Revised: As shown in Table DADT surpasses the prior arts across all the environments and evaluation scenarios.

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 ofenvironments, 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 also emphasize that DADT does not use any dynamics annotation in training, while other baselines are trained with trajectories annotated with dynamics-id.

Original: 5.3 How does DADT perform compared to DT and BC?
Revised: 3 How does DADT perform compared to DT and BC?

Original: 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.
Revised: To investigate the effect of the proposed auxiliary next state prediction for improving dynamics generalization, we compare DADT with DT in Table 2 in a zero-shot setup, where the same network architecture and hyperparameters are used for both DADT and DT.

Original: We find that DADT outperforms
Revised: We find that DADT outperforms

Original: [Table]
Revised: [Table]

Original: DT, demonstrating  effectiveness of the proposed prediction loss.
Revised: DT, demonstrating the effectiveness of the proposed prediction loss.

Original:  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).
Revised: Moreover, we also compare DT and DADT with Transformer-BC, which does not condition on return-to-go and is trained in a manner of behavior cloning (BC).

Original: We find that both DT and DADT significantly outperform Transformer-BC.
Revised: We find that both DT and DADT significantly outperform Transformer-BC.

Original: 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.
Revised: 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.

Original: 5.4 Jointly optimizing state and action prediction boosts dynamics generalization
Revised: 5.4 Jointly optimizing state and action prediction boosts dynamics generalization

Original: 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.
Revised: To further investigate whether including the proposed auxiliary state prediction loss is helpful for dynamics generalization, we compare our method and alternative training scheme, which also learns state prediction but in a different manner, with respect to next state prediction error in unseen dynamics.

Original: 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.
Revised: In this context, we design DT+, a variant of DT; it is trained using only action prediction loss (as the original DT does), and then the additional linear layer for state prediction is stacked and trained 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 ofdynamics generalization.
Revised: DT+ could support how jointly optimizing state prediction loss can help transformer’s capabilities of dynamics generalization.

Original: 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: As shown in Table 3, DADT exhibits far lower next state prediction error across the environments even if DT+ uses the same amount of data and same objective function.

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

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

Original: [Figure]
Revised: [Figure]

Original: Figure 3 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 measured as the number of updated gradient steps per second on a 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 requires about 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 metric learning [Figure] and MerPO trains separated networks for dynamics model and policy in order, while DADT jointly trains both return-conditioned policy and dynamics information in linear time complexity with single transformer.

Original: 5.6 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 fordynamics generalization of DADT, we comparenext state prediction error of DADT and alternative architectures.
Revised: To further investigate the source of capabilities for dynamics generalization of DADT [Figure] we compare next state prediction error of DADT and alternative architectures.

Original: 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.
Revised: 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 context from history of state-action pairs to predict next state.

Original: 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.
Revised: We remark that MLP only consumes the current state without considering past transitions, while LSTM and CaDM watch a fixed number of past transitions along with recurrent model and ensemble of MLPs, respectively.

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

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

Original: [Figure]
Revised: [Figure]

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

Original: 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.
Revised: To investigate how DADT performs in a low data-quality regime, we compare DADT and baselines with a subset of offline dataset , which is consisted of suboptimal trajectories.

Original: 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.
Revised: Specifically, we build a new offline dataset by selecting trajectories of which the average return is in the bottom 30% of the original offline dataset; note that the dataset for adaptation is not re-built.

Original: As shown in Figure 5,
Revised: As shown in Figure 5,

Original: DADT shows the best performance both in zero-shot and adaptation in the low data-quality regime.
Revised: DADT shows the best performance both in zero-shot and adaptation in the low data-quality regime.

Original: 5.8 How does balancing coefficient affect performance?
Revised: 5.8 How does balancing coefficient affect performance?

Original: of DADT with varying balancing

Original: coeffi-

Original: [Table]
Revised: [Table]

Original: Finally, we investigate how the the balancing coefficient λ ,which determines how state prediction loss occupies totalloss, affects the performance of DADT.
Revised: Finally, we investigate how the the balancing coefficient λ , which determines how state prediction loss occupies total loss, affects the performance of DADT.

Original: Figure 4 shows theoverall performance in unseen dynamics by aggregatingthe average return across all environments with different λ .
Revised: Figure 4 shows the overall performance in unseen dynamics by aggregating the average return across all environments with different λ .

Original: One can observe that the difference in performance by λ isnot significant.
Revised: One can observe that the difference in performance by λ is not significant.

Original: This result indicates that our performanceof DADT is robust to the choice of hyperparameter λ .
Revised: This result indicates that our performance of DADT is robust to the choice of hyperparameter λ .

Original:  Conclusion
Revised:  Conclusion

Original: In this work, we proposed Dynamics-Augmented Decision
Revised: In this work, we proposed Dynamics-Augmented Decision

Original: Transformer (DADT), a simple yet effective transformer-based model for offline dynamics generalization with auxiliary state prediction loss.
Revised: Transformer (DADT), a simple yet effective transformer-based model for offline dynamics generalization with auxiliary state prediction loss.

Original: We have shown that our model outperforms state-of-the-artalgorithms in various control tasks.
Revised: We have shown that our model outperforms state-of-the-art algorithms in various control tasks.

Original: Remarkably, DADT without additional adaptation surpassesfine-tuned baselines in unseen dynamics.
Revised: Remarkably, DADT without additional adaptation surpasses fine-tuned baselines in unseen dynamics.

Original: We believe that DADT will guide new interesting directionsin offline dynamics generalization.
Revised: We believe that DADT will guide new interesting directions in offline dynamics generalization.

Original: Limitation.
Revised: Limitation.

Original: While our work shows impressive results in offline dynamics generalization with variouscontinuous control tasks, we have considered the limited context of varying dynamics (i.e., mass,friction, and inertia).
Revised: While our work shows impressive results in offline dynamics generalization with various continuous control tasks, we have considered the limited context of varying dynamics (i.e., mass, friction, and inertia).

Original: 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.
Revised: However, in the real world, agents often encounter a more diverse and difficult situation of dynamics generalization rather than varying dynamics coefficient, i.e., navigating various shapes of mazes.

Original: 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)).
Revised: Therefore, it would be an interesting future direction to extend our work into a more complex setting with more practical experimental setups like autonomous driving (i.e., CARLA (Dosovitskiy et al., 2017)).

Original: Extending our work into setups where reward function also changes alongwith varying dynamics would be an interesting future step.
Revised: Extending our work into setups where reward function also changes along with varying dynamics would be an interesting future step.

Original: A 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: We largely follow the implementation details of the Gym 1 for resetting agent state and shaping reward function 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 environment code of Humanoid in Gym, following the description in (Wang et al., 2019).

Original: Walker.
Revised: Walker.

Original: Walker is

Original:  a two -dimensional two-legged figure with 7 rigid links including a torso and 2 legs.
Revised: Walker is a two-dimensional two-legged figure with 7 rigid links, including a torso and 2 legs.

Original: The goal is to move forward as fast as possible while maintaining the standing height andconsuming as small control input as possible
Revised: The goal is to move forward as fast as possible while maintaining the standing height and consuming 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.
Revised: Observation is a 17-dimensional vector that includes 1) root joint’s position, 2) angular position and velocities.

Original: • Action.
Revised: • Action.

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

Original: • Reward.
Revised: • 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: 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.

Original: Hopper.
Revised: Hopper.

Revised: Hopper is a

Revised: two dimensional

Original: Hopper is a two-dimensional one-legged figure consisting of 4 rigid links, including a torso,thigh, leg, and foot.
Revised: 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.
Revised: The goal is to move forward as fast as possible while maintaining the standing height and consuming as small control input as possible.

Original: • Observation.
Revised: • 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 and velocity 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: We also 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 of  rigid links ( for torso
Revised: It consists of 7 rigid links ( for torso, 3 for forlimb„ 3 for hindlimb).

Original: 3 for

Original: forlimb„

Original: 3 for

Original: hindlimb).

Original: It is connected by 6 joints, to which an actuator is attached per each joint.
Revised: 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 for
Revised: The goal is to move forward as quickly as possible while minimizing the cost for

Original: controlling agent as small as possible.
Revised: 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 arms unlike 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 move forward 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 of use.

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-based quantities, 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 velocity of the torso, and z t is the height of the torso.
Revised: z t < = 2 . 0) , where ˙ x torso,t represents forward velocity of the torso, and z t is the height of the torso.

Original: We also add an alive bonus of 5 to the agents 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 with[Table] 2 joints.
Revised: Ant is a 3D ant-like robot consisting of a torso with 4 legs, each of which is connected with 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[Table] possible while maintaining the standing height and consuming as small control input as possible.[Table]
Revised: The goal is to move forward as fast as possible while maintaining the standing height and consuming as small control input as possible.

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.[Table]
Revised: Observation is a 27-dimensional vector that includes angular position and velociyu of all 8 joints, except for the x and y positions of the root joints.

Original: • Action.
Revised: • Action.

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

Original: [Table] each leg and torso.• Reward.
Revised: • 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 forwardvelocity of the torso.
Revised: 001 × + 1.0, where ˙ x torso,t represents forward velocity 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.

Original: B Offline dataset details
Revised: B Offline dataset details

Original: [Table]

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 by changing 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 first 50K 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 rollouts from 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 5.
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: We used the minimum and maximum value among test dynamics of each environment for normalizing the average return in evaluation.

Original: [Table]
Revised: [Table]

Original: HalfCheetah

Original: Dynamics Dynamics
Revised: Dynamics Dynamics

Original: Dynamics

Original: Dynamics
Revised: 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: C 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 and evaluating 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 decay of 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, we train 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 the code 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. , 2018)for each test dynamics.
Revised: For random _ 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., 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) number of 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 performance periodically during training and report the best return.

Original: Hyperparamters.
Revised: Hyperparamters.

Original: Hyperparameters for DADT are shown in Table 6.
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 slightly above expert performance.

Original: We heuristically find that training (1) larger DADT and (2) DADT withmore gradient steps boost the model’s performance but use the current version for faster training andevaluation.
Revised: We heuristically find that training (1) larger DADT and (2) DADT with more gradient steps boost the model’s performance but use the current version for faster training and evaluation.

Original: [Table]
Revised: [Table]

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

Original: [Table]
Revised: [Table]

Original: β

Original: [Equation]
Revised: [Equation]

Original: = 0 .

Original: 9 , β

Original: [Equation] = [Table] [Table]
Revised: [Table]

Original: Resources.
Revised: Resources.

Original: For training and evaluating our model, we use a single NVIDIA GeForce RTX 2080 Ti
Revised: 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 for training 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.
Revised: DT (Chen et al., 2021), DADT shows no significant increase in training time.

Original: D Comparison with Trajectory Transformer
Revised: D Comparison with Trajectory Transformer

Original: 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: While our next state prediction loss has a connection with that of Trajectory Transformer (Janner et al., 2021) in that both of them predict next state given past transitions, we remark that DADT does not use discretization, which might be a negative factor in dynamics generalization.

Original: Specifically,approximating a set of values from different dynamics using the exact discrete quantities could hurt[Table] capturing subtle changes in distinguishing the difference between dynamics.[Table]
Revised: Specifically, approximating a set of values from different dynamics using the exact discrete quantities could hurt capturing 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 underperforming of TT is caused by (1) reduced visible context due to per-dimension processing and (2) information loss from discretization.

Original: [Table]
Revised: [Table]