Venue: ICLR 2023 Conference

Reviews:

final_score: rating: 8: accept, good paper
review_content: 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.
strength_and_weaknesses: Strengths:
- The idea of learning a latent discrete action space is very interesting and highly relevant
- The approach appears sound and novel
- The paper is well-written with a thorough ablation study

Weaknesses:
- 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?
- 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.

clarity,_quality,_novelty_and_reproducibility: The paper is mostly well written, novel and appears reproducible.

Some things that could be improved:
- Framing: The intro implies planning-based RL has only shown strong performance on low-dim problems, but as e.g. Go, Chess are quite high-dimensional, this claim seems questionable. I would consider rephrasing that, I think the paper can make its point about improved efficiency in high-dim planning without it. 

- How much data was actually used to train the offline model in the experiments? Has to be quite a bit?

- The paper is simultaneously also talking about RL success and compute, but unless the model is actually better (and not just more compact) RL success for earlier approaches should also depend on the compute spent on sampling the (less compact) space? I feel these aspects could have been better examined separately, but I realize space is at a premium.  

Some minor language issues:
- "future trajectories estimated by a dynamics model" -> sampled from?
- The second paragraph of the intro is a bit wordy, and many arguments seem to overlap. I think it could be condensed a bit.
- I didn't see the VQ-VAE paper cited. The explanation of VQ-VAE was also terse to the point of being mysterious.
- Typo on p.4: "probabiilty"

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.

EDIT: 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.
review_title: null

final_score: rating: 6: marginally above the acceptance threshold
review_content: 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. 
The key points are: 
  * (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.
  * (2)  TAP uses a Transformer to autoregressively model the prior distribution of the latent action codes. 
  * (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$.

strength_and_weaknesses: **Strength:**
  *  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.
  * 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.
  * The paper is written clearly and easy to follow.
  

**Weaknesses:**
  * 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?
  * 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?
 * 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. 
  * It's better to provide the overall training loss function as well.
  * [1] may be a missing related work.
  * Dreamer V2 [2] also use discrete latent representations. It's better to discuss connection between the two.
  * Minors: 
    * "After adding the positional embedding, the encoder then reconstructs the trajectory",  the 'encoder'  should be 'decoder'.
    * "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'.



**Reference**

[1] Continuous Control with Action Quantization from Demonstrations

[2] Mastering Atari with Discrete World Models
clarity,_quality,_novelty_and_reproducibility: * The paper is written clearly and easy to follow.
* The idea of learning and planning in the temporal abstracted latent action space is novel. Applying the overall method to Offline RL setting is reasonable.
* The code is also attached in the supplementary material. So, I guess the experimental results can be reproduced properly.
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.
review_title: null

final_score: rating: 5: marginally below the acceptance threshold
review_content: 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. 
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. 

However,  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.

1. Missing references
There 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]. 

[1] IRIS: Implicit Reinforcement without Interaction at Scale for Learning Control from Offline Robot Manipulation Data, Mandlekar et al., 2019
[2] Learning latent plans from play data, Lynch et al., 2019
[3] LASER: Learning a Latent Action Space for Efficient Reinforcement Learning, Allshire and Martin-Martin et al, 2021
[4] PLAS: Latent Action Space for Offline Reinforcement Learning, Zhou et al., 2020
[5] Learning latent actions to control assistive robots, Losey et al., 2022
[6] Learning visually guided latent actions for assistive teleoperation, Karamcheti et al., 2021

2. Learning models vs. learning behaviors.
As 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.

My 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. 
clarity,_quality,_novelty_and_reproducibility: Clarity: very clear. great writing! 

Quality: solid execution of an intuitive idea, although the evaluation task complexity can be improved. 

Originality: somewhat questionable. See the weakness comments on missing references.
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.
review_title: null

Meta-review: metareview:_summary,_strengths_and_weaknesses: Summary:
The paper extends trajectory transformers with a latent action encoding for model-based / planning-based RL. The approach is evaluated in two benchmark suits.

Strengths:
- Interesting approach that works well
- Well written paper
- Most concerns of the reviewers have been successfully addressed in the discussion (also of the one that didn't reply)

Weaknesses:
- The locomotion problems add only limited insights
- The general idea of action abstractions is well known, this particular setting seems new
justification_for_why_not_higher_score: The paper builds on well known ideas and is addressing a somewhat narrow problem.
justification_for_why_not_lower_score: A solid paper, the most negative reviewer did not react but I think the concerns have been well addressed.
Decision: Accept: poster