Original: Behavior Transformers: Cloning k modes with one stone
Revised: Behavior Transformers: Cloning k modes with one stone

Original: Abstract
Revised: Abstract

Original: While behavior learning has made impressive progress in recent times, it lags behind computer vision and natural language processing due to its inability to leveragelarge, human-generated datasets.
Revised: While behavior learning has made impressive progress in recent times, it lags behind computer vision and natural language processing due to its inability to leverage large, human-generated datasets.

Original: Human behaviors have wide variance, multiplemodes, and human demonstrations typically do not come with reward labels.
Revised: Human behaviors have wide variance, multiple modes, and human demonstrations typically do not come with reward labels.

Original: Theseproperties limit the applicability of current methods in Offline RL and Behavioral
Revised: These properties limit the applicability of current methods in Ofﬂine RL and Behavioral

Original: Cloning to learn from large, pre-collected datasets.
Revised: Cloning to learn from large, pre-collected datasets.

Original: In this work, we present Behavior Transformer (BeT), a new technique to model unlabeled demonstrationdata with multiple modes.
Revised: In this work, we present Behavior Transformer (BeT), a new technique to model unlabeled demonstration data with multiple modes.

Original: BeT retrofits standard transformer architectures withaction discretization coupled with a multi-task action correction inspired by offsetprediction in object detection.
Revised: BeT retroﬁts standard transformer architectures with action discretization coupled with a multi-task action correction inspired by offset prediction in object detection.

Original: This allows us to leverage the multi-modal modelingability of modern transformers to predict multi-modal continuous actions.
Revised: This allows us to leverage the multi-modal modeling ability of modern transformers to predict multi-modal continuous actions.

Original: Weexperimentally evaluate BeT on a variety of robotic manipulation and self-drivingbehavior datasets.
Revised: We experimentally evaluate BeT on a variety of robotic manipulation and self-driving behavior datasets.

Original: We show that BeT significantly improves over prior state-of-theart work on solving demonstrated tasks while capturing the major modes present inthe pre-collected datasets.
Revised: We show that BeT signiﬁcantly improves over prior state-of-theart work on solving demonstrated tasks while capturing the major modes present in the pre-collected datasets.

Original: Finally, through an extensive ablation study, we analyzethe importance of every crucial component in BeT. Videos of behavior generatedby BeT are available here: https://submission0.github.io .
Revised: Finally, through an extensive ablation study, we analyze the importance of every crucial component in BeT. Videos of behavior generated by BeT are available here: https://notmahi.github.io/bet .

Original:  Introduction
Revised:  Introduction

Original: Creating agents that can behave intelligently in complex environments has been a longstandingproblem in machine learning.
Revised: Creating agents that can behave intelligently in complex environments has been a longstanding problem in machine learning.

Original: Although Reinforcement Learning (RL) has made significant advancesin behavior learning, its success comes at the cost of high sample complexity [51, 21, 1].
Revised: Although Reinforcement Learning (RL) has made signiﬁcant advances in behavior learning, its success comes at the cost of high sample complexity [56, 23, 1].

Original: Withoutpriors on how to behave, state-of-the-art RL methods require online interactions on the order of1-10M ‘reward-labeled’ samples for benchmark control tasks [75].
Revised: Without priors on how to behave, state-of-the-art RL methods require online interactions on the order of 1-10M ‘reward-labeled’ samples for benchmark control tasks [81].

Original: This is in stark contrast to visionand language tasks, where pretrained models and data-driven priors are the norm [17, 10, 28, 5] which allows for efficient downstream task solving.
Revised: This is in stark contrast to vision and language tasks, where pretrained models and data-driven priors are the norm [19, 11, 32, 6], which allows for efﬁcient downstream task solving

Original: So how do we learn behavioral priors from pre-collected data?
Revised: So how do we learn behavioral priors from pre-collected data?

Original: One option is offline RL [41], whereoffline datasets coupled with conservative policy optimization can learn task-specific behaviors.
Revised: One option is ofﬂine RL [46], where ofﬂine datasets coupled with conservative policy optimization can learn task-speciﬁc behaviors.

Original: However, such methods have yet to tackle domains where task-specific reward labels are not present.
Revised: However, such methods have yet to tackle domains where task-speciﬁc reward labels are not present.

Original: Without explicit reward labels, imitation learning, particularly behavior cloning, is a more fittingoption [61, 8, 71].
Revised: Without explicit reward labels, imitation learning, particularly behavior cloning, is a more ﬁtting option [67, 9, 77].

Original: Here, given behavior data D ≡ { s t , a t } , behavior models can be trained topredict actions f θ ( s t ) → a t through supervised learning.
Revised: Here, given behavior data D ≡ { s t , a t } , behavior models can be trained to predict actions f θ ( s t ) → a t through supervised learning.

Original: When demonstration data is plentiful,such approaches have found impressive success in a variety of domains from self-driving [61, 13] torobotic manipulation [78, 55].
Revised: When demonstration data is plentiful, such approaches have found impressive success in a variety of domains from self-driving [67, 14] to a r X i v : s ub m it/ [ c s . L G ]

Original: Importantly, it requires neither online interactions nor reward labels tooptimize behaviors. However, state-of-the-art behavior cloning methods often make a fundamental assumption – that thedata is drawn from a unimodal expert solving a single task. This assumption is often baked in to the
Revised: J un

Original: Rollout
Revised: Rollout

Original: Start
Revised: Start

Original: Reach red
Revised: Reach red

Original: Push red to red
Revised: Push red to red

Original: Reach green
Revised: Reach green

Original: Push green to green
Revised: Push green to green

Original: Start
Revised: Start

Original: Push red to green Push green to red Reach green
Revised: Push red to green Push green to red Reach green

Original: Reach red
Revised: Reach red

Original: Ro l lout
Revised: Ro [Figure] lout

Original: Rollout
Revised: Rollout

Revised: [Figure]

Original: Start
Revised: Start

Original: Microwave
Revised: Microwave

Original: Kettle
Revised: Kettle

Original: Bottom knob
Revised: Bottom knob

Original: Hinge
Revised: Hinge

Original: Start
Revised: Start

Original: Kettle
Revised: Kettle

Original: Light switch
Revised: Light switch

Original: Slide
Revised: Slide

Original: Hinge
Revised: Hinge

Original: Rollout
Revised: Rollout

Revised: [Figure]

Original: Start position
Revised: Start position

Original: Rollout Rollout
Revised: Rollout Rollout

Original: Turn right Turn right
Revised: Turn right Turn right

Original: Turn right
Revised: Turn right

Original: Trajectory Turn left Start
Revised: Trajectory Turn left Start

Original: Start
Revised: Start

Original: Turn right Turn left
Revised: Turn right Turn left

Original: Turn left Turn left
Revised: Turn left Turn left

Original: Trajectory
Revised: Trajectory

Original: CARLA
Revised: [Figure]

Original: Block
Revised: robotic manipulation [84, 61]. Importantly, it requires neither online interactions nor reward labels to optimize behaviors. However, state-of-the-art behavior cloning methods often make a fundamental assumption – that the data is drawn from a unimodal expert solving a single task.

Original: Push Franka Kitchen architecture design, such as using a Gaussian prior.
Revised: This assumption is often baked in to the architecture design, such as using a Gaussian prior.

Original: On the other hand, natural pre-collected data issub-optimal, noisy, and contains multiple modes of behavior, all entangled in a single dataset.
Revised: On the other hand, natural pre-collected data is sub-optimal, noisy, and contains multiple modes of behavior, all entangled in a single dataset.

Original: Thisdistributionally multi-modal experience is most prominent in human demonstrations.
Revised: This distributionally multi-modal experience is most prominent in human demonstrations.

Original: Not only do weperform a large variety of behaviors every day, our personal biases result in significant multi-modalityeven for the same behavior [27, 46].
Revised: Not only do we perform a large variety of behaviors every day, our personal biases result in signiﬁcant multi-modality even for the same behavior [31, 51].

Original: This raises an important question: How do we train models thatcan “clone” multi-modal behavior data?
Revised: This raises an important question: How do we train models that can “clone” multi-modal behavior data?

Original: In this work, we present Behavior Transformers (BeT), a new method for learning behaviors from rich,distributionally multi-modal data.
Revised: In this work, we present Behavior Transformers (BeT), a new method for learning behaviors from rich, distributionally multi-modal data.

Original: BeT is based of three key insights.
Revised: BeT is based of three key insights.

Original: First, we leverage the contextbased multi-token prediction ability of transformer-based sequence models [72] to predict multimodal actions.
Revised: First, we leverage the context based multi-token prediction ability of transformer-based sequence models [78] to predict multimodal actions.

Original: Second, since transformer-based sequence models are naturally suited to predictingdiscrete classes, we cluster continuous actions into discrete bins using k-means [47].
Revised: Second, since transformer-based sequence models are naturally suited to predicting discrete classes, we cluster continuous actions into discrete bins using k-means [52].

Original: This allows usto model high-dimensional, continuous multi-modal action distributions as categorical distributionswithout learning complicated generative models [36, 18].
Revised: This allows us to model high-dimensional, continuous multi-modal action distributions as categorical distributions without learning complicated generative models [41, 20].

Original: Third, to ensure that the actions sampledfrom BeT are useful for online rollouts, we concurrently learn a residual action corrector to producecontinuous actions for a specific sampled action bin.
Revised: Third, to ensure that the actions sampled from BeT are useful for online rollouts, we concurrently learn a residual action corrector to produce continuous actions for a speciﬁc sampled action bin.

Original: We experimentally evaluate BeT on five datasets ranging from simple diagnostic toy datasets tocomplex datasets that include simulated robotic pushing [23], sequential task solving in kitchen environments [30], and self-driving with visual observations in CARLA [19].
Revised: We experimentally evaluate BeT on ﬁve datasets ranging from simple diagnostic toy datasets to complex datasets that include simulated robotic pushing [25], sequential task solving in kitchen environments [34], and self-driving with visual observations in CARLA [21].

Original: The two main findingsfrom these experiments can be summarized as:
Revised: The two main ﬁndings from these experiments can be summarized as:

Original: On multi-modal datasets, BeT achieves significantly higher performance during online rollouts compared to prior behavior modelling methods.
Revised: On multi-modal datasets, BeT achieves signiﬁcantly higher performance during online rollouts compared to prior behavior modelling methods.

Original: Rather than collapsing or latching onto one mode, BeT is able to cover the major modes present in the training behavior datasets.
Revised: Rather than collapsing or latching onto one mode, BeT is able to cover the major modes present in the training behavior datasets.

Original: Unconditional rollouts from this model can be seen in Fig.
Revised: Unconditional rollouts from this model can be seen in Fig.

Original: All of our datasets, code, and trained models will be made publicly available.
Revised: All of our datasets, code, and trained models will be made publicly available.

Original:  Behavior Transformers
Revised:  Behavior Transformers

Original: Given a dataset of continuous observation and action pairs D ≡ { ( o, a ) } ⊂ O × A that containsbehaviors we are interested in, our goal is to learn a behavior policy π : O (cid:55)→
Revised: Given a dataset of continuous observation and action pairs D ≡ { ( o, a ) } ⊂ O × A that contains behaviors we are interested in, our goal is to learn a behavior policy π : O (cid:55)→

Original: A that models thisdata without any online interactions with the environment or reward labels.
Revised: A that models this data without any online interactions with the environment or reward labels.

Original: This setup followsthe Behavior Cloning formulation, where policies are trained to model demonstrations from expertrollouts.
Revised: This setup follows the Behavior Cloning formulation, where policies are trained to model demonstrations from expert rollouts.

Original: Often, such policies are chosen from a hypothesis class parametrized by parameter set θ .
Revised: Often, such policies are chosen from a hypothesis class parametrized by parameter set θ .

Original: Following this convention, our objective is to find the parameter θ that maximizes the probability ofthe observed data
Revised: Following this convention, our objective is to ﬁnd the parameter θ that maximizes the probability of the observed data

Original: [Equation]
Revised: [Equation]

Original: When the model class is restricted to unimodal isotropic Gaussians, this maximum likelihoodestimation problem leads to minimizing the Mean Squared Error (MSE)
Revised: When the model class is restricted to unimodal isotropic Gaussians, this maximum likelihood estimation problem leads to minimizing the Mean Squared Error (MSE), (cid:80) t (cid:107) a t − π ( o t ; θ ) (cid:107)

Original: [Equation]
Revised: [Equation]

Original: [Figure]
Revised: [Figure]

Original: BeT Limitations of traditional MSEbased BC: While MSE-based BChas been able to solve a variety oftasks [8, 71], it assumes that the datadistribution is unimodal.
Revised: MSE. Limitations of traditional MSEbased BC: While MSE-based BC has been able to solve a variety of tasks [9, 77], it assumes that the data distribution is unimodal.

Original: Clean datafrom an expert demonstrator solving a particular task in a particularway satisfies this assumption, but precollected intelligent behavior oftenmay not [46, 30].
Revised: Clean data from an expert demonstrator solving a particular task in a particular way satisﬁes this assumption, but precollected intelligent behavior often may not [51, 34].

Original: While more recent behavior generation models havesought to address this problem, theyoften require complex generative models [70], an exponential number of bins for actions [48], complicated training schemes [59], ortime-consuming test-time optimization [23].
Revised: While more recent behavior generation models have sought to address this problem, they often require complex generative models [76], an exponential number of bins for actions [53], complicated training schemes [65], or time-consuming test-time optimization [25].

Original: An experimental analysis of some of these prior worksis presented in Section 3.
Revised: An experimental analysis of some of these prior works is presented in Section 3.

Original: Overview of Behavior Transformers (BeT): We address two critical assumptions in regular BC.
Revised: Overview of Behavior Transformers (BeT): We address two critical assumptions in regular BC.

Original: First, we relax the assumption that the behavior we are cloning is purely Markovian, and insteadmodel P ( a t | o t , o t − 1 , · · · , o t − h +1 ) for some horizon h .
Revised: First, we relax the assumption that the behavior we are cloning is purely Markovian, and instead model P ( a t | o t , o t − 1 , · · · , o t − h +1 ) for some horizon h .

Original: Second, instead of assuming that actionsare generated by a unimodal action distribution, we model our action distribution as a mixture ofgaussians.
Revised: Second, instead of assuming that actions are generated by a unimodal action distribution, we model our action distribution as a mixture of gaussians.

Original: However, unlike previous efforts similar to Mixture Density Networks (MDN) to doso, whose limitations have been explored in Florence et al.
Revised: However, unlike previous efforts similar to Mixture Density Networks (MDN) to do so, whose limitations have been explored in Florence et al.

Original: [23], we do not explicitly predict modecenters, which significantly improves our modeling capacity.
Revised: [25], we do not explicitly predict mode centers, which signiﬁcantly improves our modeling capacity.

Original: To operationalize these two features ina single behavior model, we make use of transformers since (a) they are effective in utilizing priorobservational history, and (b) they are naturally suited to output multi-modal tokens through theirarchitecture.
Revised: To operationalize these two features in a single behavior model, we make use of transformers since (a) they are effective in utilizing prior observational history, and (b) their architecture is naturally suited to output multi-modal tokens. [Figure]

Original: 2.1 Action discretization for distribution learning
Revised: 2.1 Action discretization for distribution learning

Original: Although transformers have become standard as a backbone for sequence-to-sequence models [17, 10],they are designed to process discrete tokens and not continuous values.
Revised: Although transformers have become standard as a backbone for sequence-to-sequence models [19, 11], they are designed to process discrete tokens and not continuous values.

Original: In fact, modeling multi-modaldistributions of high-dimensional continuous variables in a tractable manner is in itself a challengingproblem, especially if we want the trained behavior model to cover the modes present in the dataset. 
Revised: In fact, modeling multi-modal distributions of high-dimensional continuous variables in a tractable manner is in itself a challenging problem, especially if we want the trained behavior model to cover the modes present in the dataset. To address this

Original: [Figure]
Revised: [Figure]

Original: To address this, we propose a new factoring of the action prediction task by dividing each action intwo parts: a categorical variable denoting an ‘action center’, and a corresponding ‘residual action’.
Revised: we propose a new factoring of the action prediction task by dividing each action in two parts: a categorical variable denoting an ‘action center’, and a corresponding ‘residual action’.

Original: To this end, given the actions in our dataset, we first optimize for a set of k action centers,{ A 1 , A 2 , · · · , A k } ⊂ A .
Revised: To this end, given the actions in our dataset, we ﬁrst optimize for a set of k action centers, { A 1 , A 2 , · · · , A k } ⊂ A .

Original: We then decompose each action into two parts: a categorical variablerepresenting the closest action bin, ⌊ a ⌋ := arg min i ∥ a − A i ∥ 2 , and a continuous residual action⟨ a ⟩ := a − A ⌊ a ⌋ .
Revised: We then decompose each action into two parts: a categorical variable representing the closest action bin, (cid:98) a (cid:99) := arg min i (cid:107) a − A i (cid:107) 2 , and a continuous residual action (cid:104) a (cid:105) : = a − A (cid:98) a (cid:99) .

Original: If we are given the set of action centers { A i } ki =1 , an action bin index ⌊ a ⌋ andthe residual action ⟨ a ⟩ , we can deterministically reconstruct the true action a :=
Revised: If we are given the set of action centers { A i } ki =1 , an action bin index (cid:98) a (cid:99) and the residual action (cid:104) a (cid:105) , we can deterministically reconstruct the true action a :

Original:  A ⌊ a ⌋ + ⟨ a ⟩ .
Revised: = A (cid:98) a (cid:99) + (cid:104) a (cid:105) .

Original: Oncelearned, these k-means based encoder and decoders for this action factorization process are fixed forthe rest of the train and testing phases.
Revised: Once learned, these k-means based encoder and decoders for this action factorization process are ﬁxed for the rest of the train and testing phases.

Original: The action factorization procedure is illustrated in Fig.
Revised: The action factorization procedure is illustrated in Fig.

Original: 3 (A).
Revised: 3 (A).

Original: 2.2 Attention-based behavior mode learning
Revised: 2.2 Attention-based behavior mode learning

Original: Once we have the clustering based autoencoder learned from the actions in the dataset, we model ourdemonstration trajectories with BeT. We use a transformer decoder model, namely minGPT [10], withminor modifications, as our backbone.
Revised: Once we have the clustering based autoencoder learned from the actions in the dataset, we model our demonstration trajectories with BeT. We use a transformer decoder model, namely minGPT [11], with minor modiﬁcations, as our backbone.

Original: The transformer T takes in a sequence of continuous observations ( o i , o i +1 , · · · , o i + h − 1 ) and learns a sequence-to-sequence model mapping each observationto a categorical distribution over k discrete action bins.
Revised: The transformer T takes in a sequence of continuous observations ( o i , o i +1 , · · · , o i + h − 1 ) and learns a sequence-to-sequence model mapping each observation to a categorical distribution over k discrete action bins.

Original: The predicted probability sequence is thencompared with the ground truth labels, ( ⌊ a i ⌋ , ⌊ a i +1 ⌋ , ⌊ a i +2 ⌋ , · · · , ⌊ a i + h − 1 ⌋ ) .
Revised: The predicted probability sequence is then compared with the ground truth labels, ( (cid:98) a i (cid:99) , (cid:98) a i +1 (cid:99) , (cid:98) a i +2 (cid:99) , · · · , (cid:98) a i + h − 1 (cid:99) ) .

Original: We use a negativelog-likelihood-based Focal loss [43] between the predicted categorical distribution probabilities andthe ground truth labels to train the transformer head.
Revised: We use a negative log-likelihood-based Focal loss [48] between the predicted categorical distribution probabilities and the ground truth labels to train the transformer head.

Revised: Focal loss is a simple modiﬁcation over the standard cross entropy loss.

Revised: While the standard cross entropy loss for binary classiﬁcation can be thought of L ce ( p t )

Revised: = − log( p t ) , Focal loss adds a term (1 − p t ) γ to this, to make the new loss

Original:  The model is illustrated in Fig.
Revised: [Equation] This loss has the interesting property that its gradient is more steep for smaller values of p t , while ﬂatter for larger values of p t . Thus, it penalizes and changes the model more for making errors in the low-probability classes, while is more lenient about making errors in the high probability classes. The model is illustrated in Fig.

Original: 3 (B).
Revised: 3 (B).

Original: 2.3 Action correction: from coarse to finer-grained predictions
Revised: 2.3 Action correction: from coarse to ﬁner-grained predictions

Original: Using a transformer allows us to model multi-modal actions.
Revised: Using a transformer allows us to model multi-modal actions.

Original: However, discretizing the continuousaction space in any way invariably causes loss of fidelity [34].
Revised: However, discretizing the continuous action space in any way invariably causes loss of ﬁdelity [38].

Original: Discretization error may cause onlinerollouts of the behavior policy to go out of distribution from the original dataset [67], which can inturn cause critical failures.
Revised: Discretization error may cause online rollouts of the behavior policy to go out of distribution from the original dataset [73], which can in turn cause critical failures.

Original: To predict the complete continuous action, we add an extra head to thetransformer decoder that offsets the discretized action centers based on the observations.
Revised: To predict the complete continuous action, we add an extra head to the transformer decoder that offsets the discretized action centers based on the observations.

Original: For each observation o i in the sequence, the head produces a k × dim ( A ) matrix with k proposedresidual action vectors, (cid:16) ⟨ a ( j ) i ⟩ (cid:17)
Revised: For each observation o i in the sequence, the head produces a k × dim ( A ) matrix with k proposed residual action vectors, (cid:16) (cid:104) a ( j ) i (cid:105) (cid:17)

Original: k
Revised: k

Revised: j =1 =

Revised: ( (cid:104)

Original: j =1 = ( ⟨ ˆ a (1) i ⟩ , ⟨ ˆ a (2) i ⟩ , ⟨ ˆ a (3) i ⟩ , · · · , ⟨ ˆ a ( k ) i ⟩ ) , where the residual actionscorrespond to bin centers A 1 , A 2 , A 3 , · · · , A k .
Revised: ˆ a (1) i (cid:105) , (cid:104) ˆ a (2) i (cid:105) , (cid:104) ˆ a (3) i (cid:105) , · · · , (cid:104) ˆ a ( k ) i (cid:105) ) , where the residual actions correspond to bin centers A 1 , A 2 , A 3 , · · · , A k .

Original: These residual actions are trained with a loss akin to the masked multi-task loss [26] from object detection.
Revised: These residual actions are trained with a loss akin to the masked multi-task loss [30] from object detection.

Original: In our case, if the ground truth action is a , theloss is:
Revised: In our case, if the ground truth action is a , the loss is:

Original: MT-Loss (cid:18) a , (cid:16) ⟨ ˆ a ( j ) i ⟩ (cid:17) k

Original: j = (cid:19) = k (cid:88) j = I [ ⌊ a ⌋ =
Revised: [Equation]

Original: j ] · ∥⟨ a ⟩ − ⟨ ˆ a ( j ) ⟩∥

Original: Where I [] denotes the Iverson bracket, ensuring the offset head of BeT only incurs loss from theground truth class of action a .
Revised: Where I [] denotes the Iverson bracket, ensuring the offset head of BeT only incurs loss from the ground truth class of action a .

Original: This mechanism prevents the model from trying to fit the ground truthaction using the offset at every index.
Revised: This mechanism prevents the model from trying to ﬁt the ground truth action using the offset at every index.

Original: 2.4 Test-time sampling from BeT
Revised: 2.4 Test-time sampling from BeT

Original: During test time, at timestep t we input the latest h observations ( o t , o t − 1 , · · · , o t − h +1 ) to thetransformer, combining the present observation o t with h − 1 previous observations.
Revised: During test time, at timestep t we input the latest h observations ( o t , o t − 1 , · · · , o t − h +1 ) to the transformer, combining the present observation o t with h − 1 previous observations.

Original: Our trained
Revised: Our trained

Original: MinGPT model gives us h × 1 × k bin center probability vectors, and h × k × dim ( A ) offset matrix.
Revised: MinGPT model gives us h × 1 × k bin center probability vectors, and h × k × dim ( A ) offset matrix.

Original: Tosample an action at timestep t , we first sample an action center according to the predicted bin centerprobabilities on the t th index.
Revised: To sample an action at timestep t , we ﬁrst sample an action center according to the predicted bin center probabilities on the t th index.

Original: Once we have chosen an action center A t,j , we add the correspondingresidual action ⟨ ˆ a ( j ) t ⟩ to it to recover a predicted continuous action ˆ a t =
Revised: Once we have chosen an action center A t,j , we add the corresponding residual action (cid:104) ˆ a ( j ) t (cid:105) to it to recover a predicted continuous action ˆ a t =

Original: A t,j + ⟨ ˆ a ( j ) t ⟩ .
Revised: A t,j + (cid:104) ˆ a ( j ) t (cid:105) .

Original: This samplingprocedure is illustrated in Fig.
Revised: This sampling procedure is illustrated in Fig.

Original: 3 (C).
Revised: 3 (C).

Original:  Experiments
Revised:  Experiments

Original: We now study the empirical performance of BeT on a variety of behavior learning tasks.
Revised: We now study the empirical performance of BeT on a variety of behavior learning tasks.

Original: Ourexperiments are designed to answer the following questions: (a) Is BeT able to imitate multi-modaldemonstrations?
Revised: Our experiments are designed to answer the following questions: (a) Is BeT able to imitate multi-modal demonstrations?

Original: (b) How well does BeT capture the modes present in behavior data?
Revised: (b) How well does BeT capture the modes present in behavior data?

Original: (c) Howimportant are the individual components of BeT?
Revised: (c) How important are the individual components of BeT?

Original: 3.1 Environments and datasets
Revised: 3.1 Environments and datasets

Original: We experiment with five broad environments.
Revised: We experiment with ﬁve broad environments.

Original: While full descriptions of these environments, datasetcreation procedure, and overall statistics are in Appendix A, a brief description of them are as follows.
Revised: While full descriptions of these environments, dataset creation procedure, and overall statistics are in Appendix A, a brief description of them are as follows.

Original: (a) Point mass environment #1: Our first set 
Revised: (a) Point mass environment #1: Our ﬁrst set of experiments in Fig.

Original: of experiments in Fig.

Original: used

Original: to get

Original: a

Original: qualitativeunderstanding of BeT were performed in a simple Pointmass environment 
Revised: 2, used to get a qualitative understanding of BeT, were performed in a simple Pointmass environment with a 2D observation and action space with two hundred demonstrations.

Original: with a

Original: 2D

Original: observation

Original: and action space with two hundred demonstrations.

Original: The pre-collected demonstrations start at a

Original: fixed

Original: point

Original:  and then make their way to another point while avoiding a block in the middle.
Revised: The pre-collected demonstrations start at a ﬁxed point, and then make their way to another point while avoiding a block in the middle.

Original: The two primary modes in this dataset are taking a left turn versus a right turn.
Revised: The two primary modes in this dataset are taking a left turn versus a right turn.

Original: (b) Point mass environment #2: The setup is similar to the previous environment with the exceptionof one straight line and two complicated prolonged ‘Z’ shaped modes of demonstration (Fig. 5.)
Revised: (b) Point mass environment #2: The setup is similar to the previous environment with the exception of one straight line and two complicated prolonged ‘Z’ shaped modes of demonstration (Fig. 5.)

Original: (c) CARLA self-driving environment: CARLA [19] uses the Unreal Engine to provide a sim- ulated driving environment in a visually realistic landscape.
Revised: (c) CARLA self-driving environment : CARLA [21] uses the Unreal Engine to provide a simulated driving environment in a visually realistic landscape.

Original: The agent action space is 2D (accelerate/brake and left/right steer), while the observation space is (224,224,3)-dimensional
Revised: The agent action space is 2D (accelerate/brake and left/right steer), while the observation space is (224,224,3)-dimensional

Original: RGB image from the car.
Revised: RGB image from the car.

Original: A hundred total demonstrations drive around a building block in
Revised: A hundred total demonstrations drive around a building block in

Original: two distinct modes.
Revised: two distinct modes.

Original: This environment highlights the challenge of behavior learning from high- dimensional observations as shown in Fig.
Revised: This environment highlights the challenge of behavior learning from highdimensional observations as shown in Fig.

Original: 1 (a).
Revised: 1 (a).

Original: For visual observations with BeT, we use a frozen ResNet-18 [31] pretrained on ImageNet [16] as an encoder.
Revised: For visual observations with BeT, we use a frozen ResNet-18 [35] pretrained on ImageNet [18] as an encoder.

Original: (d) Multi-modal block-pushing environment: For more complicated interaction data, we use the multi-modal block-pushing environment from Implicit Behavioral Cloning (IBC)[23], where
Revised: (d) Multi-modal block-pushing environment: For more complicated interaction data, we use the multi-modal block-pushing environment from Implicit Behavioral Cloning (IBC)

Original:  an XArm robot needs to push two blocks into two squares in any order.
Revised: [25], where an XArm robot needs to push two blocks into two squares in any order.

Original: The blocks and target squares are colored red and green.
Revised: The blocks and target squares are colored red and green.

Original: The positions of the blocks are randomized at episode start.
Revised: The positions of the blocks are randomized at episode start.

Original: We collect 1,000 demonstrations using a deterministic controller with two independent axes of multi-modality: (a) it starts by reaching either the red or the green block, with 50% probability, and (b) it pushes the blocks to (red, green) or (green, red) squares respectively with 50% probability.
Revised: We collect 1,000 demonstrations using a deterministic controller with two independent axes of multi-modality: (a) it starts by reaching either the red or the green block, with 50% probability, and (b) it pushes the blocks to (red, green) or (green, red) squares respectively with 50% probability.

Original: (e) Franka kitchen environment: To highlight the complexity of performing long sequences of actions, we use the Relay Kitchen Environment [30] where a Franka robot manipulates a virtual kitchen environment.
Revised: (e) Franka kitchen environment: To highlight the complexity of performing long sequences of actions, we use the Relay Kitchen Environment [34] where a Franka robot manipulates a virtual kitchen environment.

Original: We use the relay policy learning dataset with 566 demonstrations collectedby human participants wearing VR headsets.
Revised: We use the relay policy learning dataset with 566 demonstrations collected by human participants wearing VR headsets.

Original: The participants completed a sequence of four[Table] object-interaction tasks in each episode [30].
Revised: The participants completed a sequence of four object-interaction tasks in each episode [34].

Original: There are a total of seven interactable objects inthe kitchen: a microwave, a kettle, a slide cabinet, a hinge cabinet, a light switch [Table] and two burnerknobs.
Revised: There are a total of seven interactable objects in the kitchen: a microwave, a kettle, a slide cabinet, a hinge cabinet, a light switch, and two burner knobs.

Original: This dataset contains two different kinds of multi-modality: one from the inherent noisein human demonstrations, and another from the demonstrators’ intent.
Revised: This dataset contains two different kinds of multi-modality: one from the inherent noise in human demonstrations, and another from the demonstrators’ intent.

Original: 3.2 Baseline behavior learning methods
Revised: 3.2 Baseline behavior learning methods

Original: While a full description of our baselines are in Appendix B.1, a brief description of them is here
Revised: While a full description of our baselines are in Appendix B. a brief description of them is here:

Original: (a) Multi-layer Perceptron with MSE (RBC): We use MLP networks trained with MSE loss asour first baseline, since this is the standard way of performing behavioral cloning for a new task [71].
Revised: (a) Multi-layer Perceptron with MSE (RBC): We use MLP networks trained with MSE loss as our ﬁrst baseline, since this is the standard way of performing behavioral cloning for a new task [77].

Original: A comparison with transformer-based behavior cloning is discussed in Section 3.5.
Revised: A comparison with transformer-based behavior cloning is discussed in Section 3.5.

Original: (b) Nearest neighbor (NN): Nearest neighbor based algorithms are easy to implement, and hasrecently shown to have strong performance on complicated behavioral cloning tasks [3].
Revised: (b) Nearest neighbor (NN): Nearest neighbor based algorithms are easy to implement, and has recently shown to have strong performance on complicated behavioral cloning tasks [3].

Original: (c) Locally Weighted Regression (LWR): This non-parametric approach provides better regularization compared to NN and is a strong alternative to parametric BC [4, 55].
Revised: (c) Locally Weighted Regression (LWR): This non-parametric approach provides better regularization compared to NN and is a strong alternative to parametric BC [4, 61].

Original: (d) Variational auto-encoders (VAE): Inspired by SPiRL [59], where behavioral priors are learnedthrough a VAE [36], we compare with continuous actions generated from the VAE and the prior.(e)
Revised: (d) Variational auto-encoders (VAE): Inspired by SPiRL [65], where behavioral priors are learned through a VAE [41], we compare with continuous actions generated from the VAE and the prior.

Original:  Normalizing Flow (Flow): Inspired by PARROT [70], where state-conditioned action priors arelearned through a Flow model [18], we compare with actions generated from the Flow model.(f)
Revised: (e) Normalizing Flow (Flow): Inspired by PARROT [76], where state-conditioned action priors are learned through a Flow model [20], we compare with actions generated from the Flow model.

Original:  Implicit Behavioral Cloning (IBC): Instead of modeling the conditional distribution P ( a | o ) ,IBC models the joint probability distribution P ( a, o ) using energy-based models [23].
Revised: (f) Implicit Behavioral Cloning (IBC): Instead of modeling the conditional distribution P ( a | o ) , IBC models the joint probability distribution P ( a, o ) using energy-based models [25].

Original: WhileIBC is slower than explicit BC models because of their sampling requirements, they have been shown to learn well on multi-modal data, and outperform earlier work such as MDNs [7].
Revised: While IBC is slower than explicit BC models because of their sampling requirements, they have been shown to learn well on multi-modal data, and outperform earlier work such as MDNs [8].

Original: 3.3 Is BeT able to imitate multi-modal demonstrations?
Revised: 3.3 Is BeT able to imitate multi-modal demonstrations?

Original: The first question we ask is whether BeT can actually clone behaviors given a mixed dataset ofunlabeled, multi-modal behaviors.
Revised: The ﬁrst question we ask is whether BeT can actually clone behaviors given a mixed dataset of unlabeled, multi-modal behaviors.

Original: To examine that, we look at the performance of our model in
Revised: To examine that, we look at the performance of our model in

Original: CARLA, Block push, and Kitchen environments compared with our baselines in Table 1.
Revised: CARLA, Block push, and Kitchen environments compared with our baselines in Table 1.

Original: We see that BeT outperforms all other methods in all environments except CARLA, where it isnarrowly outperformed by LWR.
Revised: We see that BeT outperforms all other methods in all environments except CARLA, where it is narrowly outperformed by LWR.

Original: Since the models are all behavioral cloning algorithms, they sharethe failure mode of failing once the observations go out of distribution (OOD).
Revised: Since the models are all behavioral cloning algorithms, they share the failure mode of failing once the observations go out of distribution (OOD).

Original: However, they vary inthe tolerance.
Revised: However, they vary in the tolerance.

Original: For example, BeT shines in the Block push environment, where alongside extremeenvironment randomness and multi-modality, the models also have to learn significant long-termbehaviors and commit to a single mode over a long period.
Revised: For example, BeT shines in the Block push environment, where alongside extreme environment randomness and multi-modality, the models also have to learn signiﬁcant long-term behaviors and commit to a single mode over a long period.

Original: While all baselines can somewhatsuccessfully reach one block, they fail to complete the long-horizon, multi-modal task of pushingtwo blocks into two different bins.
Revised: While all baselines can somewhat successfully reach one block, they fail to complete the long-horizon, multi-modal task of pushing two blocks into two different bins.

Original: On the other hand, we observe that BeT’s primary failure mode isnot realizing a block has not completely entered the target yet, while other methods either go OODquickly, or keep switching between modes.
Revised: On the other hand, we observe that BeT’s primary failure mode is not realizing a block has not completely entered the target yet, while other methods either go OOD quickly, or keep switching between modes.

Original: We also observe that BeT performs well even in complex 
Revised: We also observe that BeT performs well even in complex observation and action spaces. In the CARLA environment, the model takes in visual observations, [Table]

Original: [Figure]
Revised: [Figure]

Original: m b l s m t b h t b h s m
Revised: [Table] [Figure] [Table] [Figure]

Original:  k b s m bh s t b l s k b l s k [Figure]
Revised: [Figure] Task sequence k b . . m b s .

Original: k t
Revised: k b s .

Original: [Figure] k [Figure] [Figure]
Revised: k b l . m . . .

Original: m b . 
Revised: m b . .

Original: [Figure] t . [Figure] k . . 
Revised: k b t s k . . .

Original: [Figure] [Figure] [Figure] Demonstration Dataset k-NN + LWR
Revised: k b t .

Original: Implicit BC
Revised: m b t .

Original: Behavior Transformer CARLA Block: first block reached
Revised: 200 Frequency 0 1 2 3 Task sequence k b s h b t s .

Original: Push: red block target Push: green block target
Revised: b t s h b t . . k b s .

Original: Kitchen
Revised: k b . .

Original: Baselines Left Right Red
Revised: k b t s k b t .

Original: Green Red Green
Revised: b t l s b s . .

Original: Red Green
Revised: [Figure]

Original: Task entropy
Revised: Frequency

Original: RBC 0.

Original: [Table]
Revised: [Figure]

Original: observation and action spaces. In the CARLA environment, the model takes in visual observations,while in the Franka Kitchen environment, the action space corresponds to a 9-DOF torque controlledrobot.
Revised: Behavior Transformer while in the Franka Kitchen environment, the action space corresponds to a 9-DOF torque controlled robot.

Original: BeT handles both cases with the same ease as it does environments with lower-dimensionalobservation or action spaces.
Revised: BeT handles both cases with the same ease as it does environments with lower-dimensional observation or action spaces.

Original: 3.4 Does BeT capture the modes present in behavior data?
Revised: 3.4 Does BeT capture the modes present in behavior data?

Original: Next, we examine the question of whether, given a dataset where multi-modal behavior exists, ourmodel learns behavior that is also multi-modal.
Revised: Next, we examine the question of whether, given a dataset where multi-modal behavior exists, our model learns behavior that is also multi-modal.

Original: Here, we are interested in seeing the variance of thebehavior of the model over different rollouts.
Revised: Here, we are interested in seeing the variance of the behavior of the model over different rollouts.

Original: In each of our environments, the demonstrations containdifferent types of multi-modality.
Revised: In each of our environments, the demonstrations contain different types of multi-modality.

Original: As a result, we show a comprehensive analysis of multi-modalityseen in our agent behaviors.
Revised: As a result, we show a comprehensive analysis of multi-modality seen in our agent behaviors.

Original: We see in Table 2 that in CARLA and Block push, BeT covers all the modes of the demonstrationdata, even in the few cases where it does not perfectly match the demonstrated task probabilities.
Revised: We see in Table 2 that in CARLA and Block push, BeT covers all the modes of the demonstration data, even in the few cases where it does not perfectly match the demonstrated task probabilities.

Original: Forthe Kitchen environment, we see in Fig.
Revised: For the Kitchen environment, we see in Fig.

Original: 4 that BeT visits certain strings of tasks more frequently thanin the original demonstrations.
Revised: 4 that BeT visits certain strings of tasks more frequently than in the original demonstrations.

Original: However, compared to other strong baselines, BeT generates longertask strings more often while maintaining diversity and not collapsing to a single mode.
Revised: However, compared to other strong baselines, BeT generates longer task strings more often while maintaining diversity and not collapsing to a single mode.

Original: 3.5 How important are the individual components of BeT?
Revised: 3.5 How important are the individual components of BeT?

Original: There are four key differences between BeT architecture and standard BC: (a) binning actions intodiscrete clusters, (b) using offsets to faithfully reconstruct actions later, (c) learning sequentially to Dataset MLP + MSE BC BeT with no history
Revised: There are four key differences between BeT architecture and standard BC: (a) binning actions into discrete clusters, (b) using offsets to faithfully reconstruct actions later, (c) learning sequentially to use historical context, and (d) using an attention-based MinGPT trunk.

Original: BeT with history use historical context, and (d) using an attention-based MinGPT trunk.

Original: In this section, we discuss theimpacts they have in BeT performance.
Revised: In this section, we discuss the impacts they have in BeT performance.

Original: [Table]
Revised: [Table]

Original:  Impact of discrete binning: Intuitively, having discrete options for bin centers is what enables BeT to express multi-modal behavior evenwhen starting from an identical starting state.
Revised: [Figure] [Table] Impact of discrete binning: Intuitively, having discrete options for bin centers is what enables BeT to express multi-modal behavior even when starting from an identical starting state.

Original: Indeed, if there is no binning, we see from Table 3 that the performance of BeT drops significantly.
Revised: Indeed, if there is no binning, we see from Table 3 that the performance of BeT drops significantly.

Original: More tellingly, in the Franka Kitchenenvironment, the model only ever completed asubsequence of (kettle, top/bottom burner, light switch, slide cabinet) tasks after 100 random rollouts.
Revised: More tellingly, in the Franka Kitchen environment, the model only ever completed a subsequence of (kettle, top/bottom burner, light switch, slide cabinet) tasks after 100 random rollouts.

Original: This result shows us that having discrete bins helps BeT achieve multi-modality.
Revised: This result shows us that having discrete bins helps BeT achieve multi-modality.

Original: To find the best valueof the number of bins, k , we fit a Bayesian GMM [66] to our action dataset and use the estimatednumber of clusters as a guide.
Revised: To ﬁnd the best value of the number of bins, k , we ﬁt a Bayesian GMM [72] to our action dataset and use the estimated number of clusters as a guide.

Original: We show a plot of k vs. performance on the tasks in the Appendix C.2.
Revised: We show a plot of k vs. performance on the tasks in the Appendix C.2.

Original: Necessity of action offsets: An important feature of BeT is the residual action offset that correctsthe discrete actions coming from the bins.
Revised: Necessity of action offsets: An important feature of BeT is the residual action offset that corrects the discrete actions coming from the bins.

Original: While the bin centers may be quite expressive, Table 3shows that the inability to correct them causes a performance degrade.
Revised: While the bin centers may be quite expressive, Tableshows that the inability to correct them causes a performance degrade.

Original: Interestingly, the largestdegradation comes in the Kitchen environment, which also has the highest dimensional action space.
Revised: Interestingly, the largest degradation comes in the Kitchen environment, which also has the highest dimensional action space.

Original: Intuitively, we can understand how in higher dimension the loss of fidelity from discretizing wouldbe higher, and the relative performance loss across three environments support that hypothesis.
Revised: Intuitively, we can understand how in higher dimension the loss of ﬁdelity from discretizing would be higher, and the relative performance loss across three environments support that hypothesis.

Original: Importance of historical context: While RL algorithms traditionally assume environments are
Revised: Importance of historical context: While RL algorithms traditionally assume environments are

Original: Markovian, human behavior in an open-ended environment is rarely so.
Revised: Markovian, human behavior in an open-ended environment is rarely so.

Original: Thus, using historicalcontext helps BeT to perform well.
Revised: Thus, using historical context helps BeT to perform well.

Original: We show a simple experiment in Fig.
Revised: We show a simple experiment in Fig.

Original: 5 on the second point massenvironment.
Revised: 5 on the second point mass environment.

Original: Here, training and evaluating with some historical context allows BeT to follow thedemonstrations better.
Revised: Here, training and evaluating with some historical context allows BeT to follow the demonstrations better.

Original: We experience the same in the CARLA, Block push, and Kitchen environments,where training with some historical context raises performance across the board as seen in Table 3.
Revised: We experience the same in the CARLA, Block push, and Kitchen environments, where training with some historical context raises performance across the board as seen in Table 3.

Original: Since transformer-based models are generally able to learn from long-term historical context, webelieve BeT should also be able to model real-world long-term behavior patterns.
Revised: Since transformer-based models are generally able to learn from long-term historical context, we believe BeT should also be able to model real-world long-term behavior patterns.

Original: Importance of transformer architecture: Despite transformers’ success in other fields of machinelearning, it is natural to wonder whether the tasks BeT solves here really requires one.
Revised: Importance of transformer architecture: Despite transformers’ success in other ﬁelds of machine learning, it is natural to wonder whether the tasks BeT solves here really requires one.

Original: We ablated
Revised: We ablated

Original: BeT by replacing the MinGPT trunk withan MLP,  and found that they have lower performance whilealso being difficult to train stably. This performance reduction remains even if the MLP is givensome historical context by concatenatingh observations before passing it to the MLP.
Revised: BeT by replacing the MinGPT trunk with MLP, temporal convolution, and LSTM architecture based models, and found that they have lower performance while also being difﬁcult to train stably. We gave the MLP equal historical context by concatenating the last few frames.

Original: See Table.
Revised: See Table.

Original: 3for results and Appendix C.3 for further details.
Revised: 3 for results and Appendix C.3 for further details.

Original: Computation considerations: While transformers in usual contexts are large models, we downscalethem for our application in BeT (See Appendix B.4).
Revised: Computation considerations: While transformers in usual contexts are large models, we downscale them for our application in BeT (See Appendix B.4).

Original: Our models contain on the order of 10 4 − 10 6parameters, and even with a small batch size and on a single desktop GPU trains within an hour forour largest datasets (Block push).
Revised: Our models contain on the order of 10 4 − 10 6 parameters, and even with a small batch size and on a single desktop GPU trains within an hour for our largest datasets (Block push).

Original: In contrast, for the same task, our strongest baseline IBC takesabout 14 hours.
Revised: In contrast, for the same task, our strongest baseline IBC takes about 14 hours.

Original: On the same environment, 100 evaluation rollouts take about 165 seconds with BeT,as opposed to 1770 seconds with IBC.
Revised: On the same environment, 100 evaluation rollouts take about 165 seconds with BeT, as opposed to 1770 seconds with IBC.

Original:  Related Work
Revised:  Related Work

Original: This paper builds upon a rich literature in imitation learning, offline learning, generative models, andtransformer architectures.
Revised: This paper builds upon a rich literature in imitation learning, ofﬂine learning, generative models, and transformer architectures.

Original: The most relevant ones to our work are discussed here.
Revised: The most relevant ones to our work are discussed here.

Original: Learning from offline data: Since Pomerleau [60] showed the possibility of driving an autonomousvehicle using offline data and a neural network, learning behavior from offline data has been acontinuous topic of research for scalable behavior learning [2, 6, 69].
Revised: Learning from ofﬂine data: Since Pomerleau [66] showed the possibility of driving an autonomous vehicle using ofﬂine data and a neural network, learning behavior from ofﬂine data has been a continuous topic of research for scalable behavior learning [2, 7, 75].

Original: The approaches can be dividedinto two broad classes: Offline RL [25, 38, 39, 74, 41, 24], focusing on learning from datasets ofa mixed quality that also have reward labels; and imitation learning [54, 56, 57, 32], focusing onlearning behavior from a dataset of expert behavior without reward labels.
Revised: The approaches can be divided into two broad classes: Ofﬂine RL [27, 43, 44, 80, 46, 26], focusing on learning from datasets of a mixed quality that also have reward labels; and imitation learning [60, 62, 63, 36], focusing on learning behavior from a dataset of expert behavior without reward labels.

Original: BeT falls under thesecond category, as it is a behavior cloning model.
Revised: BeT falls under the second category, as it is a behavior cloning model.

Original: Behavior cloning is a form of imitation learningthat tries to model the action of the expert given the observation which is often used in real-worldapplications [78, 79, 78, 62, 22, 77].
Revised: Behavior cloning is a form of imitation learning that tries to model the action of the expert given the observation which is often used in real-world applications [84, 85, 84, 68, 24, 83].

Original: As behavior cloning algorithms are generally solving a fullysupervised learning problem, they tend to be faster and simpler than reinforcement learning or offline RL algorithms and in some cases show competitive results [24, 29].
Revised: As behavior cloning algorithms are generally solving a fully supervised learning problem, they tend to be faster and simpler than reinforcement learning or ofﬂine RL algorithms and in some cases show competitive results [26, 33].

Original: Generative models for behavior learning: One approach for imitation learning is Inverse Reinforcement Learning or IRL [68, 53], where given expert demonstrations, a model tries to construct thereward function.
Revised: Generative models for behavior learning: One approach for imitation learning is Inverse Reinforcement Learning or IRL [74, 58], where given expert demonstrations, a model tries to construct the reward function.

Original: This reward function is then used to generate desirable behavior.
Revised: This reward function is then used to generate desirable behavior.

Original: GAIL [32], an IRLalgorithm, connects generative adversarial models with imitation learning to construct a model thatcan generate expert-like behavior.
Revised: GAIL [36], an IRL algorithm, connects generative adversarial models with imitation learning to construct a model that can generate expert-like behavior.

Original: Under this IRL framework, previous works have tried to predictmulti-modal, multi-human trajectories [40, 33].
Revised: Under this IRL framework, previous works have tried to predict multi-modal, multi-human trajectories [45, 37].

Original: Similarly, other works have tried Gaussian Processes[64] for creating dynamical models for human motion [73].
Revised: Similarly, other works have tried Gaussian Processes [70] for creating dynamical models for human motion [79].

Original: Another class of algorithms learn agenerative action decoder [58, 46, 70] from interaction data to make downstream reinforcementlearning faster and easier, which inspired BeT’s action factorization.
Revised: Another class of algorithms learn a generative action decoder [64, 51, 76] from interaction data to make downstream reinforcement learning faster and easier, which inspired BeT’s action factorization.

Original: Finally, a class of algorithms,most notably [44, 23, 37, 52] do not directly learn a generative model but instead learn energy basedmodels.
Revised: Finally, a class of algorithms, most notably [49, 25, 42, 57] do not directly learn a generative model but instead learn energy based models.

Original: These energy based models can then be sampled to generate desired behavior.
Revised: These energy based models can then be sampled to generate desired behavior.

Original: Since [23] isa BC model capable of multi-modality, we compare against it as a baseline in Sec. 3.
Revised: Since [25] is a BC model capable of multi-modality, we compare against it as a baseline in Sec. 3.

Original: Transformers for control: With the stellar success of transformer models [72] in natural languageprocessing [17, 10] and computer vision [20], there has been significant interest in using transformermodels to learn behavior and control.
Revised: Transformers for control: With the stellar success of transformer models [78] in natural language processing [19, 11] and computer vision [22], there has been signiﬁcant interest in using transformer models to learn behavior and control.

Original: Among those, [11, 34] applies them to Reinforcement Learningand Offline Reinforcement Learning, respectively, while [12, 15, 48] use them for imitation learning.[15, 48] use transformers mostly to summarize historical visual context, while [12] relies on theirlong-term extrapolation abilities.
Revised: Among those, [12, 38] applies them to Reinforcement Learning and Ofﬂine Reinforcement Learning, respectively, while [13, 16, 53] use them for imitation learning. [16, 53] use transformers mostly to summarize historical visual context, while [13] relies on their long-term extrapolation abilities to collect human-in-the-loop demonstrations more efﬁciently.

Original: BeT is inspired by both of these use cases, as we use a transformerto summarize historical context while leveraging its generative abilities.
Revised: BeT is inspired by both of these use cases, as we use a transformer to summarize historical context while leveraging its generative abilities.

Original: Datasets for distributionally multi-modal data: Similar to computer vision [16, 42, 45] andnatural language processing [9, 63], there has been a recent interest in collecting behavior datasetsthat may aid in downstream behavior learning.
Revised: Datasets for distributionally multi-modal data: Similar to computer vision [18, 47, 50] and natural language processing [10, 69], there has been a recent interest in collecting behavior datasets that may aid in downstream behavior learning.

Original: Some of them are labeled with agent goals or rewardsfor downstream tasks [49, 24, 50], while others are more open ended [30, 46, 76] and come withoutreward or task labels.
Revised: Some of them are labeled with agent goals or rewards for downstream tasks [54, 26, 55], while others are more open ended [34, 51, 82] and come without reward or task labels.

Original: In our work, we focus towards the latter class.
Revised: In our work, we focus towards the latter class.

Original: The lack of labeled goal or rewardlabels in the second category implies that there is more multi-modality in the action distributionscompared to action distributions of goal or reward conditioned datasets, which is the same reason alot of work learning from multi-modal datasets try to learn a goal-conditioned model [30, 46, 15].
Revised: The lack of labeled goal or reward labels in the second category implies that there is more multi-modality in the action distributions compared to action distributions of goal or reward conditioned datasets, which is the same reason a lot of work learning from multi-modal datasets try to learn a goal-conditioned model [34, 51, 16].

Original: Finally, the lack of labelling requirements mean that the unlabelled datasets are cheaper to obtain,which should help BeT scale further in the future.
Revised: Finally, the lack of labelling requirements mean that the unlabelled datasets are cheaper to obtain, which should help BeT scale further in the future.

Original:  Discussion, limitations, and scope for future work
Revised:  Discussions

Original: In this work, we introduce Behavior Transformers (BeT), which uses a transformer-decoder basedbackbone with a discrete action mode predictor coupled with a continuous action offset correctorto model continuous actions sequences from open-ended, multi-modal demonstrations.
Revised: In this work, we introduce Behavior Transformers (BeT), which uses a transformer-decoder based backbone with a discrete action mode predictor coupled with a continuous action offset corrector to model continuous actions sequences from open-ended, multi-modal demonstrations.

Original: While
Revised: While

Original: BeT shows promise, the truly exciting use of it would be to learn diverse behavior from humandemonstrations or interactions in the real world.
Revised: BeT shows promise, the truly exciting use of it would be to learn diverse behavior from human demonstrations or interactions in the real world.

Original: In parallel, extracting a particular, unimodal behaviorpolicy from BeT during online interactions, either by distilling the model or by generating the right‘prompts’ [65], would make BeT tremendously useful as a prior for online Reinforcement Learning.Checklist
Revised: In parallel, extracting a particular, unimodal behavior policy from BeT during online interactions, either by distilling the model or by generating the right ‘prompts’ [71], would make BeT tremendously useful as a prior for online Reinforcement Learning.

Original: The checklist follows the references.

Original: Please read the checklist guidelines carefully for information onhow to answer these questions.

Original: For each question, change the default [TODO] to [Yes] , [No] , or[N/A] .

Original: You are strongly encouraged to include a justification to your answer , either by referencingthe appropriate section of your paper or providing a brief inline description.

Original: For example:

Original: • Did you include the license to the code and datasets?

Original: [Yes]

Original: • Did you include the license to the code and datasets?

Original: [No] The code and the data areproprietary.•

Original: Did you include the license to the code and datasets?

Original: [N/A] Please do not modify the questions and only use the provided macros for your answers. Note that the Checklist section does not count towards the page limit.
Revised: Acknowledgments

Original: In your paper, please delete this instructionsblock and only keep the Checklist section heading above along with the questions/answers below.

Original: For all authors...

