Efficient Planning in a Compact Latent Action Space
Authors: zhengyao jiang
Tianjun Zhang
Michael Janner
Yueying Li
Tim Rocktäschel
Edward Grefenstette
Yuandong Tian
Venue: ICLR 2023 poster (2023)
Keywords: Model-based RL
Planning
Sequence Modelling RL
Generative Model
Offline Reinforcement Learning

Abstract: Planning-based reinforcement learning has shown strong performance in tasks in discrete and low-dimensional continuous action spaces. However, planning usually brings significant computational overhead for decision making, so scaling such methods to high-dimensional action spaces remains challenging. To advance efficient planning for high-dimensional continuous control, we propose Trajectory Autoencoding Planner (TAP), which learns low-dimensional latent action codes with a state-conditional VQ-VAE. The decoder of the VQ-VAE thus serves as a novel dynamics model that takes latent actions and current state as input and reconstructs long-horizon trajectories. During inference time, given a starting state, TAP searches over discrete latent actions to find trajectories that have both high probability under the training distribution and high predicted cumulative reward. Empirical evaluation in the offline RL setting demonstrates low decision latency which is indifferent to the growing raw action dimensionality. For Adroit robotic hand manipulation tasks with high-dimensional continuous action space, TAP surpasses existing model-based methods by a large margin and also beats strong model-free actor-critic baselines.

Published as a conference paper at ICLR 2023
EFFICIENT PLANNING IN A COMPACT LATENT ACTION SPACE
Zhengyao Jiang'
Tianjun Zhang²
Michael Janner*
Yueying Li
Tim Rocktäschel!
Edward Grefenstette
Yuandong Tian University of California, Berkeley
University College London *Cohere Meta AI (FAIR)
*Cornell University
ABSTRACT
Planning-based reinforcement learning has shown strong performance in tasks in discrete and low-dimensional continuous action spaces. However, planning usu- ally brings significant computational overhead for decision-making and scaling such methods to nigh-dimensiooal action spaces remains challenging. To advance efficient planning for high-dimensional continuous control, we propose Trajectory Autoencoding Planner (TAP), which learns low-dimensional latent action codes with conditionl VQ- TQ-V The decoder of the VQ- TQ-V thus serves a novel dynamics model that takes latent actions and current state as input and reconstructs long-horizon trajectories. During inference time, given a starting state, TAP searches over discrete latent actions to find trajectories that have both high probability under the training distribution and high predicted cumulative re- ward. Empirical evaluation in the offine RI. setting demonstrates low decision latency which is indifferent the growing raw action dimensionality. For Adroit robotic hand manipulation tasks with high-dimensional continuous action space, TAP surpasses existing model-based methods by a large margin and also beats strong model-free actor-critic baselines.
INTRODUCTION
Planning-basee based reinforcement learning (RL) methods have shown strong performance on board games (Silver et al., 2018; Schrittwieser et al., 2020), video games (Schritwieser al., 2020; Ye et al., 2021) and low-dimensional continuous control (Janner et al., 2021). Planning convention- ally occurs in the action space of the Markov Decision Process (MDP), by rolling-out future trajectories with a dynamics model of the environment, which is either predefined or learned.
While such a planning procedure is intuitive, planning in action space can be inefficient and inflexible. Firstly, the optimal plan in a high-dimensional raw action space can be difficult to find. Even if the optimizer is powerful enough to find the optimal plan, it still difficult make sure the learned model is accurate in the whole raw action space. In such cases, the planner can exploit the weakness of the model and lead to over-optimistie planning. Secondly, planning in raw action space means the planning procedure is tied to the temporal structure of the environment. However, human planning much more flexible, for example, humans can introduce temporal abstractions and plan with high-level actions; humans can plan backwards from the goal; the plan can also start from high-level outline and get refined step-by-step. The limitations of raw action space planning cause slow decision speeds, which hamper adoption in real-time control.
In this paper, we propose the Trajectory Autoencoding Planner (TAP), which learns a latent action space and latent-action model from offline data. A latent-action model takes state $1 and latent actions z as input and predicts a segment of future trajectories T = (a1, T1, R1, 82, 02, T2, R2....) This latent action space be much smaller than the raw action space since it only captures plausi- ble trajectories in the dataset, preventing ant-of-distribution actions. Furthermore, the latent action decouples the planning from the original temporal structure of MDP. This enables the model, for example, to predict multiple steps of future trajectories with a single latent action.
*Correspond to z.jiang @cs w......l.ac..k.. The webpage is at: sites. aites. google. www.//////ewaateeeneee Source code available g thub. com www.//hengaa...iee....o..oo http///aag/////lllaaa.aa

---

Published as a conference paper at ICLR 2023
301
3.0
TAPICOL
Transformer
plaminn time
2.5
TAPITT
[F]. M)-
plais))
2.0
3.5
Encader
C22
C1
64
Ca
C2
10
reladive
5.0
Decoder
0.5
aci
2,
0.0
a Ppyrre)
dimensionality
dimenstonality
(a) TAP Modelling
(b) Decision Latency
(c) Relative Performance
Figure (a) gives an overview of TAP modelling. where blocks represent the latent actions. (b) shows decision time growth with the dimensionality D. Tests are done on single GPU. The number of planning steps for (b) is 15 and both models apply a beam search with a beam width of 64 and expansion factor of 4. (c) shows the relative performance between TAP and baselines when dealing with tasks with increasing raw action dimensionalitics.
The model and latent action space are learned in an unsupervised manner. Given the current state, the encoder of TAP learned map the trajectories to sequence of discrete latent codes (and the decoder maps it back), using a state-conditioned Vector Quantised Variational AutoEncoder (VQ- VAEV. As illustrated in Figure I(a), the distribution of these latent codes then modelled autoregressively with a Transformer, again conditioned on the current state. During inference, instead of sampling the actions and next state sequentially, TAP samples latent codes, reconstructs the trajectory via the decoder, and executes the first action in the trajectory with the highest objective score. These latent codes of the VQ-VAE are thus latent actions and the conditioniii decoder acts as a latent-action dynamics model. In practice, TAP uses single discrete latent variable to model multiple (L 3) steps of transitions, creating compact latent action space for downstream planning. Planning in this compact action space reduces the decision latency and makes high-ding -igh-dimensional planning easier. In addition, reconstructing the entire trajectories after all latent codes are sampled also helps alleviate the compound errors of step-by-step rollouts.
We evaluate TAP extensively in the offline RL setting. Our results on low-dimensional locomotion control tasks show that TAP enjoys competitive performance actor-critic, and sequence modelling baselines. On tasks with higher dimensionality, TAP not only surpasses model-based methods like MOPO (Yu et al., 2020) Trajectory Transformer(TT) (Janner et al., 2021) ) but also significantly outperforms strong model-free ones (e.g., CQL (Kumar et al., 2020)and IQL (Kostrikov al., 2022)). In Figure ((c). we show how the relative performance between TAP and basclines changes according to the dimensionality of the action space. One can see that the advantage of TAP starts to pronounce when the dimensionality of raw actions grows and the margin becomes large for high dimensionality, especially when compared to the model-based method TT. This can be explained by the innate difficulty of policy optimization in high-dimensional raw action space, which is avoided by TAP since its planning happens in low-dimensional discrete latent space. At the same time, sampling and planning of TAP are significantly faster than prior work that also uses Transformer as dynamics model: the decision time of TAP meets the requirement of deployment on a real robot (20Hz) (Reed et al., 2022), while TT is much slower and the latency grows along the state-action dimensionality in Figure 1(b).
strong model-based, model-free
2 BACKGROUND
2.1 VECTOR QUANTISED- VARIATIONAL AUTOENCODER
The Vector Quantised Variational Autoencoder (VQ-VAE) (van den Oord et al., 2017) is generative model designed based on the Variational Autoencoder (VAE). There are three components in VQ- VAEs: 1) an encoder network mapping inputs collection of discrete latent codes; decoder that maps latent codes to reconstructed inputs; 3) a learned prior distribution over latent variables. With the prior and the decoder, we can draw samples from the VQ-VAE.
2

---