Original: (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s

Original: contributions and scope?

Original: [Yes]

Original: (b) Did you describe the limitations of your work?

Original: [Yes] See Sec. 5.

Original: (c) Did you discuss any potential negative societal impacts of your work?

Original: [N/A]

Original: (d) Have you read the ethics review guidelines and ensured that your paper conforms to

Original: them?

Original: [Yes]

Original: If you are including theoretical results...

Original: (a) Did you state the full set of assumptions of all theoretical results?

Original: [N/A]

Original: (b) Did you include complete proofs of all theoretical results?

Original: [N/A]

Original: If you ran experiments...

Original: (a) Did you include the code, data, and instructions needed to reproduce the main experi-

Original: mental results (either in the supplemental material or as a URL)?

Original: [Yes] See abstract. (b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?
Revised: We thank Ben Evans, David Brandfonbrener, Abitha Thankaraj, Jyo Pari, and Anthony Chen for valuable feedback and discussions.

Original: [Yes] See appendix

Original: B

Original: (c) Did you report error bars (e.g., with respect to the random seed after running experi-

Original: ments multiple times)?

Original: [N/A]

Original: (d) Did you include the total amount of compute and the type of resources used (e.g., type

Original: of GPUs, internal cluster, or cloud provider)?

Original: [Yes] See Sec 3.5.

Original: If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...

Original: (a) If your work uses existing assets, did you cite the creators? [Yes] See linked code. (b) Did you mention the license of the assets? [Yes] See linked code.
Revised: This work was supported by grants from Honda, and ONR award numbers N00014-21-1-2404 and N00014-21-1-2758.

Original: (c) Did you include any new assets either in the supplemental material or as a URL?

Original: [No]

Original: (d) Did you discuss whether and how consent was obtained from people whose data you’re

Original: using/curating?

Original: [N/A]

Original: (e) Did you discuss whether the data you are using/curating contains personally identifiable

Original: information or offensive content?

Original: [N/A]

Original: If you used crowdsourcing or conducted research with human subjects...

Original: (a) Did you include the full text of instructions given to participants and screenshots, if

Original: applicable?

Original: [N/A]

Original: (b) Did you describe any potential participant risks, with links to Institutional Review

Original: Board (IRB) approvals, if applicable?

Original: [N/A]

Original: (c) Did you include the estimated hourly wage paid to participants and the total amount

Original: spent on participant compensation?

Original: [N/A]

Original: Appendix
Revised: Appendix

Original: Diverse, multi-modal behaviors generated by our models on different environment are best experienced and understood in a video.
Revised: Diverse, multi-modal behaviors generated by our models on different environment are best experienced and understood in a video.

Original: We invite you to visit https://submission0.github.io to see
Revised: We invite you to visit https://notmahi.github.io/bet to see

Original: BeT models in action.
Revised: BeT models in action.

Original: A Environment and Dataset Details
Revised: A Environment and Dataset Details

Original: Point mass environments: In the point mass environment, we have a simple point-mass agentwith two-dimensional observation and action spaces.
Revised: Point mass environments: In the point mass environment, we have a simple point-mass agent with two-dimensional observation and action spaces.

Original: The observation of the agent denotes the ( x, y )position of the agent, while the action sets the immediate (∆ x, ∆ y ) displacement of the agent in thenext timestep.
Revised: The observation of the agent denotes the ( x, y ) position of the agent, while the action sets the immediate (∆ x, ∆ y ) displacement of the agent in the next timestep.

Original: To show the effects of unimodal and multimodal behavioral cloning algorithms more cleanly, wealso add a “snapping" effect to the environment which moves the agent close to the nearest integercoordinates after each step.
Revised: To show the effects of unimodal and multimodal behavioral cloning algorithms more cleanly, we also add a “snapping" effect to the environment which moves the agent close to the nearest integer coordinates after each step.

Original: We generate random trajectories for each of our Multipath experiment datasets.1.
Revised: We generate random trajectories for each of our Multipath experiment datasets.

Original: In the first one (Fig. 2), our dataset has two modes, which are colored differently in thefigure based on the path taken at the fork.(a)
Revised: In the ﬁrst one (Fig. 2), our dataset has two modes, which are colored differently in the ﬁgure based on the path taken at the fork. (a)

Original: In the first set of demonstrations, the point mass follows the trajectory(1 , 2) , (2 , 2) , (2 , 3) , (2 , 4) , (3 , 4) , (4 , 4) , (4 , 3) , (4 , 2) , (5 , 2) .(b) In the second set of demonstrations, the point mass follows(1 , 2) , (2 , 2) , (2 , 1) , (2 , 0) , (3 , 0) , (4 , 0) , (4 , 1) , (4 , 2) , (5 , 2) .2.
Revised: In the ﬁrst set of demonstrations, the point mass follows the trajectory (1 , 2) , (2 , 2) , (2 , 3) , (2 , 4) , (3 , 4) , (4 , 4) , (4 , 3) , (4 , 2) , (5 , 2) . (b) In the second set of demonstrations, the point mass follows (1 , 2) , (2 , 2) , (2 , 1) , (2 , 0) , (3 , 0) , (4 , 0) , (4 , 1) , (4 , 2) , (5 , 2) .

Original: For the second Multipath environment (Fig. 5), there are three modes of demonstration,which are colored in the figure according to their first step direction.(a)
Revised: For the second Multipath environment (Fig. 5), there are three modes of demonstration, which are colored in the ﬁgure according to their ﬁrst step direction. (a)

Original: In the first set of demonstration, the point mass follows x = y from (0 , 0) to (8 , 8) with√ 2 size step increments.(b)
Revised: In the ﬁrst set of demonstration, the point mass follows x = y from (0 , 0) to (8 , 8) with √ 2 size step increments.

Original:  In the second set of demonstration, the point mass follows straight lines from (0 , 0) →(0 , 4) → (4 , 4) → (8 , 4) → (8 , 8) with step size 1 .(c)
Revised: (b) In the second set of demonstration, the point mass follows straight lines from (0 , 0) → (0 , 4) → (4 , 4) → (8 , 4) → (8 , 8) with step size 1 . (c)

Original: In the third set of demonstration, the point mass follows straight lines from (0 , 0) →(4 , 0) → (4 , 4)
Revised: In the third set of demonstration, the point mass follows straight lines from (0 , 0) → (4 , 0) → (4 , 4)

Original: → (4 , 8) → (8 , 8) with step size 1 .
Revised: → (4 , 8) → (8 , 8) with step size 1 .

Original: CARLA environment: We use the CARLA [19] self-driving environment to examine BeT performance in environments with high-dimensional observation spaces.
Revised: CARLA environment: We use the CARLA [21] self-driving environment to examine BeT performance in environments with high-dimensional observation spaces.

Original: CARLA uses the Unreal Engineto provide a photo-realistic driving simulation.
Revised: CARLA uses the Unreal Engine to provide a photo-realistic driving simulation.

Original: We create our environment on the Town04 map in
Revised: We create our environment on the Town04 map in

Original: CARLA 0.9.13.
Revised: CARLA 0.9.13.

Original: The observation space is 224 × 224 × 3 RGB images from the vehicle, which areprocessed by an ImageNet-pretrained, frozen ResNet-18 to a 512-dimensional real-valued vector.
Revised: The observation space is 224 × 224 × 3 RGB images from the vehicle, which are processed by an ImageNet-pretrained, frozen ResNet-18 to a 512-dimensional real-valued vector.

Original: The action space is [ − 1 , 1] 2 with an accelerator-brake axis and a steering axis.
Revised: The action space is [ − 1 , 1] 2 with an accelerator-brake axis and a steering axis.

Original: The dataset on this environment is collected with the built-in PID agent with minor tuning.
Revised: The dataset on this environment is collected with the built-in PID agent with minor tuning.

Original: We fixwaypoints in the trajectory that the demonstration agent needs to follow.
Revised: We ﬁx waypoints in the trajectory that the demonstration agent needs to follow.

Original: The waypoints fork aroundtwo central blocks: one set of trajectories thus go to the left, while another set of demonstrationtrajectories go to the right.
Revised: The waypoints fork around two central blocks: one set of trajectories thus go to the left, while another set of demonstration trajectories go to the right.

Original: For the demonstrations, we add some noise in the environment beforeexecuting an action so that there is some variation in the set of 100 total demonstrations that wecollect in the environment.
Revised: For the demonstrations, we add some noise in the environment before executing an action so that there is some variation in the set of 100 total demonstrations that we collect in the environment.

Original: Block-push environment: We use a simulated environment similar to Multimodal Push environment described in [23].
Revised: Block-push environment: We use a simulated environment similar to Multimodal Push environment described in [25].

Original: We take the environment implementation directly from the PyBullet [14] basedimplementation provided by Florence et al.
Revised: We take the environment implementation directly from the PyBullet [15] based implementation provided by Florence et al.

Original: [23] in https://github.com/google-research/ibc/tree/master/environments .
Revised: [25] in https://github.com/google-research/ ibc/tree/master/environments .

Original: In our environment, an XArm robot is situated in front of two blocks in a 0 .
Revised: In our environment, an XArm robot is situated in front of two blocks in a 0 .

Original: 75 × 1 plane.
Revised: 75 × 1 plane.

Original: On theplane there are also two square targets.
Revised: On the plane there are also two square targets.

Original: The goal of the agent is to push the blocks inside of thesquares.
Revised: The goal of the agent is to push the blocks inside of the squares.

Original: However, the exact order of the block being pushed, or the combination of which block ispushed in which square doesn’t matter.
Revised: However, the exact order of the block being pushed, or the combination of which block is pushed in which square doesn’t matter.

Original: A block is considered successfully pushed if the center of theblock is less than 0 .
Revised: A block is considered successfully pushed if the center of the block is less than 0 .

Original: away from a square.
Revised: away from a square.

Original: On initialization, the blocks’ positions are randomly shifted within a rectangle of side lengths(0 . 2 , 0 . 3) , while the squares are randomly shifted within a rectangle of size (0 . 01 , 0 . 015) .
Revised: On initialization, the blocks’ positions are randomly shifted within a rectangle of side lengths (0 . 2 , 0 . 3) , while the squares are randomly shifted within a rectangle of size (0 . 01 , 0 . 015) .

Original: Additionally, the blocks were rotated at an uniformly arbitrary angle, while the target squares were rotated atan angle between ( π 6 , − π 6 ) .
Revised: Additionally, the blocks were rotated at an uniformly arbitrary angle, while the target squares were rotated at an angle between ( π 6 , − π 6 ) .

Original: The demonstrations in this environments were collected with a hard-coded controller.
Revised: The demonstrations in this environments were collected with a hard-coded controller.

Original: There are twomodes of multimodality inherent in the controller generated demonstartions.
Revised: There are two modes of multimodality inherent in the controller generated demonstartions.

Original: The controller:1.
Revised: The controller:

Original: Selects a block to start pushing first,2.
Revised: Selects a block to start pushing ﬁrst, 2.

Original: At the same time, independently chooses a target for the block to be pushed into.3.
Revised: At the same time, independently chooses a target for the block to be pushed into.

Original: Once the first block is pushed to a target, it pushes the second block to the remaining target.
Revised: Once the ﬁrst block is pushed to a target, it pushes the second block to the remaining target.

Original: Thus combinatorially, the controller is capable of four different modes of behavior.
Revised: Thus combinatorially, the controller is capable of four different modes of behavior.

Original: There areadditional stochasticity in the controller behavior since there are many ways of pushing the sameblock into the same target.
Revised: There are additional stochasticity in the controller behavior since there are many ways of pushing the same block into the same target.

Original: The controller pushes the blocks to their targets following specific behavior primitives, such asmoving to origin position, moving to a place collinear with a block and its target, and making astraight motion from that position towards the target unless the block rotates too much from itsstarting position.
Revised: The controller pushes the blocks to their targets following speciﬁc behavior primitives, such as moving to origin position, moving to a place collinear with a block and its target, and making a straight motion from that position towards the target unless the block rotates too much from its starting position.

Original: Our models were trained on 1,000 demonstrations, all generated from the controller under the aboverandomized modes.
Revised: Our models were trained on 1,000 demonstrations, all generated from the controller under the above randomized modes.

Original: Franka kitchen environment: For the final set of experiments, we use the Franka Kitchen environment originally introduced in the Relay Policy Learning [30] paper.
Revised: Franka kitchen environment: For the ﬁnal set of experiments, we use the Franka Kitchen environment originally introduced in the Relay Policy Learning [34] paper.

Original: In that paper, the authorsintroduce a virtual kitchen environment where human participants in VR manipulated seven differentobjects in the kitchen: one kettle, one microwave, one sliding door, one hinged door, one light switch,and two burners.
Revised: In that paper, the authors introduce a virtual kitchen environment where human participants in VR manipulated seven different objects in the kitchen: one kettle, one microwave, one sliding door, one hinged door, one light switch, and two burners.

Original: In total, we use 566 demonstrations collected by the researchers in that paper, wherein each demonstration episode, each participant performed four manipulation task specified by theresearchers in advanced.
Revised: In total, we use 566 demonstrations collected by the researchers in that paper, where in each demonstration episode, each participant performed four manipulation task speciﬁed by the researchers in advanced.

Original: The manipulator agent in simulator is a Franka Emika Panda robot, which is controlled through a 9dimensional action space controlling the robot’s joint and end-effector position.
Revised: The manipulator agent in simulator is a Franka Emika Panda robot, which is controlled through a 9dimensional action space controlling the robot’s joint and end-effector position.

Original: The 60-dimensionalobservation space is split into two parts, the first 30 dimension contains information about thecurrent position of the interesting factors in the environment, while the last 30 dimensions containinformation about the goal of the demonstrator or the agent.
Revised: The 60-dimensional observation space is split into two parts, the ﬁrst 30 dimension contains information about the current position of the interesting factors in the environment, while the last 30 dimensions contain information about the goal of the demonstrator or the agent.

Original: Note that in our demonstrations and ourenvironments, we zero out the last 30 dimensions in all cases since we assume goal is not labelled inthe demonstrations and is not specified in the unconditioned rollouts of the model.
Revised: Note that in our demonstrations and our environments, we zero out the last 30 dimensions in all cases since we assume goal is not labelled in the demonstrations and is not speciﬁed in the unconditioned rollouts of the model.

Original: One thing to note that, while the D4RL [24] paper also has three versions of the dataset, we choseto use the original version of the collected data from the Relay Policy Learning [30] paper.
Revised: One thing to note that, while the D4RL [26] paper also has three versions of the dataset, we chose to use the original version of the collected data from the Relay Policy Learning [34] paper.

Original: Thatis because the relay policy learning dataset is not labeled with intended tasks of the participants orrewards, while the D4RL dataset is geared towards that.
Revised: That is because the relay policy learning dataset is not labeled with intended tasks of the participants or rewards, while the D4RL dataset is geared towards that.

Original: B Implementation Details and HyperparametersB.1 Baselines
Revised: B Implementation Details and Hyperparameters B.1 Baselines

Original: Multi-layer Perceptron with MSE For our MLP with MSE baselines, we trained fully connectedneural networks with optionally BatchNorm layers.
Revised: Multi-layer Perceptron with MSE For our MLP with MSE baselines, we trained fully connected neural networks with optionally BatchNorm layers.

Original: In each of our environment, we varied the depthand the width of the MLPs to fit them best according to the bias-variance trade-off, while trainingthem on 95% of the dataset and testing on the remaining 5% on the dataset in terms of MSE loss.
Revised: In each of our environment, we varied the depth and the width of the MLPs to ﬁt them best according to the bias-variance trade-off, while training them on 95% of the dataset and testing on the remaining 5% on the dataset in terms of MSE loss.

Original: Nearest Neighbor Nearest Neighbor is conceptually the simplest baseline we show in this paper.
Revised: Nearest Neighbor Nearest Neighbor is conceptually the simplest baseline we show in this paper.

Original: During training, our Nearest Neighbor model simply stores all the ( o, a ) pairs.
Revised: During training, our Nearest Neighbor model simply stores all the ( o, a ) pairs.

Original: During test time,given a query observation, o , we find the observation o ′ with the minimum Euclidean distance to thatin the representation space, and execute the associated action a ′ in the environment.
Revised: During test time, given a query observation, o , we ﬁnd the observation o (cid:48) with the minimum Euclidean distance to that in the representation space, and execute the associated action a (cid:48) in the environment.

Original: While it is a simple baseline, we show that it has a surprisingly high effectiveness in simple environments like CARLA, or dense environments like Kitchen where there is less of a chance in going
Revised: While it is a simple baseline, we show that it has a surprisingly high effectiveness in simple environments like CARLA, or dense environments like Kitchen where there is less of a chance in going

Original: OOD simply by executing seen actions.
Revised: OOD simply by executing seen actions.

Original: On the other hand, in environments like Block-push wherethe model needs to interpolate or extrapolate more, the NN model fails more.k-Nearest Neighbor with Locally Weighted Regression
Revised: On the other hand, in environments like Block-push where the model needs to interpolate or extrapolate more, the NN model fails more. k-Nearest Neighbor with Locally Weighted Regression

Original: A slightly more robust version of NN forregression problems, k-NN with locally weighted regression or LWR, is the next baseline we use.
Revised: A slightly more robust version of NN for regression problems, k-NN with locally weighted regression or LWR, is the next baseline we use.

Original: Inthis baseline, we take the k-nearest neighbors (in all our cases, 5 ) in the observation representationspace, and take a weighted average of their associated actions.
Revised: In this baseline, we take the k-nearest neighbors (in all our cases, 5 ) in the observation representation space, and take a weighted average of their associated actions.

Original: The weighting is based on the negativeexponent of the distance, or namely, exp −|| o − o ′ || , as seen in [55].
Revised: The weighting is based on the negative exponent of the distance, or namely, exp −|| o − o (cid:48) || , as seen in [61].

Original: This model is better than simple
Revised: This model is better than simple

Original: Nearest Neighbors in interpolations, and thus we see a higher success in the Kitchen environment.
Revised: Nearest Neighbors in interpolations, and thus we see a higher success in the Kitchen environment.

Original: Continuous Generative Model: VAE with Gaussian Prior
Revised: Continuous Generative Model: VAE with Gaussian Prior

Original: Following prior works[59], we usevariational auto-encoders (VAE) for encoding and decoding sequences of actions into a smaller latentspace.
Revised: Following prior works[65], we use variational auto-encoders (VAE) for encoding and decoding sequences of actions into a smaller latent space.

Original: The VAE here learns to compress a sequence of T = 10 actions into a single latent variable zof 10 dimensions.
Revised: The VAE here learns to compress a sequence of T = 10 actions into a single latent variable z of 10 dimensions.

Original: The hyperparameters for training the VAE has been taken directly from Pertschet al.
Revised: The hyperparameters for training the VAE has been taken directly from Pertsch et al.

Original: Concurrently with training the VAE, we train a state-conditioned latent prior model that tries to predict
Revised: Concurrently with training the VAE, we train a state-conditioned latent prior model that tries to predict

Original: P ( z | o ) .
Revised: P ( z | o ) .

Original: This latent generator produces a vector of µ and σ which is sampled to find latent z , andwe feed a Gaussian distributed variable z back into the decoder network where the action sequenceis reconstructed.
Revised: This latent generator produces a vector of µ and σ which is sampled to ﬁnd latent z , and we feed a Gaussian distributed variable z back into the decoder network where the action sequence is reconstructed.

Original: For the current observation o t , sequence of reconstructed actions a t , · · · , a t +9 areperformed in a simulated environment.
Revised: For the current observation o t , sequence of reconstructed actions a t , · · · , a t +9 are performed in a simulated environment.

Original: The design choices of this algorithm has been heavily inspired by [59].
Revised: The design choices of this algorithm has been heavily inspired by [65].

Original: Although this model showspromise in theory, we found in practice that unconditional rollout from this model is not verysuccessful.
Revised: Although this model shows promise in theory, we found in practice that unconditional rollout from this model is not very successful.

Original: We believe the shortcoming is a result of random sampling from the z space that does nottake into account the recently executed actions, and using a single-mode Gaussian as the state priorsimilar to [59], and thus this baseline is only slightly better than the MLP-MSE model.
Revised: We believe the shortcoming is a result of random sampling from the z space that does not take into account the recently executed actions, and using a single-mode Gaussian as the state prior similar to [65], and thus this baseline is only slightly better than the MLP-MSE model.

Original: Continuous Generative Models: Normalizing Flow with and without Prior Similar to Singhet al.
Revised: Continuous Generative Models: Normalizing Flow with and without Prior Similar to Singh et al.

Original: [70], we use a Normalizing Flow [18] based generative model.
Revised: [76], we use a Normalizing Flow [20] based generative model.

Original: We follow the architecturalchoices and the hyperparameters from [70] in our baseline implementation.
Revised: We follow the architectural choices and the hyperparameters from [76] in our baseline implementation.

Original: Our observation-conditioned Flow model is trained on the distribution P ( a | o ) to continuouslytransform it into an identity Gaussian distribution of the same dimensions as a .
Revised: Our observation-conditioned Flow model is trained on the distribution P ( a | o ) to continuously transform it into an identity Gaussian distribution of the same dimensions as a .

Original: To find a better priorthan simply an identity Gaussian, we also trained a prior model that generates µ, σ of a Gaussiandistribution given the observation o .
Revised: To ﬁnd a better prior than simply an identity Gaussian, we also trained a prior model that generates µ, σ of a Gaussian distribution given the observation o .

Original: We found that the prior improves the quality of the rollouts,however slightly.
Revised: We found that the prior improves the quality of the rollouts, however slightly.

Original: We believe the under-performance of these continuous generative approaches were based on twomajor problems.
Revised: We believe the under-performance of these continuous generative approaches were based on two major problems.

Original: One is that they fail to take historical context in concern, and by being a continuousdistribution, returned less likely actions that led to more rollouts going OOD.
Revised: One is that they fail to take historical context in concern, and by being a continuous distribution, returned less likely actions that led to more rollouts going OOD.

Original: Second, they weredesigned with a focus of making RL approachable by compressing the action space, which requireshaving a prior that is not so strict.
Revised: Second, they were designed with a focus of making RL approachable by compressing the action space, which requires having a prior that is not so strict.

Original: However, most of BeT’s performance comes from having a strongprior over the actions, which is only augmented by the action offset prediction.
Revised: However, most of BeT’s performance comes from having a strong prior over the actions, which is only augmented by the action offset prediction.

Original: Implicit Behavioral Cloning Implicit Behavioral Cloning (IBC) [23] takes a different approach inbehavioral cloning, where instead of learning a model f ( o ) :
Revised: Implicit Behavioral Cloning Implicit Behavioral Cloning (IBC) [25] takes a different approach in behavioral cloning, where instead of learning a model f ( o ) :

Original: = a , we learn an energy based model
Revised: = a , we learn an energy based model

Original: E ( o, a ) where the intended action a at any observation is defined as arg min a E ( o, a ) .
Revised: E ( o, a ) where the intended action a at any observation is deﬁned as arg min a E ( o, a ) .

Original: While thissuffers from all the classic issues of training an EBM, like higher sample complexity and highercomplexity in sampling, IBC models have been shown to have higher success in learning multi-modaland discontinuous actions.
Revised: While this suffers from all the classic issues of training an EBM, like higher sample complexity and higher complexity in sampling, IBC models have been shown to have higher success in learning multi-modal and discontinuous actions.

Original: As a baseline, we use the official implementation provided in https://github.com/google-research/ibc For the CARLA environment, we use equivalent hyperparameters from the“pushing from pixels” hyperparameters.
Revised: As a baseline, we use the ofﬁcial implementation provided in https://github.com/ google-research/ibc For the CARLA environment, we use equivalent hyperparameters from the “pushing from pixels” hyperparameters.

Original: For the Block-pushing environment, we use the “pushingfrom states” hyperparameters.
Revised: For the Block-pushing environment, we use the “pushing from states” hyperparameters.

Original: Finally, for the Kitchen environment, we use the “D4RL kitchen”hyperparameters.
Revised: Finally, for the Kitchen environment, we use the “D4RL kitchen” hyperparameters.

Original: While IBC is our strongest baseline, in our experience it is also one that is quite easy to overfit toour datasets.
Revised: While IBC is our strongest baseline, in our experience it is also one that is quite easy to overﬁt to our datasets.

Original: As a result, we monitored test performance over the training and had to employ earlystopping for both the CARLA and the Block-pushing tasks.
Revised: As a result, we monitored test performance over the training and had to employ early stopping for both the CARLA and the Block-pushing tasks.

Original: B.2 Algorithm Details
Revised: B.2 Algorithm Details

Original: Loss function details:
Revised: Loss function details:

Original: In this paper, we use two loss functions that are inspired by practices incomputer vision, in particular object detection.
Revised: In this paper, we use two loss functions that are inspired by practices in computer vision, in particular object detection.

Original: The first of them is the Focal loss [43], and the secondone is the Multi-task loss [26].
Revised: The ﬁrst of them is the Focal loss [48], and the second one is the Multi-task loss [30].

Original: The Focal loss is a simple modification over the cross entropy loss.
Revised: The Focal loss is a simple modiﬁcation over the cross entropy loss.

Original: While the normal cross entropyloss for binary classification can be thought of L ce ( p t ) = − log( p t ) , the Focal loss adds a term(1 − p t ) γ to this, to make the new loss
Revised: While the normal cross entropy loss for binary classiﬁcation can be thought of L ce ( p t ) = − log( p t ) , the Focal loss adds a term (1 − p t ) γ to this, to make the new loss

Original: [Equation]
Revised: [Equation]

Original: This loss has the interesting property that its gradient is more steep for smaller values of p t , whileflatter for larger values of p t .
Revised: This loss has the interesting property that its gradient is more steep for smaller values of p t , while ﬂatter for larger values of p t .

Original: Thus, it penalizes and changes the model more for making errorsin the low-probability classes, while is more lenient about making errors in the high probabilityclasses.
Revised: Thus, it penalizes and changes the model more for making errors in the low-probability classes, while is more lenient about making errors in the high probability classes.

Original: Using this error in the object detection world has helped with class imbalance betweendifferent classes, and here it helps BeT learn to predict different k -means from the dataset even iftheir appearance in the dataset is not completely balanced.
Revised: Using this error in the object detection world has helped with class imbalance between different classes, and here it helps BeT learn to predict different k -means from the dataset even if their appearance in the dataset is not completely balanced.

Original: For the multi-task loss, we use the formulation
Revised: For the multi-task loss, we use the formulation

Original: [Equation]
Revised: [Equation]

Original: This helps us penalize only the offset for the ground truth class, thus making sure the MinGPT is nottrying to predict the right action offset through all classes and instead only trying to predict the actionoffset through the right class.
Revised: This helps us penalize only the offset for the ground truth class, thus making sure the MinGPT is not trying to predict the right action offset through all classes and instead only trying to predict the action offset through the right class.

Original: In practice, we optimize the combined loss, L focal + α L mt while α is a hyperparameter that justmakes sure at initialization the two losses are of the same order of magnitude.
Revised: In practice, we optimize the combined loss, L focal + α L mt while α is a hyperparameter that just makes sure at initialization the two losses are of the same order of magnitude.

Original: Compute details: All of our code was run in a single NVIDIA RTX 3080 GPU for state-basedenvironments and RTX 8000 for image-based environments.
Revised: Compute details: All of our code was run in a single NVIDIA RTX 3080 GPU for state-based environments and RTX 8000 for image-based environments.

Original: Performance measurement details: We measured the performance reported in the Section 3.5 inan NVIDIA RTX 3080 machine with AMD Threadripper 5950x CPUs.
Revised: Performance measurement details: We measured the performance reported in the Section 3.5 in an NVIDIA RTX 3080 machine with AMD Threadripper 5950x CPUs.

Original: We took the average overthree runs to minimize inter-run variances, and measured wall-clock time to report in the paper.
Revised: We took the average over three runs to minimize inter-run variances, and measured wall-clock time to report in the paper.

Original: Hyperparameters list: We present the BeT hyperparameters in Table 4 below:
Revised: Hyperparameters list: We present the BeT hyperparameters in Table 4 below:

Original: [Table]
Revised: [Table]

Original: However, we have found that as long as the model does not overfit, a wide range of parameters allyield favorable results for BeT; thus, this table should be taken as reference values for reproducingour results rather than the only parameter sets that work.
Revised: However, we have found that as long as the model does not overﬁt, a wide range of parameters all yield favorable results for BeT; thus, this table should be taken as reference values for reproducing our results rather than the only parameter sets that work.

Original: Apart from that, we have some hyperparameters that are shared across all BeT experiments.
Revised: Apart from that, we have some hyperparameters that are shared across all BeT experiments.

Original: They arereproduced in Table 5.
Revised: They are reproduced in Table 5.

Original: B.3 Pseudocode
Revised: B.3 Pseudocode

Original: See the pseudocode described on Algorithm 1.
Revised: See the pseudocode described on Algorithm 1.

Original: [Table]
Revised: [Table]

Original: Input: Dataset ( o t,i , a t,i ) t,i for 0 ≤ i ≤ number of demonstrations, 0 ≤ t ≤ maximum episode lengths, intended number of clusters k and context history length
Revised: Input: Dataset ( o t,i , a t,i ) t,i for 0 ≤ i ≤ number of demonstrations, 0 ≤ t ≤ maximum episode lengths, intended number of clusters k and context history length

Original: h .
Revised: h .

Original: Initialize: θ M the parameters for MinGPT, { A i } ki =1 cluster centers randomly in the action space.
Revised: Initialize: θ M the parameters for MinGPT, { A i } ki =1 cluster centers randomly in the action space.

Original: Learn k-means encoder/decoder:
Revised: Learn k-means encoder/decoder:

Original: Using all possible a t,i , learn the k cluster centers using the k means algorithm.
Revised: Using all possible a t,i , learn the k cluster centers using the k means algorithm.

Original: Set { A i } ki =1 as the learned cluster centers.
Revised: Set { A i } ki =1 as the learned cluster centers.

Original: Define functions: 
Revised: Deﬁne functions: (cid:98) a (cid:99) :=

Original: [Equation]
Revised: arg min ki =1 || a − A i || (cid:104) a (cid:105) : = a − (cid:98) a (cid:99)

Original: Enc [Equation]
Revised: Enc ( a ) :=

Original:  Dec [Equation]
Revised: ( (cid:98) a (cid:99) , (cid:104) a (cid:105) ) Dec ( (cid:98) a (cid:99) , (cid:104) a (cid:105) ) :

Original: [Equation] A ⌊ [Equation] + ⟨ a ⟩
Revised: = A (cid:98) a (cid:99) + (cid:104) a (cid:105)

Original: Train MinGPT trunk of BeT:
Revised: Train MinGPT trunk of BeT:

Original: while Not converged do
Revised: while Not converged do

Original: Sample trajectory subsequence ( o t , a t ) , · · · , ( o t + h − 1 , a t + h − 1 ) from the dataset.
Revised: Sample trajectory subsequence ( o t , a t ) , · · · , ( o t + h − 1 , a t + h − 1 ) from the dataset.

Original: Feed in the observations ( o t , o t +1 , · · · , o t + h − 1 ) into the MinGPT.
Revised: Feed in the observations ( o t , o t +1 , · · · , o t + h − 1 ) into the MinGPT.

Original: Get categorical distribution probabilities p τ,c for t ≤ τ ≤ t + h − 1 , 1 ≤ c ≤ k .
Revised: Get categorical distribution probabilities p τ,c for t ≤ τ ≤ t + h − 1 , 1 ≤ c ≤ k .

Original: Compute focal loss L ce of p τ,c against ground truth class ⌊ a τ ⌋ , for all τ, c .
Revised: Compute focal loss L ce of p τ,c against ground truth class (cid:98) a τ (cid:99) , for all τ, c .

Original: Get the residual action offset per class, ⟨ a τ,c ⟩ , for all τ, c from MinGPT.
Revised: Get the residual action offset per class, (cid:104) a τ,c (cid:105) , for all τ, c from MinGPT.

Original: Calculate the multi-task loss, L mt , against true class predicted offset, (cid:80) τ ||⟨ a τ, ⌊ a τ ⌋ ⟩ − ⟨ a τ ⟩||
Revised: Calculate the multi-task loss, L mt , against true class predicted offset, (cid:80) τ [Table] [Table]

Original: Backprop using the normalized loss, L ce + α L mt where α makes the losses of equal magnitude.
Revised: |Backprop using the normalized loss, L ce + α L mt where α makes the losses of equal magnitude.

Original: Running on the environment: while Episode not completed do
Revised: Running on the environment: while Episode not completed do

Original: Stack the last h observations in the environment, ( o t , o t +1 , · · · , o t + h − 1 ) and feed into MinGPT.
Revised: Stack the last h observations in the environment, ( o t , o t +1 , · · · , o t + h − 1 ) and feed into MinGPT.

Original: Get categorical probabilities p τ,c for t ≤ τ ≤ t + h − 1 , 1 ≤ c ≤ k from the MinGPT.
Revised: Get categorical probabilities p τ,c for t ≤ τ ≤ t + h − 1 , 1 ≤ c ≤ k from the MinGPT.

Original: Sample a class c from p t + h − 1 ,c for 1 ≤ c ≤ k .
Revised: Sample a class c from p t + h − 1 ,c for 1 ≤ c ≤ k .

Original: Get the associated action offset, ⟨ a t + h − 1 ,c ⟩ from the MinGPT.
Revised: Get the associated action offset, (cid:104) a t + h − 1 ,c (cid:105) from the MinGPT.

Original: Decode into full continuous action, a t + h − 1 :=
Revised: Decode into full continuous action, a t + h − 1 :=

Original: Dec ( c, ⟨ a t + h − 1 ,c ⟩ ) Execute decoded action a t + h − 1 into environment.
Revised: Dec ( c, (cid:104) a t + h − 1 ,c (cid:105) ) Execute decoded action a t + h − 1 into environment.

Original: B.4 Architecture and Implementation
Revised: B.4 Architecture and Implementation

Original: For our implementation, we used the MinGPT [35] repository almost as-is.
Revised: For our implementation, we used the MinGPT [40] repository almost as-is.

Original: We modified the inputtoken conversion layer to a linear projection layer to handle our continuous, instead of discrete, inputs.
Revised: We modiﬁed the input token conversion layer to a linear projection layer to handle our continuous, instead of discrete, inputs.

Original: Apart from that, we followed the MinGPT architecture quite exclusively, with successive attentionlayers with a number of attention head and embedding dimensions.
Revised: Apart from that, we followed the MinGPT architecture quite exclusively, with successive attention layers with a number of attention head and embedding dimensions.

Original: Between the layers, we useddropout regularization same as [35].
Revised: Between the layers, we used dropout regularization same as [40].

Original: For the smallest tasks, like point-mass environments, we used models with approximately 10 4parameters, which went up to around 10 6 for Kitchen environments.
Revised: For the smallest tasks, like point-mass environments, we used models with approximately 10parameters, which went up to around 10 6 for Kitchen environments.

Original: C Ablation studies
Revised: C Ablation studies

Original: [Figure]

Original: In this section, we provide more details about the ablation studies presented in the main paper, aswell as present detailed plots of our ablation studies that compare different versions of the BeT[Figure] architecture.[Figure]
Revised: In this section, we provide more details about the ablation studies presented in the main paper, as well as present detailed plots of our ablation studies that compare different versions of the BeT architecture.

Original: C.1 Ablating historical context
Revised: C.1 Ablating historical context

Original: [Figure]

Original: One of the reasons why we used transformer-based generative networks in our work is because ofour hypothesis that having historical context helps our model learn better behavioral cloning.
Revised: One of the reasons why we used transformer-based generative networks in our work is because of our hypothesis that having historical context helps our model learn better behavioral cloning.

Original: Ourexperiments are performed by using the same model and simply providing sequences of length oneon training and test time.
Revised: Our experiments are performed by using the same model and simply providing sequences of length one on training and test time.

Original: As we can see on Sec. 3.5, having some historical context helps our modellearn much better.
Revised: As we can see on Sec. 3.5, having some historical context helps our model learn much better.

Original: C.2 Ablating the number of discrete bin centers, k
Revised: C.2 Ablating the number of discrete bin centers, k

Original: [Figure]

Original: Since BeT is trained with a sum of focal loss for the binning head and MSE loss for the offset head,the number of cluster centers present a trade-off in the architecture.
Revised: Since BeT is trained with a sum of focal loss for the binning head and MSE loss for the offset head, the number of cluster centers present a trade-off in the architecture.

Original: Concretely, as the number of binsgo up, the log-likelihood loss goes up but the MSE loss goes down.
Revised: Concretely, as the number of bins go up, the log-likelihood loss goes up but the MSE loss goes down.

Original: In Sec. 3.5, we showed that usingonly one bin ( k = 1 ) decreases the performance level of BeT.
Revised: In Sec. 3.5, we showed that using only one bin ( k = 1 ) decreases the performance level of BeT.

Original: [Figure]

Original: In this section, we present the plot of the variation in performance as k value changes.
Revised: In this section, we present the plot of the variation in performance as k value changes.

Original: [Figure]

Original: C.3 Ablating the core model in the architecture
Revised: C.3 Ablating the core model in the architecture

Original: [Figure] To ablate the core MinGPT transformer model in the architecture, we  replace it  with a fully-connected MLP network.
Revised: To ablate the core MinGPT transformer model in the architecture, we run three ablations, where we replace it respectively with a fully-connected multi-layer perceptron (MLP) network, a temporal convolution network, and an LSTM-based recurrent neural network.

Original:  Since generally MLP networks are not capable of taking in historical context inconsideration, we instead stack the last t frames of observation to pass into the MLP network.
Revised: Multi-Layer Perceptrons: Since generally MLP networks are not capable of taking in historical context in consideration, we instead stack the last t frames of observation to pass into the MLP network.

Original: Nearthe beginning of a trajectory, the stack of observation is zero-padded to t frames.
Revised: Near the beginning of a trajectory, the stack of observation is zero-padded to t frames.

Original: For the intermediatelayers in the MLP, we keep the same width and the number of layers as the corresponding MinGPT.
Revised: For the intermediate layers in the MLP, we keep the same width and the number of layers as the corresponding MinGPT.

Revised: Temporal Convolution: Convolutions over the sequence length has been used in numerous prior works [59, 39, 17, 28, 5] for sequence modeling.

Revised: As a baseline, we implement such temporal convolutional network to replace our MinGPT-based trunk.

Revised: We perform a temporal convolution over the same period of history that is provided to our transformer models.

Revised: We found that the performance of the temporal convolution models are constantly lower than our MinGPT based models.

Revised: However, temporal convolutional networks are easier to ﬁt on our data compared to RNNs.

Revised: LSTM-based RNN: Recurrent neural networks (RNNs) were the previous state-of-the-art for sequence modeling before transformer-based models.

Revised: In this work, we compare against an Longshort term memory (LSTM) [29] based RNN instead of a transformer based trunk.

Revised: We ﬁnd that even with sufﬁcient model capacity, the RNN based model took signiﬁcantly longer than our MinGPT model to ﬁt the same dataset.

Revised: Moreover, the quality of ﬁt was worse, both in training and test time.

Revised: Finally, in open-ended rollouts, this performance downgrade is reﬂected in a far lower success rate for completing tasks in the environment (Table. 3).