Published as a conference paper at ICLR 2023
In order to represent the input with discrete latent variables, VQ- (Q-VAE employs an approach inspired by vector quantization (VQ). The encoder outputs continuous vectors z; which are not imput to the decoder directly but used to query vector in codebook R R with K embedding vectors es E IR k E 1, 1,2.3.... K. The query is done by simply finding the nearest neighbour from K embedding vectors, and the resultant vectors
2; = ek; where argming ||= - e,ll2
are imput to the decoder. The autoencoder is trained by jointly minimizing the reconstruction error and the distances ||=, exll2 and ||z ex The gradients copied directly through the bottle- neck during backpropagation. A PixelCNN (van den Oord et al., 2016) is used to model the prior distribution of the latent codes autoregressively.
3
METHOD
To enable flexible planning in learned latent action space, we propose Trajectory Autoencoding Planner (TAP) based on the VQ- this section, we will describe the latent action model framework and provide the model architecture details for TAP. Then we elaborate on how to plan with this latent-action model.
3.1 LATENT-ACTION MODEL AND PLANNING
Consider the following trajectory т of length T, sampled from an MDP with a fixed stochas- tic behaviour policy, consisting of sequence of states, actions, rewards and return-to-go R₂ Ei "rrrr; as proxies for future rewards:
F ( 41, F1, R1, 02 r2, R2,... ......... ST, rT, RT).
We model the conditional distribution of the trajectory p(r|s, 2) with a sequence of latent variables = (=1...., ). Assume the state and latent variables (s, =) can be deterministically mapped to a trajectory v so that p(T]s, =) 1(T = =))p((()) We refer to z as latent actions and p(z(s) as the latent policy. The mapping h(s,2) 2) from state and latent actions (s, 2) trajectory T is thus a latent-actionn acttin model. In deterministic MDP, the trajectory for an arbitrary h(s, 2) with p(2|s) 0 will be executable plan, namely, the trajectory
be recovered by following the action sequences, starting from state s. Therefore, we can optimize latent actions 2 in order to find an optimal plan.
Planning in the latent action space has two advantages compared to planning in the raw action space. Firstly, the latent-action model only captures possible actions in the support of the behaviour policy, allowing the latent action space to potentially be much smaller than the raw action space. For example, considering the policy is a mixture of policies, then the latent action space can be a discrete space with only X actions, no matter how nigh-dimensional the raw action space is. In addition, only allowing the in-distribution actions prevents the planner from exploiting the weakness of the model by querying the actions with high uncertainty whose values are most susceptible to overestimation. Secondly, the latent-action model allows the decoupling between planning and modelling. When planning in the raw action space, the time resolution of planning must the same as the predicted trajectories. In contrast, planning in the latent action space can be much more flexible, the transition. This property allows the length of the latent action sequence A to be smaller than the planning horizon T, leading more efficient planning.
the temporal structure
latent action does not have be tied to particular step of
3.2 LEARNING A LATENT-ACTION MODEL WITH VQ-VAES
A discrete action space makes dealing with the full distribution of actions easier and also allows a spectrum of advanced planning algorithms to be applied Silver et al. (2018); Tillmann et al. (1997). To learn a compact discrete latent action space, we propose Trajectory Autoencoding Planner (TAP) to model the trajectories with state-conditioned VQ-VAE. We treat æ; (st, a₁, Tt, Re) single token for the Transformer encoder and decoder. Both the autoencoder and the prior over latent variables conditioned on the first state 81 of the trajectory.
3

---

Published as a conference paper at ICLR 2023
(e mrrarrr (irr op do. (o
Encoderl
Sample Latents
Reconstruct Trejectories
Codabook
Pooling
Vector Quantisation
Prior
Decoder
es
Tiling
e2e15
onal
(2), Rsinm))
Act According to the Plan
[ aj. P1.
So. doy, Ro. (is)
Figure 2: Illustration of the training and test time inference process of TAP. The left-hand side shows the training process, highlighting the design the bottleneck. The right-hand side figure shows how we generate plans during the test time, with vanilla sampling.
Encoder g and decoder h. For the encoder, token æ; is processed by a causal Transformer, leading to T feature vectors. We then apply 1 I-dimensional max pooling with both kemel size and stride of I,, followed by linear layer. This results M = T/ vectors (2, ...; &M). corresponding to the number of discrete latent variables. After vector quantization in Equation (I), we get embedding vectors for the latent variables 22., MM)
For the decoder, the latent variables are then tiled L times match the number of input/output tokens T. For L = 3, this is written as Z2; (21, 21, 21, 22; 22, ...; 2T/3: 23/3; 33)))- We concatenate the state and embedding vector for the codes and apply linear projection to provide state information to the decoder. After adding the positional embedding. the decoder then reconstructs the trajectory f == (Ø), i2,.... LT), where i; = (8,, âe, ft, Re). Finally, to train the encoder/decoderr, the reconstruction loss is the mean squared error over the input trajectories {T} and reconstructed ones (+).
Prior distribution of latent codes. While the derivation of VQ- VAE assumes a uniform prior over latents, in practice, learning a parameterised prior over latents will usually lead better sample quality, Similarly, TAP learns prior policy p(z[8)) to inform sampling and planning. TAP uses Transformer to autoregressively model the distribution of the latent action codes given the initial state 81. The state information is also blended into the architecture by the concatenation of token embeddings. We denote the autoregressive prior policy as p(z, (2cc, = p(z $1 21, 22, ..., 21-
3.3 PLANNING
DISCRETE LATENT ACTION SPACE
We now describe how to plan with the learned TAP model. This includes how to evaluate trajectories and optimize latent actions.
Evaluation Criterion. Given the initial state and latent actions ($1, 2). we obtain a sample trajectory from decoder h, where a; := (&c, âr, fr, Rt). We score 7 with the following
f
function g:
g(81, 21, 22, 2M)
)= = En' S/ + Rr + a In (n...ll)).......xx11))...
(3)
The first part (in red) of the score function is the predicted return-to-go following the action se- quences in the decoded trajectory f. The second part (in blue) is a term that penalizes out-of- distribution plans, where the prior distribution p gives low probability. The hyperparameter is set to be larger than the maximum of the discounted returns to select plausible trajectory when the conditional probability of latent action sequences is lower than the threshold 8d When the probability is higher than the threshold, the objective will then encourage choosing a high-return plan.
4

---

Published as a conference paper at ICLR 2023
Vanilla Sampling The most straightforward way to search in the latent space of TAP is to sample according to the prior distribution and select the trajectory that gives the best return. Conditioned on current state $1,
sample N latent variable sequences autoregressively according to the prior model: (21, 22,...,,.M n ~ , 2M 81) where n. E 1, 2, ...., N. These latent variable sequences will then be turned into trajectories 7n that the objective function Equation (3). This naive approach works well when the overall length of the trajectory is small since the discrete latent space significantly smaller than the raw action space. However, its performance quickly degenerates with long planning horizons.
select the optimal one according to
Beam Search in the Latent Space We use causal Transformers for encoder, decoder and the prior policy, preventing the decoding depending the future latent actions. This allows us to partially decode the trajectory of length NL with first N latent variables. Therefore we can apply a planning algorithm like bean search for more efficient optimization by only expanding promising branches of the search tree. TAP beam search always keeps B best partial latent code sequences, and samples E new latent actions conditioned expansion factor. The pseudocode of the TAP beam search is shown in Algorithm in the Appendix. In practice, find TAP with beam search is usually computationally more efficient because it can find better trajectories with fewer queries to the prior model and decoder.
the partial codes. Here B is the beam width and E is the
4 EXPERIMENTS
The empirical evaluation of TAP consists of three sets of tasks from D4RL (Fu et al., 2020): gym locomotion control, AntMaze, and Adroit. We compare TAP to number of prior offline RL algo- rithms, including both model-free actor-critic methods (Kumar et al., 2020; Kostrikov et al., 2022) and model-based approaches (Kidambi et al., 2020; Yu et al., 2020; Lu et al., 2022). Our work is conceptually the most related to the Trajectory Transformer (TT; Janner et al. 2021), model-based planning method that predicts and plans in the raw state and action space, so this baseline serves as our main point of comparison. Gym locomotion tasks serve as proof of concept in the low- dimensional domain test if TAP can accurately reconstruct trajectories for use in decision-making and control. We then test TAP on Adroit, which has a high state and action dimensionality that causes TT to struggle because of its autoregressive modelling of long token sequences. Finally, we also test TAP on AntMaze, a sparse-reward continuous-come control problem. TAP achieves similar performance as TT on AntMaze, surpassing model-free methods, where the details can be found in Appendix F.
4.1 SETUP FOR TRAJECTORY AUTOENCODING PLANNER
We keep most of the design decisions the same as in TT in order to show the advantage of the latent- action model. Namely, (1) no bootstrapped value estimation is used; (2) the trajectory treated as a sequence and the latent structure of the model does not follow the structure of the MDP; (3) beam search is used for planning: and (4) the model architectures for the encoder, decoder, and prior policy are Transformers.
As for the TAP-specific hyperparameters: we set the number of the steps associated with each latent variable to be L = 3 and each latent variable has K = 512 candidate values. The planning horizon the raw action space 15 for gym locomotion tasks and 24 for Adroit tasks. As such, the planning horizon in the latent space will be reduced to 5 or 8. Other hyperparameters including architectures can found the Appendix. We test each task with training seeds, each evaluated for 20 episodes. Following the evaluation protocol of TT and IQL, use the v2 version of the datasets for locomotion control and vO for the other tasks.
4.2 RESULTS
On fow-dimensional gym locomotion control tasks, all the model-based methods show comparable performance to model-free ones. However, the performance of model- baselines high- dimensional adroit tasks is much worse than the low-dimensional case, showing inferior scalability with respect to state and action dimensions. In contrast, TAP shows consistently strong performance
5

---

Published as a conference paper at ICLR 2023
among low-dimeniional and high high-dimensional tasks and surpasses other model-basee methods with a large margin on Adroit.
To show how the relative performance between TAP and basclines vary with action dimension, average the relative performance on tasks of the same action dimensionality and plotted them in Figure I(c). The x-axis is the action dimensionality of the tasks and the horizontal red line shows the relative performance of 1, meaning two methods achieve the same score. We can see that TAP scales better terms of decision latency compared TT. Also, TAP shows better performance for tasks with higher action dimensionality, compared to both TT and strong model-free actor-critic methods (CQL/IQL). This can be explained by the increasing difficulty of policy optimization in larger action space due to the curse of dimensionality. In contrast, TAP does the optimization in compact latent space with handful of discrete latent variables.
Table I: Locomotion control results. Numbers are normalised scores following the protocol of Fu al. (2020).
Type
Model-free
Model-based
Dataset
Environment
BC
CQL
IQL
DT
MOReL
TT
TAP (Ours)
Medium-Exper Expert
HalfCheetah
59.9
91.6
86.7
86.8
53.3
95.0
91.8 #:08
Medium-Expe Expeert
Hopper
79.6
105.4
91.5
107.6
108.7
110.0
105.5 #1.7
Medium- Expeert
Walker2d
36.6
108.8
109.6
108.1
95.6
101.9
107.4 #0.9
Medium- Expert
Ant
114.2
115.8
125.6
122.3
116.1
128.8 #2.4
Medium
HalfCheetah
44.4
47.4
42.6
42.1
46.9
45.0 #0.1
Medium
Hopper
63.9
58.0
66.3
67.6
95.4
61.1
63.4 #1.4
Medium
Walker2d
77.3
72.5
78.3
74.0
77.8
79.0
64.9 #2.1
Medium
Ant
92.1
90.5
102.3
94.2
83.1
92.0 #:24
Medium-
Medium-Replay
HalsCheetah
45.5
44.2
40.2
40.8
Medium-Repp Replay
Hopper
27.6
95.0
94.7
82.7
93.6
91.5
87.3 #2.3
Medium-Reee Replay
Walker2d
36.9
77.2
73.9
66.6
49.8
82.6
66.8 #3.1
Bedium-RRelaa
Ant
89.2
93.9
88.8
88.7
77.0
96.7 #:1.4
Mean
84.1
82.2
Table Adroit robotic hand control results. These tasks have high action dimensionality (24 degrees of freedom).
Type
Model-free
Model-based
Dataset
Environment
BC
CQL
IQL
MOPO
Opt-MOPO
TT
TAP (Ours)
Human
Pen
34.4
37.5
71.5
6.2
19.0
36.4
76.5 #8.5
Human
Hammer
1.5
4.4
1.4
0.2
0.5
0.8
1.4 #0.1
Human
Door
0.5
9.9
4.3
0.1
8.8 #1.1
Human
Relocate
0.0
0.2
0.1
0.0
0.2 #0.1
Cloned
Pen
56.9
39.2
37.3
6.2
23.0
11.4
57.4 18.7
Cloned
Hammer
0.8
2.1
2.1
0.2
5.2
0.5
1.2 #0.1
Cloned
Door
-0.1
0.4
1.6
-0.1
11.7 1.5
Cloned
Relocate
0.1
0.1
0.2
0.1
0.2 20.0
Expert
Pen
85.1
107.0
15.1
50.6
72.0
127.4 47.77
Expert
Hammer
125.6
86.7
6.2
23.3
15.5
127.6 1.7
Expert
Door
34.9
101.5
94.1
104.8 +0.8
Expert
Relocate
101.3
95.0
10.3
105.8 +2.7
Mean (without Expert)
11.7
11.7
14.8
6.1
19.6
Mean (all settings)
36.7
40.3
20.1
$1.9
4.3 DECISION LATENCY
Transformer-baseee trajectory generative models (such TT and Gato (Reed et al., 2022)) treat each dimension of the state and action as an individual token. Denoting S as the dimensionality of the
We did not include expert datasets for Adroit when computing average since the common protocol for gym Incomotion evaluation also doesn't include expert datasets and the IQI. paper didn't report these

---

Published as a conference paper at ICLR 2023
TAP
1--1
0.5
Bartora:1
Morbor
Prkor
Oullion
Pall Modi
Tatent
Libiltoo Thmnbali
Planning Barton
mover Sumpligg
Enonditiooall Dazaker
Figure 3: Results of ablation studies, where the height of the bar is the mean normalised scores on gym locomotion control tasks.
state space and A as that of the action space, TT requires D = S A 2 tokens to model a step of trajectory. Since the time complexity of the Transformer ( for sequence of length N the computational cost of TAP is significantly lower especially when the aattinaction dimensionality is high. Specifically, for trajectory with T steps, TAP uses tokens (in autoencoders) and TT uses TD tokens. Therefore, the inference cost of TAP is O(T²) rather than O(D²T²) of TT, leading to more efficient training and planning.
Besides the computational complexity of the forward pass and sampling, we also tested how exactly the computational costs and decision lateney of TT and TAP grow along with the state-action dimen- sionality of the tasks. Tests are done platform with an i5 12900K CPU and single RTX3090 GPU. We fix the planning horizon be 15 (the default planning horizon for TT) and test the decision latency hopper (D 14), halfcheetah (D 23), ant (D 37) and Adroit-pen (D 71) tasks. Even for the task with the lowest dimensionality. TT needs 1.5 seconds to make a decision and the latency grows 32 seconds when dealing with adroit. On the other hand, TAP manages make a decision within 0.05 seconds and meets the real-world robotics deployment criterion proposed by Reed et al. (2022). Moreover, the latency remains nearly constant for different dimensionalities, making it computationally feasible to be applied to high-dimensional control tasks.
4.4 ANALYSIS
TAP is novel model-based RI. method which involves several designs that can potentially be transferred to other model-based or offline RL methods. order to provide more insights into these design choices, provide analyses here, together with ablations of common hyper-parameters like planning horizon. In Figure 3, we showed a summary of the results of ablation studies on gym locomotion tasks. The full results for more hyper- aaper-parameters and scores for each individual task can be found in the Appendix.
Latent Steps A prominent feature of TAP is that planning happens in a latent action space with temporal abstraction. Whenever single latent code is sampled, L steps of transitions that continue the previous trajectory can be decoded. This design increases the efficiency of planning since the number of steps of unrolling is reduced to 1: therefore, the search space is exponentially reduced in size. However, it is unclear how such design affects the decision-makinn method. We thus tested TAP with different latent steps L. namely, the number of steps associated with a latent variable. As shown in the yellow bars in Figure 3, reducing the latent step L to significantly undermines the performance of TAP. We hypothesize that the performance drop is related to the overfitting of the VQ-VAE since we observe the reduced latent step leads to higher estimated return and also a higher prediction error; see Appendix for more details. On the other hand, we found in Appendix Table 8 that the optimal latent steps vary across tasks. This indicates fixed latent step may also be suboptimal. We believe this shows a direction for future work to further decouple the temporal structure between latent actions and decoded trajectories; Appendix D for further discussion.
performance of the
OOD Penalty By varying the threshold for probability clipping (3), we also tested the benefits of having both OOD penalty and return estimations in the objective function Equation (3). When

---

Published as a conference paper at ICLR 2023
3 = -10-5 the estimated return term dominates the optimization objective for the reward scale of the tasks considered. When 3 = 0.5, the OOD penalty is always activated and encourages the agent to choose the most likely trajectory. namely, imitating the behaviour policy. In these two cases, the performance of the TAP agent drops 30.2% and 20.2%, respectively. This experiment demonstrates that solely optimizing the return prior probability leads to suboptimal performance, so both ingredients are necessary for the objective function. Nevertheless, besides these extreme values, the performance of TAP is robust to wide range of choices of B. ranging from 0.002 0.1, as shown in Table 4.
As shown in the blue bars in Figure 3, TAP is not very sensitive the planning horizon, and achieves a mean score of 71.9 (-12.9%) by simply expanding a single latent variable, thus doing steps of planning. This conclusion might be bound to the tasks, rewae-reward locomotion control may require less long- horizon reasoning than more complex decision-making problems. An ablation (Hamrick et al., 2021) of MuZero also shows that test-time planning is not as helpful in more reactive games like Atari compared to board games.
Planning Horizon
Beam Search We also investigate whether structured decoding in the form of beam search is helpful, compared to an alternative strategy of simply autoregressive sampling from the prior policy to generate action proposals. We test the TAP agent with this form of random shooting by sampling 2048 trajectories from the policy prior. As shown in the green bar with the "Prior" label in Figure 3, beam search still performs slightly better. Moreover, its decision latency is lower because the beam width (64) is much smaller than the number of samples needed for direct sampling (2048). On the other hand, the fact that even direct sampling can generate decent plans for TAP shows that the latent space for TAP is compact and can be used for efficient planning.
Sampling from a uniform action prior When TAP performs beam search or direct sampling, the latent action codes are sampled from the learned prior policy rather than the uniform distribution. Such design is straightforward for a VQ-VAE since the objective, in that case, is to sample from the dataset distribution. However, in the case of RL, the aim is different since we would like to find an optimal trajectory within the support of the data. We therefore also test direct sampling where the latent codes are uniformly sampled. The performance of this approach is illustrated as the green bar with "Uniform" label in Figure 3. Sampling according to the uniform distribution largely damages (-37.9%) the performance of the agent, even with an OOD penalty. To further investigate the role of sampling from the prior, in Appendix I in visualize the distribution of trajectories modelled by TAP in two locomotion tasks.
Unconditional Decoder One of the key design choices of TAP is conditioning the decoder on the initial state of trajectory. This allows us model the (conditional) distribution of trajectories with a very small number of latent variables in order to construct compact action space for downstream planning. While it might be obvious that representing multiple steps of nigh-dimensional states and actions with a single latent variable is difficult, we present an empirical investigation of the performance of TAP with the unconditional decoder. The orange bars in Figure 3 show that without the first state as input to the decoder, the performance of TAP drops drastically. This is because, given the same number of latent variables, the reconstruction accuracy of the unconditional decoder much lower. Slightly increasing the number of latent variables from L = 3 to L 1 also does not alleviate the performance drop. Details about this ablation can be found in the Appendix Table 10.
RELATED WORK
TAP fits into a line of work on model-based RL method (Sutton, 1990; Silver et al., 2008; Fairbank,
2008; Deisenroth & Rasmussen, 2011; Lampe & Riedmiller, 2014; Heess et al., 2015; Wang & Ba, 2020; Schrittwieser et al., 2020; Amos et al., 2021; Hafner et al., 2021) because makes decisions by predicting into the future. In such approaches, prediction often occurs in the original state space of an MDP, meaning that models take a current state and action as input and return distribution over future states and rewards. TAP is different from these conventional model-based methods since the TAP takes latent actions and the current state as inputs, returning a whole trajectory. These latent variables act learned latent action space which is compact and allows efficient planning.
8

---

Published as a conference paper at ICLR 2023
Hafner et al. (2019) and Ozair al. (2021) proposed to learn latent state space and dynamics function in the latent space. However, in their cases, the action space of the planner is still the same as the raw MDP and the unrolling of the plan is still tied to the original temporal structure of the environment. Wang et al. (2020); Yang et al. (2021) proposed to learn representations of actions on-the-fly for black-box optimization and path planning setting. While related at a high level, these works operate on environment dynamics that are assumed to be known in advance. TAP extends this idea to the RL setting where the true dynamics of the environment are unknown, requiring a dynamics model and latent action representation space to be learned jointly.
TAP follows a recent line of work that treats RL as a sequence modelling problem (Chen et al., 2021: Janner et al., 2021; Zheng et al., 2022). Such works have used GPT-2 (Radford et al., 2019) style Transformers (Vaswani et al., 2017) to model the whole trajectory of states, actions, rewards and values, and turn the prediction ability into a policy. DT applies an Upside Down RL (Schmid- huber, 2019) approach to predict actions conditioned on rewards. TT applies planning in order to find the best trajectory that gives the highest return. The strong sequence modelling power of the Transformer enables TT to generate long-horizon plans with high accuracy. TAP provides an effi- cient planning solution for TT and all other discrete space planning algorithms to efficiently plan in complex action spaces.
The concept of learning representation of actions and doing RL in this latent action space (or skill embedding space) has been explored in model-free methods (Dadashi et al., 2022; Allshire et al., 2021; Zhou et al., 2020; Merel et al., 2019; Peng al., 2022). Unlike TAP, where latent action space is to enable efficient and robust planning, the motivations for learning a latent action space for model-free methods are varied but the key point is to provide policy constraints. For example, Merel et al. (2019) and Peng et al. (2022) apply this idea to humanoid control to make sure the learned policies are human-like. Zhou et al. (2020) used latent actions to prevent out-fo-distriuution (OOD) actions in offline RL setting. Dadashi et al. (2022) proposes discrete latent action space to make methods designed for discrete action space to be applied to the continuous case. In the literature of teleportation, Losey et al. (2022); Karamcheti et al. (2021) embed the robot's high-dimensional actions into cow-dimensional and human-controllable latent actions.
In this paper, TAP is tested in an offline RL setting (Ernst et al., 2005) where are not allowed to use online experience to improve our policy, A major challenge of offline RL is to prevent out-of- distribution (OOD) actions to be chosen by the learned policy, so that the inaccuracy of the value function and the model will not be exploited. Conservatism has been proposed as a solution to this issue (Kumar et al., 2020; Fujimoto & Gu, 2021; Lu et al., 2022; Kostrikov et al., 2022; Kidambi et al., 2020). TAP can naturally alleviate the OOD actions because it models trajectories generated by behaviour policy and decoded policy is therefore in-distribution.
In the literature of hierarchical RL, a conceptually relevant work to ours is (Co-Reyes et al., 2018), which uses a VAE project the state sequences into continuous latent variables and conditions policies on not only the current state but also the latent variable. Similar ideas also appear in goal- oriented imitation learning. Lynch et al. (2019) proposes policy on a latent variable of trajectories to avoid averaging out multiple modes of the distribution of trajectories. Mandlekar et al. (2020) use a VAE to model the distribution of reachable states and filter the sampled states according to value function, then feed the goal to downstream goal-reaching policy. Since the high-level latent actions are discrete, our work is also related to the (conceptual) options framework (Sutton et al., 1999; Stolle & Precup, 2002; Bacon et al., 2017) since both the latent codes of TAP and options introduce a mechanism for temporal abstraction.
condition a goal-reaching imitation
6 LIMITATIONS
One of the limitations of TAP that it does not distinguish between uncertainty stemming from a stochastic behaviour policy and uncertainty from the environment dynamics. Our empirical results on continuous control with deterministic dynamics suggest that TAP can handle epistemic uncer- tainty, potentially thanks to the OOD penalty in the objective function. However, it is unclear if TAP can handle tasks with stochastic dynamics without any modifications. Another problem for the latent-action model is the loss of fine- -ne-igai control over the generated trajectories. For example, it is nontrivial for TAP to predict future trajectories conditioned on a sequence of actions in the raw action space, which is straightforward to do with conventional dynamics models.

---

Published as a conference paper at ICLR 2023
The work is supported by UCL Digital Inovation Center and AWS. We thank to Minqi Jiang, Robert Kirk, Akbir Khan and Yicheng Luo for insightful discussions.
Acknowledgements
REPRODUCIBILITY STATEMENT
The code reproduce all the experiment results are available in the supplementary materials. README includes guides to setup and environment and commands to run the experiments. The hyperparameters used in our experiments are in the default config file which will be loaded automat- ically. Experiments for all environments can be run on single GPU, Besides codes, pseudo-code in Appendix Algorithm and hyperparameters described Appendix Table can also be helpful to reproduce the results if readers are interested in implementing by themselves.
REFERENCES
Arthur Allshire, Roberto Martín-Martín, Charles Lin, Shawn Manuel, Silvio Savarese, and Ani- mesh Garg. LASER: learning latent action space for efficient reinforcement learning. In IEEE International Conference on Robotics and Automation, ICRA 2021, Xi an. China, May June 5, 2021, pp. 6650-6656. IEEE, 2021. doi: 10.1109/CCA44806... 202.4448800.222155512333 URL https: / doi....... /10. 1109 SA///ICCA48885 2021. 9561232.
Brandon Amos, Samuel Stanton, Denis Yarats, and Andrew Gordon Wilson. On the model-based stochastic value gradient for continuous reinforcement learning. In Ali Jadbabaic, John Lygeros, George J. Pappas, Pablo A. Parrilo, Benjamin Recht, Claire J. Tomlin, and Melanie N. Zeilinger (eds.), Proceedings of the 3rd Annual Conference on Learning for Dynamics and Control, LADC 2021, 7-8 June 2021, Virtual Event, Switzerland, volume 144 of Proceedings of Machine Learn- ing Research, pp. 6-20. PMLR, 2021. URL http: http:://cccceedeeeeeeooo mlr press/v144/ amos2la. html.
Pierre-Lue Bacon, Jean Harb, and Doina Precup. The option-criti architecture. In Satinder Singh and Shaul Markovitch (eds.), Proceedings of the Thirty AAAI Conference Artificial Intel- ligence, February 4-9, 2017, San Francisco, California, USA, pp. 1726-1734. Press, 2017. URL http: iaaiiaaaaaaaaaa.aaaaaa.aaaaaaaae org/ /wwp///ssned////////e php/ AAAP/AAAAI/AT /tt.//AAA11111 /ttpa//AAAA//aapppeee/ae////////
Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Michael Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning sequence modeling. In A. Beygelzimer, Y. Dauphin. P. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, 2021. URL net/foren 2o=
John D. Co-Reyes, Yuxuan Liu, Abhishek Gupta, Benjamin Eysenbach, Pieter Abbeel, and Sergey Levine. Self-consistent trajectory autoencoder: Hierarchical reinforcement learning with trajec- tory embeddings. In Jennifer G. Dy and Andreas Krause (eds.), Proceedings of the 35th Inter- national Conference on Machine Learning, ICML 2018, Stockholmsmässan. Stockholm, Sweden, July 10-15, 2018, volume 80 of Proceedings of Machine Learning Research, pp. 1008-1017. PMLR, 2018. URL http:://// http///pooceeeesseeeeeeeeec mir www...pees//veee /erress//OO/e/ co reyes html.
Robert Dadashi, Léonard Hussenot, Damien Vincent. Sertan Girgin, Anton Raichuk, Matthieu Geist, and Olivier Pietquin. Continuous control with action quantization from demonstrations. In Kama- lika Chaudhuri, Stefanie Jegelka, Song, Csaba Szepesvári, Gang Niu, and Sivan Sabato (eds.), International Conference on Machine Learning. ICML. 2022, 17-23 July 2022, Baltimore, Mary- land, USA, volume 162 of Proceedings of Machine Learning Research, pp. 4537-4551 PMLR, 2022. URL https: / prwceedings. mlr. wpw.sp/sseee........... / vvlvvvlll 62 dadashi22 inchii2224 22a html.
Mare Peter Deisenroth and Carl Edward Rasmussen. PILCO: A model-based and data-efficient approach to policy search. Lise Getoor and Tobias Scheffer (eds.). Proceedings of the 28th International Conference on Machine Learning, ICML 2011, Bellevue, Washington, USA, June 28 -July 2, 201 pp. 465-472. Omnipress, 2011. URL https: cml wtw.mc/cccc.eecccccccmmmmeecce wtcp//////aapeeeeeeeeeee 323 1immmmllgg paper .dff.
Damien Ernst, Pierre Geurts, and Louis Wehenkel. Tree-based batch mode reinforcement learning. J. Mach. Learn. Res., 6:503-556, dec 2005. ISSN 1532-4435.
10

---

Published as a conference paper at ICLR 2023
Michael Fairbank. Reinforcement learning by value gradients. CoRR, abs/0803.3539, 2008. URL attt:://arxe org/ wabs//803 3539.
Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4RL: datasets for deep data-driven reinforcement learning. CoRR, abs/2004.07219, 2020. URL https: //axxiv.
org/abs/2001 07219.
Scott Fujimoto and Shixiang Gu. A minimalist approach to offline reinforcement learning. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, 2021. URL https:/ http///pppeeveeeeeeeeeeeeeeeeeee ew. net. forum? i= Q32U7dzwXpc.
Danijar Hafner, Timothy P. Lillicrap, Ian Fischer, Ruhen Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. In Kamalika Chaudhuri and Ruslan Salakhutdinov (eds.), Proceedings of the 36th International Conference on Machine Learning, ICML 2019, 9-15 June 2019, Long Beach, California, USA, volume 97 of Proceedings of Machine Learning Research, pp. 2555-2565. PMLR, 2019. S htttt::/prrceedings. malr.press/ve ha fnerl 19.. html.
Danijar Hafner, Timothy P. Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering atari with discrete world models. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021. URL net/ wwruurrummm.. barrumiid=ooabbyybboo
Jessica B. Hamrick, Abram L. Friesen, Feryal Behbahani, Arthur Guez, Fabio Viola, Sims Wither- spoon, Thomas Anthony, Lars Holger Buesing, Petar Velickovic, and Theophane Weber. On the role of planning in model-based deep reinforcement learning. In 9th International Conference Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021. URL
wwt.:e/en///nmmeeeeee IrrM 111M4444GC0 4DGB21.
Nicolas Heess, Gregory Wayne, David Silver, Timothy P. Lillicrap, Tom Erez, and Yuval Tassa. Learning continuous control policies by stochastic value gradients. In Corinna Cortes, Neil D. Lawrence, Daniel D. Lee, Masashi Sugiyama, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 28: Anmual Conference on Neural In- formation Processing Systems 2015, December 7-12, 2015, Montreal, Quebec, Canada, pp- 2944-2952, 2015. URL https: http::////pccceeegg..... 1455100133:6662265c0ee1111131113
wwt.:eeps...ppppee/////./.. pttppaape///1////aeeeeeeeee 2215//ash/// Abstract .tml.
Michael Janner, Qiyang Li, and Sergey Levine. Offline reinforcement learning as big sequence modeling problem. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan (eds.), Advances Neural Information Processing Systems, 2021. URL https:////// https://openreww/ /oppreviee.eeeeeeeeeeeeeee net net/Eorum? i Wggk Sw
Siddharth Karamcheti, Albert Zhai, Dylan P. Losey, and Dorsa Sadigh. Learning visually guided latent actions for assistive telcoperation. In Ali Jadbabaie, John Lygeros, George J. Pappas, Pablo A. Parrilo, Benjamin Recht, Claire J. Tomlin, and Melanie N. Zeilinger (eds.), Proceed- ings of the 3rd Annual Conference on Learning for Dynamics and Control, L4DCC 2021, 7-8 June 2021, Virtual Event, Switzerland, volume 144 of Proceedings of Machine Learning Re- search, pp. 1230-1241. PMLR, 2021. URL http: http///roceedceeee ngs. mllr. wtt.ssmmr...res/..e... press/v144/ karamchetis i 21...... html. .lmm.
karamchet
Rahul Kidambi, Aravind Rajeswaran, Praneeth Netrapalli, and Thorsten Joachims. Morel: Model- based offline reinforcement learning. In Hugo Larochelle, Mare' Aurelio Ranzato, Raia Had- sell, Maria-Florina Balcan, and Hsuan-Tien Lin (eds.), Advances in Neural Information Process- ing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6- 2020, virtual, 2020. URL https: http://ppcceeeeng.eeeep neerip ee/ paper/2020/her
f7e 4 f a
html.
Ilya Kostrikov, Ashvin Nair, and Sergey Levine. Offline reinforcement learning with implicit q- learning. In International Conference on Learning Representations, 2022. URL https:/ // openreview wttp:iiwwwwe.eeeeeeeeeeeeeeeeeee net//orumm d d= 9ZJWFB.

---

Published as a conference paper at ICLR 2023
Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. In Hugo Larochelle, Mare' Aurelio Ranzato, Raia Hadsell, Maria-Florine Balcan, and Hsuan- Tien Lin (eds.), Advances in Neural Information Processing Systems 33: An- nual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual, 2020. URL https: /ptp:ceee//////peeeeeeeeeeeeeeee ngs neur pss. cc asc/paper// (eaper/2020/1e 2 0d2b20618
Abstractttttttttcc html.
Thomas Lampe and Martin A. Riedmiller. Approximate model-assisted neural fitted q-iteration. In 2014 International Joint Conference Neural Networks, UCNN 2014, Beijing, China, July 6-11, 2014, pp. 2698-2704. IEEE, 2014. doi: 10. 111191/CNN.2211.888933333 URL https: //doi. wwooorgg//// 10. 1109/ IJCNN. 2 2014. 6889733.
Dylan P. Losey, Hong Jun Jeon, Mengxi Li, Krishnan Srinivasan, Ajay Mandlekar, Animesh Garg, Jeannette Bohg. and Dorsa Sadigh. Learning latent actions to control assistive robots. Auton, Robots, 46(1):115-147, 2022. doi: 10. 10.1007/s10514-022
100-00011---01000--- w. URL stttps://do...o
10. 1 100.1007/600 s111s115111111 10005-w.
Cong Lu, Philip Ball, Jack Parker-Holder, Michael Osborne, and Stephen Roberts. Revisiting de- sign choices in offline model based reinforcement learning. In International Conference Learn- ing Representations, 2022. URL https: opw.review.weeee/////e net/ forum ?id= 9
Corey Lynch, Mohi Khansari, Ted Xiao, Vikash Kumar, Jonathan Tompson, Sergey Levine, and Pierre Sermanet. Learning latent plans from play. Conference on Robot Learning (CoRL), 2019. URL hwwphhtpt///..t.. /ttp:///axvvvvoo
abs 1903. 01973.
Ajay Mandlekar, Fabio Ramos, Byron Boots. Silvio Savarese, Li Fei-Fei, Animesh Garg, and Di- eter Fox. IRIS: implicit reinforcement without interaction at scale for learning control from of- fline robot manipulation data. 2020 IEEE International Conference on Robotics and Automa- fion, ICRA 2020, Paris, France, May 31 August 31, 2020, pp. 4414 4420. IEEE, 2020. doi: A1010/9//RR4444.....2 2020.919995. URL https: / doi org 1ttp:/0..11191//./11./ 2020. 9196935.
Josh Merel, Leonard Hasenclever, Alexandre Galashov, Arun Ahuja, Vu Pham, Greg Wayne, Yee Whye Teh, and Nicolas Heess. Neural probabilistic motor primitives for humanoid con- trol. In 7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019. OpenReview.net, 2019. URL. https: //tp://eopp/eeeeeeeeeeeeeeee net forum? id-BJ16T
Sherjil Ozair, Yazhe Li, Ali Razavi, Ioannis Antonoglou. Aäron van den Oord, and Oriol Vinyals. Vector quantized models for planning. Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning. ICML 2021, 18-24 July 2021, Virtual Event, volume 139 of Proceedings of Machine Learning Research, pp. 8302-8313, PMLR, 2021. URL http: /est::/ppreet http:///ccceedingeee mlr. wttres//v13/oozarr/r
html.
Xue Bin Peng, Yunrong Guo, Lina Halper, Sergey Levine, and Sanja Fidler. ASE: large-scale reusable adversarial skill embeddings for physically simulated characters. ACM Trans. Graph., (4):94:1-94:17, 2022. doi: 10. 2,1144555282233333 10. URI. http::ts::/// :/doi.. org//0.11445 3528223. 3530110.
Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and llya Sutskever. Language models are unsupervised multitask learners. 2019.
Scott Reed, Konrad ZoIna, Emilio Parisotto, Sergio Gomez Colmenarejo, Alexander Novikov, Gabriel Barth-Maron, Mai Gimenez, Yury Sulsky, Jackie Kay, Jost Tobias Springenberg, Tom Eccles, Jake Bruce, Ali Razavi, Ashley Edwards, Nicolas Heess, Yutian Chen, Raia Had- sell, Oriol Vinyals, Mahyar Bordbar, and Nando de Freitas. A generalist agent, 2022. URL https://arxiss http:///aaxxvv...//b... attpa//s...2..1.. 06175.
Jürgen Schmidhuber. Reinforcement learning upside down: Don't tt predict rewards just map them to actions. CoRR, abs/1912.02875,, 2019. URL / http://iwvvr////////...///. wow.oorg/aa.... /ags/////// / 91 02875.
12

---

Published as a conference paper at ICLR 2023
Julian Schrittieser, loannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, et al. Mastering atari, go, chess and shogi by planning with learned model. Nature, 588(7839):604-600 2020.
David Silver, Richard S. Sutton, and Martin Müller. Sample-based learning and search with permanent and transient memories. In William W. Cohen, Andrew McCallum, and Sam T. Roweis (eds.), Machine Learning, Proceedings of the Twenty-Fifth International Conference (ICML 2008), Helsinki, Finland, June 5-9, 2008, volume 307 ACM International Confer- Proceeding Series, pp. 968 975. ACM, 2008. doi: 01144//139115.13902222 URL https: /oi..i......
https///dii..ogg......o. www.ooo........ooooo.oo 11w.11111111111111..1 1101119911111 1390278
David Silver, Thomas Hubert, Julian Schrittwieser, Ioannis Antonoglou, Matthew Lai, Arthur Guez, Mare Lanctot, Laurent Sifre, Dharshan Kumaran, Thore Graepel, et al. A general reinforcement learning algorithm that masters chess, shogi, and go through self-play. Science, 362(6419):1140- 1144, 2018.
Martin Stolle and Doina Precup. Learning options in reinforcement learning. In Sven Koenig and Robert C. Holte (eds.), Abstraction, Reformulation and Approximation, 5th International Sympo- sium, SARA 2002, Kananaskis, Alberta, Canada, August 2-4, 2002, Proceedings, volume 2371 of Lecture Notes in Computer Science, pp. 212-223, Springer, 2002. doi: 10.10077/3-44044562--8
-16. URL https:://// /ot. 10. 1007 5005544022222 8_..
Richard S. Sutton. Integrated architectures for learning, planning, and reacting based on approximating dynamic programming. In Bruce W. Porter and Raymond J. Mooney (eds.), Machine Learning. Proceedings of the Seventh International Conference on Ma- chine Learning, Austin, Texas, USA, June 21-23, 1990, pp. 216-224, Morgan Kaufmann, URL https: /tte:///ccc.iee /tt////cii.o//00.11111100
1990. doi: 11011609977---58----113333-00000 b978 1978-1-55860-11
50030-4.
Richard S. Sutton, Doina Precup, and Satinder Singh. Between mdps and semi-mdps: A framework for temporal abstraction in reinforcement eeinfornementcaaainnnn... Artif. Intell., 111-111 1999, doi: 10. (
URL
00052-1.
C. Tillmann, Stephan Vogel, Hermann Ney, Arkaitz Zubiaga, and H. Sawaf. Accelerated dp based search for statistical translation. 5th European Conference on Speech Communication and Tech- nology (Eurospeech 1997), 1997.
Aäron van den Oord, Nal Kalehbrenner, and Koray Kavukcuoglu. Pixel recurrent neural networks. In Maria-Florina Balcan and Kilian Q. Weinberger (eds.), Proceedings of the 33nd International Conference on Machine Learning, ICML 2016, New York City. NY, USA, June 19-24, 2016, vol- ume 48 of JMLR Workshop and Conference Proceedings, pp. 1747-1756. JMLR.org, 2016. URL http:/ pr procedd ngs .m malr.pee pwwpprres/vvv v48//ord166/ html.
Aäron van den Oord, Oriol Vinyals, and Koray Kavukcuoglu. Neural discrete repre- sentation learning. In Isabelle Guyon, Ulrike von Luxburg, Samy Bengio, Hanna M. Wallach, Rob Fergus, S. V. N. Vishwanathan, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 30: Annual Conference on Neural Infor- mation Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA, pp. 6306-6315, 2017. URL 7 7a wtw.aa/eee...eeee0e2003003339999
http::/s.....oosss..pp////////// httpp/ppp/////aaas/.. http:///ppooceeeigs.eeeeee html.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Isabelle Guyon, Ulrike von Luxburg, Samy Bengio, Hanna M. Wallach, Rob Fergus, S. V. N. Vishwanathan, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA, pp. cc wttpc/ppppppp../..////ppp //t0//h/hh/h//// html.
5998-6008, 2017. URL httpss//ppcocceeigg...eeee 3f5ee2435472
Kerong Wang, Hanye Zhao, Xufang Luo, Kan Ren, Weinan Zhang, and Dongsheng Li. Bootstrapped transformer for offline reinforcement learning.
CoRR, abs/2206.08569, 2022. doi: 10.48550/
2xXiv2200.008669 URL https: h/tpss://dooi org/10 a 2206..6666 08569.
13

---

Published as a conference paper at ICLR 2023
Linnan Wang. Rodrigo Fonseca, and Yuandong Tian. Learning search space parti- tion for black-box optimization using monte carlo tree search. In Hugo Larochelle, Mare' Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Tien Lin (eds.), Advances in Neural Information Processing Systems 33: Annual Conference on Neu- ral Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, vir- tual, 2020. URL nttps://proceedii: neurips. httpg/g......p...p..paee////eeee www.eecceecccceecccccc04411
hccppappr//22///aaee Abstract. html.
Tingwu Wang and Jimmy Ba. Exploring model-based planning with policy networks. 8th International Conference on Learning Representations, ICIR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net, 2020. URL att.recttttosstttttttt....... H1exf64Kwil
net/ forum?idd
Kevin Yang, Tianjun Zhang, Chris Cummins, Brandon Cui, Benoit Steiner, Linnan Wang, Joseph E. Gonzalez, Dan Klein, and Yuandong Tian. Learning space partitions for path planning. In Mare" Aurelio Ranzato, Alina Beygelzimer, Yann N. Dauphin, Percy Liang, and Jennifer Wort- man Vaughan (eds.), Advances in Neural Information Processing Systems 34: Anmual Conference on Neural Information Processing Systems 2021, NeurIPS 2021, December 6-14, 2021, virtual, pp. 378-33883333 391, 2021. URL https: http//pooceeeengeeeeeee neurips cc http/ppprr22000 220e 1//ashhh///// (aeeee949999388 Abstrract html.
03a36555 9bdee (ew.44444999999942222
Weirui Ye, Shaohuai Liu, Thanard Kurutach, Pieter Abbeel, and Yang Gao. Mastering atari games with limited data. In Marc' Aurelio Ranzato, Alina Beygelzimer, Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan (eds.), Advances in Neural Information Processing Systems 34: Annual Conference Neural Information Processing Systems 2021, NeurIPS 2021, December 6-14, 2021, virtual, pp. 25476-25488, 2021. URL http:: /ttpctpps/pppoceeeiiii ngs. nwurips.ec.om cc/ paper/ /2021 hww.h/aah///aeaee dwwcc 3d..8ddc330ccc. cad Fba
ract .html.l hhmll.
Tianhe Yu, Garrett Thomas, Lantao Yu, Stefano Ermon, James Y. Zou, Sergey Levine, Chelsea Finn, and Tengyu Ma. MOPO: model-based offline policy optimization. In Hugo Larochelle, Mare' Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan- Tien Lin (eds.), Advances in Neural Information Processing Systems 33: Annual Con- ference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual, 2020. URL https: http:/:::/pppoceo.....eee.e Abstract html.
ccc http//pper/222////2
Qinqing Zheng, Amy Zhang, and Aditya Grover. Online decision transformer. abs/2202.05607,, 2022. URL hww.h:tt/t//ttte.eee wwtp:/ogg//a/////ee/ee 05600
CoRR,
Wenxuan Zhou, Sujay Bajracharya, and David Held. PLAS: latent action space for offline rein- forcement learning. In Jens Kober, Fabio Ramos, and Claire J. Tomlin (eds.), 4th Conference on Robot Learning, CoRI. 2020, 16-18 November 2020, Virtual Event Cambridge, MA, USA, volume 155 of Proceedings of Machine Learning Research, pp. 1719-1735. PMLR, 2020. URL https: / pttpp/eeecceg...ce//e////o mlr, war.ppres//vee vv1.vvvl5////e zhoullb. html.
14

---

Published as a conference paper at ICLR 2023
8
125
TT
1.00
0.75
$1
Decoder
0.50
0.75
0.00
(a) TT and Gato
(b) TAP
(c) Scale Planning Hori-
Figure 4: (a) Illustration of dimension wise autoregressive modelling used by Trajectory Trans- former. Blocks represent discretised state/action dimensions discrete latent variables. (b) shows TAP style modelling. (c) Horizon Scalability, tested for a low-dimensional task, hopper.
Algorithm TAP Beam Search
Input: Current state sequence length T, beam width B, expansion factor E, prior model p and decoding function h
1: Sample initial context latents Zi {=1 (2(")) p(z(s), S BE}
2: for = 2 to 7 do
3: for b= 1 to B do
4: Expand the beam C, < {(Z<<) zelze ~ plz s) E}
5: end for
6: Decode Trajectories Ti +- {r|T h(zett1, s), Zett1 Ct}
7: Get top trajectories according to return
8: Set corresponding latent sequences to be Ze
9: end for
10: return best trajectory in Tr
A SCALE UP THE PLANNING HORIZON
TAP not only makes it easier scale up the state and action dimensionality, but also helps the agent to plan over a longer horizon with much less computation. As shown in Figure 4(c), the decision latency increases with the number of planning steps, tested on the hopper-medium--eplay task. TAP is able to plan 33 steps into the future within 0.25 seconds. 'abble 3 shows how the performance on locomotion control tasks varies with the planning horizon. We can see that even just expanding single latent variable (corresponding to 3 steps in the original space), TAP can actually achieve decent performance most tasks. The longer horizon planning is helpful for hopper, especially in repilay-replyy settings. In hopper tasks, the agent can easily fall over compared other tasks. Once the agent falls over, the episode will be immediately terminated, so similar actions might lead drastic longer-term differences. This might mean learning the dynamics function and using it to help learn the value function easier than directly learning the value function. On the other hand, the medium-replay dataset has higher diversity which may make longer planning useful. Both medium and medium expert datasets are generated by one or two fixed policies. This means the possible actions for the next step are quite limited, so searching may not give very different results given the estimated return. The medium-replay dataset, on the other hand, contains the data in the replay buffer. Given the policy parameters are changing during learning, the behaviour policy is then mixture of lot of different policies, making TAP has to consider more options when acting.
B THE ROLE OOD PENALTY
Besides sampling from the prior distribution given by the transformer, also used the estimated prior probability to prevent ofttt-ddistribttion (OOD) trajectories to selected during the planning. In Table 4 we showed how the threshold of OOD penalty will affect the performance of TAP. We can
15

---

Published as a conference paper at ICLR 2023
Table 3: Locomotion performance with different planning horizon
Dataset
Environment
Horizon=3
Horizon=9
Horizon=15
Horizon=21
Medium-ENa Expeert
HalfCheetah
91.74
91.17
91.77
90.6
Medium-Expert
Hopper
88.3
104.01
105.53
96.4
Medium- Medium-Expert
Walker2d
99.4
107.25
107
108.82
Medium Mediuu--Expett
Ant
126.07
130.96
128.82
134.47
Medium
HalfCheetah
44.98
45.04
44.7
Medium
Hopper
49.7
64.44
63.44
66.7
Medium
Walker2d
65.04
63.21
64.87
52.03
Medium
Ant
87.8
84.36
92.0
89.87
Medium Medium-Replay
HalfCheetah
40.59
41.51
40.78
41.36
Medium- Seeiium-Repaay
Hopper
29.39
90.17
87.3
96.35
Medium- Replay
Walker2d
57.58
58.41
66.85
69.25
Medium- Medium-Replay
Ant
83.06
98.67
96.71
97.7
Mean
81.6
82.5
82.4
when very low threshold is chosen = 10-5 the performance of the agent drops significantly because the trajectory of such low prior probability will be not sampled in the first place and the OOD penalty is working. Some the OOD trajectories with high prediction error and high value will be selected in this case. A larger 3 in the range of [0.002, 0.1] gives consistent similar results. However, keeping improving to 0.5 damages the performance again most of the trajectories do not have such high prior probability and the penalty will force the agent to choose the most probable trajectory. This B = 0.5 case can also be treated as imitation learning baseline, showing searching for the higher return trajectory rather than cloning the behaviour policy is helpful.
Table 4: Locomotion performance with different probability threshold |
Dataset
Environment
8 = 10
8 = 0.002
8 = 0.01
= 0.05
8 = 0.1
8 = 0.5
Medium-Expert
HalfCheetah
64.3
91.0
90.3
91.8
92.2
87.7
Kadium-Experts
Hopper
66.3
95.1
94.5
105.5
107.0
76.1
Medium-Erae Expert
Walker2d
86.6
110.1
109.2
107.4
106.5
104.0
Medium-Expert
Ant
104.7
132.1
133.7
128.8
127.4
116.2
Medium
HalfChectah
42.6
44.8
44.8
45.0
45.1
42.9
Medium
Hopper
42.9
64.6
67.0
63.4
63.8
47.5
Medium
Walker2d
63.1
49.2
53.8
61.9
58.7
Medium
Ant
85.1
88.5
89.8
92.0
89.7
85.1
Medium- Replay
HalfCheetah Hopper
36.2
42.2
40.5
40.8
41.0
Kedium-RRppayy
19.9
94.4
93.2
87.3
97.0
3.0
Medium- Repllyy
Walker2d
27.1
57.3
59.6
66.8
65.6
50.3
Medium-Ree Replay
Ant
52.5
96.8
95.7
96.7
93.7
76.8
Mean
57.6
80.5
81.0
82.5
82.8
65.8
C BASELINES
As for the baselines, we gather the strong baselines for each of the three sets of tasks. We include model-based methods such as MOReL (Kidambi et al., 2020) and (Optimized) MOPO (Yu al., 2020: Lu et al., 2022); actor-critic methods CQL (Kumar et al., 2020) and IQL (Kostrikov et al., 2022); and sequence modelling methods DT and TT. We use scores reported by the papers by default with few exceptions.
All the methods besides TT did not report performance on Ant locomotion control so we run them using their official codebase.
For MaMazzeUUtta,, we run IQL with their official codebase and run TT with the code given by the author (Janner et al., 2021).
16

---

Published as a conference paper at ICLR 2023
The Ant locomotion results of TT come from their official GitHub repository ?
Since CQL is tested in the vO version of locomotion control tasks, the v2 results reported by IQL (Kostrikov et al., 2022).
All the behaviour cloning (BC) results are reported by D4RL (Fu et al., 2020).
Chen et al. (2021, DT) did not report the results for AntMaze and Adroit so we use DT AntMaze results reported by IQL (Kostrikov al., 2022).
MOPO results for Adroit come from (Lu al., 2022).
TT results for adroit come from (Wang et al., 2022). We also tested two training seeds, each 10 evaluation episodes, for pen-human and pen-cloned and get 12.949.2 and 97.3 respectively.
D OTHER DESIGN POTENTIALS FOR THE BOTTLENECK
The design of the prior model is highly conditioned the design of the bottleneck and the decoder, and will have an impact when used for planning. Following Decision Transformer (Chen et al., 2021) and Trajectory Transformer (Janner et al., 2021), used GPT-2 style transformer with causal masking for our encoder and decoder. this case, the information in the future token will not follow back recent ones. Such design is conventional for sequence modelling but necessarily optimal. For example, one could reverse the order of the masking for the decoder and make the planning goal-oriented, whereas the prior should also be trained reversed order.
To fully decouple the planning and temporal structure, we also tried using attention rather than pooling and tiling to construct the bottleneck. Such a design shows similar performance when doing vanilla sampling. Such a design also prevents us from applying beam search and is thus less efficient and we did not put that the main text. However, in Table 5 we show TAP with attention bottleneck works as well as fixed length step length indicates the temporal structure of the MDP may not be a necessary inductive bias for planning. This can motivate future works that design more flexible and efficient ways of planning based on the latent-action model.
Table 5: Attention-based bottleneck.
Dataset
Environment
beam search
L 3 sample
attention sample
Medium Medium-Expert
HalfCheetah
91.8
89.9
91.3
Medium Medium-Expert
Hopper
105.5
98.5
108.2
Medium-Faae Expuert
Walker2d
107.4
107.7
104.1
Medium-Expert
Ant
128.8
124.7
129.3
Medium Medium
HalfCheetah
45.0
44.3
45.2
Hopper
63.4
64.3
65.9
Medium
Walker2d
64.9
55.5
58.0
Medium
Ant
92.0
88.8
90.2
Medium-Reng Replay
HalfCheetah
40.8
39.8
39.8
Medium- Replay
Hopper
87.3
79.0
79.7
Medium Medium-Replay
Walker2d
66.8
66.0
68.8
Medium- Replly
Ant
96.7
81.4
89.9
Mean
82.5
78.3
80.9
E ANTMAZE DETAILS
AntMaze is a sparse-reward continuous-control task in which the agent has to control robotic ant to navigate to the target position. We include the goal position the observation space so that the generated trajectories goal-aware.? In order to extensively test different methods, also
Phttp://gtthb.com/aamrriiii hamm://amee//ttnccttrrrrrraaae "The goal position specified by the infos goal field in the d4rl dataset by default.
17

---

Published as a conference paper at ICLR 2023
add customized larger antmaze that we called Antmaze- which has doubled the size of the previously largest map (Antmaze-Large)
Antmaze has the same dimensionality locomotion ant and so will not provide too much direct evidence about the action dimension scalability, so it's more about an orthogonal evaluation on sparse reward problems. AntMaze is challenging because there are lot of sub-optimal trajectories in the dataset that is navigating to different goal positions rather than the target position in the test time. Reaching these goals will not give a reward to the agent. The reward will only be given when the agent reaches the true target, which is also the goal in the test time. On AntMaze, TAP shows better performance on more challenging large dataset, but slightly inferior to (+Q) on the medium datasets. We did not use the IQL critic as our value estimator; better value estimation approach than our current Monte-Carlo estimation would be expected to bring orthogonal improvements. Effective performance without a separate Q Q-function is partially due to the inclusion of the goal position into the observation: TAP can learn to generalize across goal positions instead of completely relying on the value estimation.
TT has shown using a separate IQL value function can help sequence generative models to solve AntMaze. Such an approach, however, further increases the computational cost for sampling tra- jectories. Here we instead propose an alternative efficient approach that let TAP jointly leverage trajectories with different objectives and also the reward signals. This is simple as putting the goal positions to the observation of the agent. We hypothesized that goal positions can be helpful for TAP it models the whole distribution of the trajectories. When acting in the environment, having the Trajectories conditioned on the goal can narrow down the sampled trajectories to focus on the correct direction, therefore simplifying the planning. Note that for actor-critic methods, the positions of the test time target position are encoded by the critic network and including the goal position will not provide extra information the agent.
In order to pressure test IQL, TT-Q and TAP in larger AntMaze environments, we introduced the AntMaze-UItra task, which has 10 14 blocks of effective size. This and 2 times as large as AntMaze-Medium (6 x 6) and AntMaze-Large (7 x 7) tasks, respectively. A visualisation of these three tasks can be found in Figure In addition to the increased size, the complexity of the walls also results in multiple possible routes to navigate from the bottom left corner to the top right. Similar to the antmaze-medium and antmaze-large tasks, we let the agent run for 1K episodes, each with 1K steps collect the trajectories. However, we find that such amount of data is far away from being enough to learn a proper value function, as IQL gives a very poor performance on this task. Using the Monte-Carlo estimation, TAP even suffers more from this issue. We find that planning with beam search, in this case, is even worse than just sampling a random trajectory: with random sampling, TAP gives 22.00 +/- 4.14 on AntMaze-Ultra-Pie and 26.00 +/- 4.39 on AntMaze-Ultra Ultra-Diverse, However, beam search will decrease its performance to 21.00 +/- 4.07 on Ultra-Play and 22.00 +/- 4.14 on Ultra-Diverse.
Table 6: Antmaze results. The Antmaze-Ultra is our customised larger environment.
Dataset
Environment
BC
CQL
IQL
DT
TT(+2)
TAP
Play
matmaze-Mediumg
0.0
61.2
71.2
0.0
93.3 #6.1
78.0 #4.14
Diverse
Antmaze-Medaa
0.0
53.7
70.0
0.0
100.0
10.0
85.0 +3.57
Play
Antmaze- Marga
0.0
15.8
39.6
0.0
66.7
112.2
74.0 14.39
Diverse
Antmaze- Large
0.0
14.9
47.5
0.0
60.0 112.7
82.0 #5.00
Play
Antmaze-Ulmr
8.3
20.0 #10.0
22.0 #4.1
Diverse
Antmaze-Ulne
15.6
33.3 +12.2
26.0 14.4
Mean
42.0
62.2
61.2
F INTERPRET LATENT CODES
A key design choice we made is to let the decoder, not only the prior of latents, be conditioned on the current state. Such design leads to compact latent action space that enables efficient planning. Without the state as input, the decoder must reconstruct the whole trajectory relying solely on the latent codes, so the latent codes contain information about the whole trajectory segments
18

---

Published as a conference paper at ICLR 2023
+
(a) AntMaze-Medium
(b) AntMaze-large
(c) AntMaze-Ultra
Figure 5: Visualization of different antmaze environments.
(states, actions, rewards, values). other words, latent codes will be discrete representation of the trajectory, where latent codes sequence can be decoded into unique trajectory. On the other hand, conditioned on the current state, the latent codes only need to carry information about how the future trajectories branch from the existing trajectory. Every L step, branch of the trajectory will fork into possible futures. Each the latent codes corresponds to a branch of this tree of possible future trajectories, where the root node is the current state. The forking the uncertainty of the dynamics model or by the stochasticity of the behaviour policy. The number of latent codes needed reconstruct trajectory given starting state will reduce long the model knows more about the dynamics of the environment and the behaviour policy. Assume trajectory distribution p(T ((T|811/8 μ²) with behaviour policy f has deterministic Markovian transition dynamics, and the model has recovered the true dynamics. The probability of a valid trajectory that fit the dynamics can be expressed as ((T(81;))) IT (=(a,ss)) So the discrete latent variables only need to approximate the distribution of the behaviour policy, where latent codes correspond to grouped components of the behaviour policy. Note that the policy is not only grouped in single step but also across multiple steps. Intuitively, the latent codes can then be interpreted as state-conde nate-conditioned options (Sutton et al., 1999), learning by decomposing the behaviour policy into weighted sum of sub-policies. The number of grouped policy components of the behaviour policy can be much less than the dimensionalities needed to represent the whole trajectory. As will show in Section 4.4, the L 3 or L = 4 gives an optimal performance of TAP on D4RL tasks, which means 3 steps of the transitions with up hundreds of dimensions can be expressed with a single discrete latent variable, with K = 512 bins. In contrast, the discrete representation of the trajectory used in TT will allocate L.D discrete codes describe the same number of transitions, This means TAP has smaller and more compact latent space for downstream planning.
be caused by
G HYPERPARAMETERS
In Table we showed the hyperparameters for TAP, both for training and planning.
19

---

Published as a conference paper at ICLR 2023
Table List of Hyper-parameteer
Environment
Hyper-pane parreepsrsmeters
Value
All
learning rate batch dropout probability number of attention heads number of steps for latent code beam expansion factor
2e
All
512
0.1
All
All
All
Embedding size latent code
512
All
g
0.05
Locomotion Control
training sequence length
Locomotion Control
discount
0.99
Locomotion Control
number of layers feature vector size
Locomotion Control
512
Locomotion Control
512
Locomotion Control
beam width
64
Locomotion Control
planning horizon
15
AntMaze
training sequence length discount
15
AntMaze
0.998
AntMaze
number of layers feature vector size
AntMaze
512
AntMaze
8192
AntMaze
beam width planning horizon
AntMaze
15
Adroit
training sequence length
24
Adroit
discount
0.99
Adroit
number of layers
Adroit
feature vector size K
256
Adroit
512
Adroit
beam width planning horizon
256
Adroit
24
Table 8: The ablation of number of steps for a latent code L.
Dataset
Environment
L
L 2
3
L
Medium-Expeni
HalfCheetah
37.4
68.1
91.8
92.5
Medium- Expert
Hopper
44.7
55.6
105.5
102.2
Medium Expert
Walker2d
94.9
108.9
107.4
108.5
Medium Medium-Expert
Ant
122.7
118.0
128.8
132.2
Medium
HalfCheetah
43.5
42.6
45.0
45.4
Medium
Hopper
55.1
70.5
63.4
70.7
Medium
Walker2d
70.6
64.9
69.3
Medium
Ant
92.3
98.2
92.0
88.9
Medium-Reple RRelaay
HalfCheetah
30.3
30.1
40.8
40.0
Mediim--Repaa
Hopper
61.9
69.9
87.3
85.6
Medium -eeiium-Reppyy
Walker2d
29.4
60.5
66.8
54.8
Medium Mcdium-Replay
Ant
80.1
90.0
96.7
95.1
Mean
60.6
73.6
82.5
82.1
H PREDICTION ERROR, SEARCH VALUE AND PERFORMANCE
To investigate the reason that temporal abstraction and state-conditional encoder will be helpful for TAP, we provide more metrics about TAP search in Figure 6. Scores show the performance of the agent. Search values optimal value(return) found by TAP in the initial states. Entors mean squared errors of state prediction given the action sequences output by the TAP. We can see: 1) For state-conditional decoder, there is negative correlation between prediction emror and scores.
20

---

Published as a conference paper at ICLR 2023
-31-
noliical wwsziba
cultical wesiia
0-
70-
18-
oudiol uxsoriiion.
79-
200
48-
48-
34
(a) score
(b) search value
(c) mean squared error
Figure 6: Scores, search value and error for TAP agent with different latent step and with/without state-conditional encoder.
This not surprising because better prediction accuracy can naturally lead to better performance. 2) For both state-condition and unconditional decoder, search values are positively correlated to the errors. This indicates extra value improvement might be caused by over-optimistic prediction. 3) Both search value and errors are minimal at I, = 3. Our interpretation about this is: Smaller T, can lead to overfitting in the training data, therefore worse generalization and finer-grained search, which will amplify gencralization error. On the other hand, larger L can cause underfitting therefore higher error. 4) Naively increasing the number of latent variables is not helpful for the unconditional decoder.
Table 9: Comparing beam search against direct sampling from the prior policy.
Dataset
Environment
beam search
sample prior
sample uniform
Medium-Fxpert
HalfChectah
91.8
89.9
41.8
Medium-Fxx Expert
Hopper
105.5
98.5
62.3
Medium- Medium-Expert
Walker2d
107.4
107.7
86.7
Medium- Expert
Ant
128.8
124.7
105.4
Medium
HalfChcetah
45.0
44.3
39.5
Medium
Hopper
63.4
64.3
39.6
Medium
Walker2d
64.9
55.5
70.2
Medium
Ant
92.0
88.8
89.8
Medium-Rephee Replay
HalfCheetah
10.8
39.8
10.2
Medium-Rephh Replay
Hopper
87.3
79.0
14.7
Medium-Replee Replay
Walker2d
66.8
66.0
7.8
Medium Medium-Replay
96.7
81.4
47.0
Mean
82.5
78.3
51.3
I SAMPLED DISTRIBUTIONS
We sampled 2048 trajectories either uniformly or according to the prior in the first step of the 10 evaluation episodes and plot their metrics in Figure 7. The x-axis of the plots is predicted returns by the agents. The y-axis shows the state's prediction mean-squared error (MSE), where the ground truth is given by the simulator following the action sequences of the predicted trajectories. Sampling from the prior does not damage the diversity of the sampled trajectories in terms of the variance of returns. On the other hand, the expected MSE reduced in both hopper and ant cases. This means that by sampling from the prior, we can get higher quality of trajectory samples. We can see that TAP successfully approximates the continuous distribution of returns and the generated trajectories are high in diversity in terms of the returns. This is especially true for hopper-medium- replay because the behaviour policy itself of high diversity, and we can see multiple modes in the distribution. It is interesting to that the trajectories with higher predicted returns do not necessarily have higher prediction errors. This a nice feature because we can search for high- return trajectories without worrying too much about these trajectories being unrealistic.
21

---

Published as a conference paper at ICLR 2023
Table 10: State unconditional decoder with different length steps.
Dataset
Environment
L L=1
==
L 3
Medium- Expeet
HalfCheetah
10.0
11.0
9.5
Medium-Exper
Hopper
27.8
29.6
42.6
Bediumm-Expee
Walker2d
50.0
46.9
15.7
Medium Modium-Expert Expert
Ant
32.5
30.9
26.1
Medium
HalfCheetah
18.2
16.5
13.6
Medium
Hopper
37.0
39.2
47.0
Medium
Walker2d
13.6
17.3
24.8
Medium
Ant
18.5
21.5
19.7
Medium-Replay
HalfCheetah
4.8
5.7
5.0
Medium-Ree Repllay
Hopper
19.6
16.6
22.6
Medium- Replay
Walker2d
8.2
8.9
9.6
Medium- Replay
Ant
20.0
24.2
21.1
Mean
21.7
22.4
21.4
1.75
1.50
2.5
25-
15
0.2
02
6.
0.00
-20
-100
0
100
200
0.00
-20
-100
100
200
08
1000
700
200
300
08
1000
300
209
300
rnturn
inturn
(a) hopper uniform
(b) hopper prior
(c) ant uniform
(d) ant prior
Figure 7: Distribution of trajectory returns and state prediction mcan-squared errors. The dataset for the hopper is medium-replay and that for ant is medium.
In the ablation of sampling from prior versus sampling from a uniform distribution, we saw sampling from the prior policy gives much better performance. Comparing the distribution of returns and emror of trajectories sampled from the prior and uniform distribution, we observe the sampling from prior policy yields trajectories with lower reconstruction error but with similar returns. This indicates sampling from prior will improve the quality of trajectories being generated from TAP, which explains better performance.
J TRAINING COST
The training time of TAP and TT are similar for lower dimensionality, TT needs 6-12 hours and TAP needs hours for gym locomotion control tasks, a single GPU. It's worth noting that the training cost of TT will also grow quickly along the dimensionality. For adroit, the same training for TT takes 31 hours, but the training cost TAP is the same given the same architecture. In fact, for adroit experiments with TAP, we used smaller network (see Table 7 for needs fewer training epochs and the whole training can finished within hour.
which
K VISUALIZE PLANS
To give intuitive understanding of the latent action space, predicted trajectories and planning of TAP. In we show a visualization of the generated plans by both direct sampling and beam search with 256 samples (beam width=64 and expansion factor=4 for beam search). Each frame shows 256 latent codes and the corresponding state in a particular step in the plan, where all the plans start from an initial state of the hopper task. The trajectories sorted according the objective score,
22

---

Published as a conference paper at ICLR 2023
(a) sample step 20
sample step 80
sample step 140
(d) beam step 20
(c) heam step 80
(f) beam step 140
Figure 8: Visualization of latent codes, decoded trajectories and planning.
so that the most front-end trajectory will be the plan to be executed. Direct sampling generates more diverse trajectories while most of them are of low prior probability. There are some the predicted trajectories (opaque) that move quickly but do not follow the environment dynamics. The plan that chosen follows the true dynamics but is suboptimal as it's falling down. On the other hand, beam search generates trajectories that both with high returns and high probability. The model predicts 144 steps of the future and trained on hopper- hopper-medium-rrplaa
23


Reviews:
[{"review_id": "oTeUTfmBfV", "reviewer": "Reviewer_Ykqw", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper describes a model-based reinforcement learning method designed to learn from offline datasets (generated by an external behavior policy). The method attempts to address two key challenges in model-based RL. First, for common continuous control problems such as locomotion and manipulation, the action space is high-dimensional and fine-grained, resulting in an expensive search problem. Second, learned environment models are often riddled with “holes” (regions with poor dynamics & cost estimates) due to uneven data coverage. Uniform trajectory sampling strategies may easily fall into these holes and yield poor performance. The proposed method addresses the first problem by learning a latent action space, where each discrete latent action maps to a short snippet of trajectory. The computational cost of searching in such action space can be adjusted using the trajectory snippet length. The latent action space is discrete and finite and is learned end-to-end with a VQ-VAE (VAE with discrete priors) trajectory reconstruction objective. Such data-informed action space also partially addresses the second problem (holes), as the search is naturally confined to trajectory patterns that exist in the training set. To further tighten the search space, the method proposes to learn the likelihood of the next latent action conditioning on the first state and the latent actions in the plan so far. The method also employs beam search to further improve the planning efficiency. The method is evaluated on a simulated offline RL benchmark (D4RL), which includes continuous control problems such as half-cheetah, and tasks with higher dimensional action space (dexterous manipulation). The result shows that the proposed method achieves comparable performances relative to a number of model-based and model-free RL baselines for tasks with low-dimension actions, and compares favorably in tasks with higher-dimensional action space. ", "summary_of_the_review": "Again, I like the overall premise and the high-level idea of the work. And I think the method is conceptually simple and would be easy to work with. I'd like to see an in-depth discussion of missing references listed above and also addressing the gap between learning environment models and learning behaviors. More complex evaluation tasks can also substantially improve the overall quality of the paper.", "strength_and_weaknesses": "Overall I quite like the overall premise of the paper: if learning and planning with an accurate and comprehensive model for a control problem is hopelessly complex and expensive, why not focus on the part of the problem that is relevant to the task? I think learning a latent action space and building a constrained dynamics model on top is a nice way to approach this problem. And the proposed method is conceptually simple and easy to dissect. I also like how comprehensive the ablation study is. \n\nHowever,  I believe there are important issues with this paper that need to be addressed. The issues are (1) missing references that are highly relevant and (2) conceptual and empirical differences between learning models and learning behaviors.\n\n1. Missing references\nThere are a number of important references that are missing. First, there is an entire set of work in hierarchical imitation learning [1][2] that tries to learn temporally-abstracted behavior primitives and use them for planning / more efficient imitation. For example, Mandlekar et al., [1] uses VAE to learn a latent space of short-horizon goal-conditional policies and use that for value-based planning. Similarly, [2] learns a latent space of short behaviors and uses them for one-shot imitation. From a practical point of view, these behavior primitives and how they are used are no different than the latent action proposed in this work. In the context of more RL-ish works, Allshire and Martin-Martin et al. [3] is highly relevant. Similar to the premise and motivation of this paper, they learn a latent action space based on offline experience and use it for more efficient RL. There is also [4], which has similar line of argument. Although they do not study temporal abstraction, it is worth discussing and contextualizing the proposed method with this work. Finally, in a slightly far away domain, a number of works [5][6] proposed to learn temporally-abstract action space for better teleoperation. I hope the authors can select a subset of these works and their predecessors / followup works and provide a more thorough account of the idea of learning task-relevant action space, instead of just focusing on offline & model-based RL. In addition, some of these works may be directly comparable and may serve as strong empirical baselines, e.g., [1]. \n\n[1] IRIS: Implicit Reinforcement without Interaction at Scale for Learning Control from Offline Robot Manipulation Data, Mandlekar et al., 2019\n[2] Learning latent plans from play data, Lynch et al., 2019\n[3] LASER: Learning a Latent Action Space for Efficient Reinforcement Learning, Allshire and Martin-Martin et al, 2021\n[4] PLAS: Latent Action Space for Offline Reinforcement Learning, Zhou et al., 2020\n[5] Learning latent actions to control assistive robots, Losey et al., 2022\n[6] Learning visually guided latent actions for assistive teleoperation, Karamcheti et al., 2021\n\n2. Learning models vs. learning behaviors.\nAs mentioned in my synopsis of this work, I believe the proposed method is more or less limited to learning from offline datasets that is generated by an external behavior policy, since it assumes there exists such a task-relevant state-action space to learn the latent action. So how does one differentiate the proposed “model-based offline RL” methods from conditional imitation learning? Based on mainstream understanding, the main advantage of model-based approach is that the model can be shared among many different tasks, as the per-step dynamics are agnostic to a task goal. And one may use the same model to synthesize behaviors that can achieve new goals in a zero-shot manner. It’d be interesting to see how much of the “model-based” quality the proposed method still retains. Concretely, I’d like to see experiments that focus on generalization to new goals and behaviors.\n\nMy final critique that is not as important as (1) and (2) is the lack of more complex evaluation task. Since one of the main argument that the paper makes is computational efficiency, I'd really wish to see some real-robot experiment to substantiate this claim. Along the same line, it'd be great if the authors could at least evaluate on simulated benchmark such as robomimic (Mandlekar et al., 2021) for which the tasks are directly modeled after setup that can be replicated in the physical world. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "OC8Jrr5WTQd", "reviewer": "Reviewer_rxeb", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The authors extend the Trajectory Transformer (TT) by learning a latent discrete action-space consisting of trajectory segments, which is used to speed up (beam-search) planning for offline RL. The authors argue that planning in this latent space scales better with dimensionality of the action space. Additionally, the proposed architecture is computationally cheaper than the TT.", "summary_of_the_review": "This seems like a good paper, presenting novel and relevant ideas for learning a discrete action space in offline RL. I think their analysis of the learned representations (and in particular how much model accuracy vs. compactness contributes the improved success rate) could have gone a bit further, but it does appear sound, results in improved or similar results compared to earlier work and has a large ablation study already.\n\nEDIT: I just noticed that some seemingly relevant but unlisted related work surfaced in the other reviews. I will adjust my score based on the outcome of that discussion.", "strength_and_weaknesses": "Strengths:\n- The idea of learning a latent discrete action space is very interesting and highly relevant\n- The approach appears sound and novel\n- The paper is well-written with a thorough ablation study\n\nWeaknesses:\n- The experiment tasks, while common in this type of paper, are rather simple. The locomotion problems are not complex enough to show improvement, and for the robotic hand experiments, the optmimal length L of the latent trajectory segments seems rather small (3). This makes it a bit unclear how much the learned latent representation helps with planning compared to just 1) reducing the time discretization of the TT by a factor of 1/L (i.e. executing the same action for L=3 steps in a row), and 2) manually discretizing the action space (e.g. clustering it using k-means), although for a fair comparison one would need to sample from a learned conditional sequence prior (or ignore it for both) 3) both of these in combination. I do not think this was covered by the existing baselines? Would 1) be feasible to add?\n- It would have been interesting with some introspection and qualitative examples of the learned latent action representation (trajectory segments). This approach should be more suitable for tasks that have a natural discrete structure. It would have been interesting to see how well it could recover that.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "LCQ9j-PFA9", "reviewer": "Reviewer_Brj7", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "To alleviate the computational overhead of the planning-based RL algorithms, this paper proposes an Offline RL method Trajectory Autoencoding Planner (TAP), which learns and plans in the compact latent action space. \nThe key points are: \n  * (1) TAP achieves both *spatial abstraction* and *temporal abstraction* by applying a Transformer based VQ-VAE to encode each transition tuples $(s_t, a_t, r_t, R_t)$ in a trajectory $\\tau$ into a latent space and then merging every $L$ adjacent encoded transtions into one latent action in the discrete codebook of VQ-VAE.\n  * (2)  TAP uses a Transformer to autoregressively model the prior distribution of the latent action codes. \n  * (3) The planning is executed in the discrete latent action space by applying the trained prior latent action sampler and beam search. As multiple environmental actions are aggregated into one latent action, the planning horizon is reduced to $1/L$.\n", "summary_of_the_review": "Overall, this paper proposes an interesting idea of learning and planning in the temporal abstracted latent action space, which significantly reduces the computational cost compared to its counterpart Trajectory Transformer. I recommend an accept.", "strength_and_weaknesses": "**Strength:**\n  *  Using Transformer based VQ-VAE to achieve both *spatial abstraction* and *temporal abstraction* can significantly reduced the search space of the planning based RL method.\n  * Compared with Trajectory Transformer, considering each transition tuple $(s_t, a_t, r_t, R_t)$ as a single token for the Transformer is much more reasonable, which could significantly reduce the computational cost.\n  * The paper is written clearly and easy to follow.\n  \n\n**Weaknesses:**\n  * While using compact action representation could reduce the search space, the model also losses the ability of doing more fine-grained control and can potentially lead to sub-optimal action selection. How to balance the two?\n  * How to apply TAP to the online RL setting where we do not have lots of and diverse existing data to train the action encoder?\n * As the authors also mentioned in the limitation part, TAP currently cannot not distinguish between uncertainty stemming from a stochastic behaviour policy and uncertainty from the environment dynamics. \n  * It's better to provide the overall training loss function as well.\n  * [1] may be a missing related work.\n  * Dreamer V2 [2] also use discrete latent representations. It's better to discuss connection between the two.\n  * Minors: \n    * \"After adding the positional embedding, the encoder then reconstructs the trajectory\",  the 'encoder'  should be 'decoder'.\n    * \"As such, the planning process only needs to optimize in a 5 or 8-dimensional discrete space with a learned prior\", the 'discrete space' should be 'planning horizon'.\n\n\n\n**Reference**\n\n[1] Continuous Control with Action Quantization from Demonstrations\n\n[2] Mastering Atari with Discrete World Models", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster