Meta-Controller: Few-Shot Imitation of Unseen Embodiments and Tasks in Continuous Control

Q1. Given that state representation learning is decoupled from policy representation learning, it would be interesting to visualize the embedding learned by the state encoder and see how distinct it is between different morphologies.

We provide a t-SNE visualization of embedding space of features obtained by structure encoder in Figure R.1 of author rebuttal. First, we can observe that the embeddings are clustered by the joints of each embodiment, and the clusters corresponding the same embodiment are located nearby. Also, we note that the embeddings of slide joints (e.g., cartpole, cartpole_two, cup, pointmass) and the embeddings of hinge joints (e.g., reacher, reacher_three, walker, acrobot, cheetah, pendulum) are separated in the right and the left regions. Thus the state encoder captures both embodiment-specific and joint-specific knowledge to give rich features to the policy network.

Q2. Does the performance of the model scale with model capacity and data?

To address the reviewer’s question, we conducted additional experiments to validate the effect of data scale and how different compositions of meta-training dataset affects the performance of downstream behavior cloning. Specifically, we select 3 combinations of training tasks, where we remove 4 embodiments from the original 10 embodiments. Then, we performed 5-shot behavior cloning experiments on the 8 tasks presented in the main table.

According to Table R.4 of the author rebuttal, we observe that using all embodiments outperforms the baselines in most tasks, indicating that a diverse set of embodiments makes the model robust to unseen embodiments. Our findings also show that it is crucial to include embodiments with morphology and dynamics similar to the downstream tasks in the meta-training dataset. For example, removing the reacher_three task (as seen in row 2 and row 3) significantly drops the performance of the reacher_four task. This result reveals that embodiments with similar dynamics or morphological features can facilitate more effective knowledge transfer during meta-testing, suggesting that the diversity of data greatly impacts performance.

W1. Only demonstrates results in DMC. The reviewer is not certain if this scales to real-world control environments, as the different embodiments operate at various scales, and the model needs to be able to be robust to environment dynamics (here the environment is static).

Our experiments were conducted exclusively within the DeepMind Control Suite (DMC) because, to our knowledge, it is the only benchmark currently available that includes both diverse embodiments and tasks. Although our experiments lack real-world validation, we clarify that our primary contribution is the introduction of a fundamental framework for simultaneous generalization to unseen embodiments and tasks with few demonstrations, which we believe addresses a significant challenge in the field.

To address concerns about the transferability of our approach to real-world control environments, we have extended our experiments to simulate more realistic conditions. Specifically, we introduced varying levels of noise to the control actions within the DMC environments, thereby making the transition dynamics both stochastic and noisy. As illustrated in Figure R.5 of the author rebuttal, the performance of our Meta-Controller does not significantly drop despite the presence of random noise. This result underscores the potential of our method to operate effectively under realistic conditions that involve complex dynamics and noise.

Q1. What is the learning curve for the models compared with PDT+PEFT?

We provide learning curves of Meta-Controller and PDT+PEFT in Figure R.2 of the author rebuttal. As the reviewer pointed out, due to the modular nature of our structure encoder, our model not only achieves better performance but also converges much more quickly than PDT+PEFT in every task.

Q2. Performance when given fewer demonstrations?

We present 3-shot behavior cloning results of baseline models and Meta-Controller in Table R.3 of the author rebuttal. We observe that our model outperforms all the baselines in most of the tasks, demonstrating its adaptability on unseen embodiments and tasks in a few-shot regime.

Q3. The episode sampling strategy is confusing.

Our meta-training dataset consists of a replay buffer of an expert agent, collected throughout the agent’s learning period, which means the behavior or policy in the buffer can vary significantly. Initially, the expert agent's policy is sub-optimal and it improves over time. We observe that this leads to diverse policies for a single task. Our matching framework, however, requires consistency between input states and output actions in both support and query data.

To ensure this consistency, we sample episodes from similar training periods within the replay buffer. For instance, in the hopper-stand task with 1000 trajectories sorted by training period, we select a query trajectory (e.g., the 100th) and support trajectories within a neighboring range of size 10 (i.e., from the 95th to the 105th). The “temporal segment size 10” means this range. This ensures that the support and query data are temporally close, providing consistent and relevant task information.

W1. The models are complex and matching-based policies are time-consuming.

We clarify that transformer architectures are commonly adopted in the literature, and the complexity of our architecture would not be a significant weakness. Notably, recent works in the field use much larger transformer-based architectures than ours, such as large language models (LLMs) [3, 5] and vision transformers (ViTs) [4, 5]. In Table R.5 of the author rebuttal, we compare the inference time of our model and VC-1 [4], which uses a ViT-L backbone. Our model achieves faster inference times than VC-1. Although we couldn’t evaluate RT-2 [5] due to lack of official code, it involves a large network with at least a 40-layer ViT and a LLM with 3B parameters, requiring much higher computation costs compared to ours.

Regarding the matching process, it is implemented with a single cross-attention layer, constituting only a small part of the overall computation. As shown in Table R.6 in the author rebuttal, most inference time is occupied by the encoders and decoders. In resource-constrained robotic platforms, transformer-based architectures may hinder real-time inference, as discussed in Appendix A. However, advancements in optimizing transformers, such as sparse attention mechanisms, model pruning, and knowledge distillation, can enhance inference speeds. Since these techniques are orthogonal to our contribution, they can be naturally incorporated into our method for resource-constrained robotic platforms, although this is not the main focus of this paper.

W2. The environments are too simple, only focuses state-based environments. Performance of image-based or manipulation tasks?

First, we clarify that our primary contribution is demonstrating simultaneous generalization to unseen embodiments and tasks using few-shot learning, which is a significant challenge in the field. This necessitated using the DeepMind Control Suite (DMC), the only dataset providing both diverse embodiments and tasks to our knowledge. Other datasets are limited in either range of embodiments [1] or tasks [2], or both. Although DMC might appear simple, it enables rigorous testing of our framework's core capabilities. Similar research, such as PromptDT and MetaMorph, used even simpler settings. To our knowledge, our setting represents one of the most challenging for few-shot out-of-distribution generalization in behavior cloning.

Second, our method is inherently designed for state-based environments, where "state" refers to proprioceptive sensors attached to each joint. Extending our method to image-based environments is beyond the scope of this paper. It is common in the literature to focus on state-based environments [2, 7], as many real-world robots are equipped with proprioceptive sensors.

Regarding "manipulation tasks," we interpret this to refer to 3D robotic arm manipulation tasks, which are also state-based environments. If our interpretation is incorrect, please let us know. Conducting experiments in these environments is feasible but challenging due to the limited availability of diverse and comprehensive 3D datasets. Most available datasets focus on specific joint compositions like robotic arms or quadrupeds, making it difficult to gather sufficient meta-training data to acquire transferable knowledge about 3D embodiments. This limitation reflects the current state of available resources rather than the capability of our approach.

Nonetheless, we recognize the importance of demonstrating the generalizability of our approach in more complex settings, and additional experiments in the author rebuttal may give some insights to this direction. Table R.2 shows that our method still performs well in more diversified embodiments. In Table R.5, we demonstrate that our method is robust to noise, highlighting its potential in noisy and stochastic real-world scenarios. Finally, as shown in Table R.4, our method's performance improves when we add more diverse embodiments and tasks in the meta-training dataset. This suggests the potential of our method in more complex scenarios, such as those with noisy transition dynamics or emerging diverse variants of embodiments, as more diverse datasets become available.

Q1. The ablation study in Table 2 is incomplete.

Table R.1 of the author rebuttal completes the ablation study as requested (see row 1 and 2). Consistent with the discussion in Section 5.3 of the main text, we observe that removing the structure encoder $f_s$​ significantly decreases performance on unseen embodiments and tasks. This result indicates that the structure encoder captures transferable and modular knowledge about various morphologies. Additionally, removing all three modules yields a fully supervised model that does not involve meta-learning. As expected, this model fails to adapt to few demonstrations and suffers from overfitting.

Q2 & W1. It raises the question of whether the policy would still be effective with variations in embodiment configurations.

We conducted experiments on six additional embodiments by manually changing joint lengths (e.g., foot length of hopper) or ratios among joints (e.g., calf-thigh ratio of hopper, front leg-back leg ratio of wolf). These modifications can make the tasks harder, as the original embodiments are optimized for specific tasks. In Table R.2 of the author rebuttal, we compare Meta-Controller with the high-performing baselines from Table 1 of the main text. Our model consistently outperforms all baselines in these challenging variants, demonstrating the robustness and adaptability of our approach to variations in embodiment configurations.

Q3. A robustness analysis would further validate the generalization ability of the proposed method.

We conducted an additional experiment by introducing noise to the transition dynamics. Random noise sampled from $\mathcal{U}[-n, n]$ was added to the agent's action at each timestep, with three noise levels $n \in$ $[2$%, $5$%, $10$%$]$ of the action range. Figure R.5 in the author rebuttal plots the rewards of our model at each noise level compared to experts. The results indicate that our method maintains its performance across many tasks as noise levels increase. This shows our model’s robustness under stochastic environments. Interestingly, for tasks such as reacher-four, the performance increases with higher noise levels, likely due to the exploration effect induced by stochastic transitions.

W2. The proposed method operates only in a 2D coordinate space, limiting its application potential in realistic 3D environments.

We clarify that our Meta-Controller, including the structure-motion state encoder and matching-based policy network, is designed to be applicable to any coordinate space. In principle, we can extend the input embedding layer of the state encoder to accommodate 3D coordinate inputs.

The primary challenge in demonstrating our method in 3D environments is the limited availability of diverse embodiments and tasks in existing reinforcement learning datasets. Since our goal is to achieve simultaneous generalization to unseen embodiments and tasks, we necessiate a meta-training dataset with diverse embodiments and tasks. This leads us to adopt the DMC dataset, which, to the best of our knowledge, the only dataset that meets the condition (other datasets are composed of either limited embodiments [1], limited tasks [2], or both) and is also widely used in the literature. Although DMC includes some 3D embodiments, their diversity is too limited to cover joint compositions in 3D coordinate space (only 4 unique 3D embodiments).

To demonstrate our method's scalability with diverse embodiments and tasks, we conducted additional experiments varying the size of the meta-training data. As shown in Table R.4 of the author rebuttal, our method's performance improves in general when adding more diverse embodiments and tasks in the meta-training dataset. This suggests the potential of our method in more realistic environments as more diverse datasets become available, which will be an important step forward.

While we have not conducted direct experiments in 3D environments, we believe our key findings in 2D embodiments—namely, the generalization capability to unseen embodiments and tasks—will extend to 3D embodiments as well.

W3. The multiple transformer-based architecture constrains the possibility of real-time inference.

We clarify that our primary focus is not on real-time inference but rather on presenting a fundamental framework for achieving robust few-shot learning across various unseen embodiments and tasks, which, to our best knowledge, has not been addressed in previous work.

Recent works often adopt transformer-based architectures much larger than ours, such as large language models (LLMs) [3, 5] and vision transformers (ViTs) [4, 5]. In Table R.5 of the author rebuttal, we compare the inference time of our model with VC-1 [4], which uses a ViT-L backbone. Our model achieves faster inference times than VC-1. Although we couldn't evaluate RT-2 [5] due to lack of official code, it involves a large network with at least a 40-layer ViT and a LLM with 3B parameters, requiring much higher computation costs. However, these works have demonstrated the feasibility of using transformer architectures in real-time applications. For example, RT-2 achieved real-time inference with their heavy network using cloud computing. Thus employing multiple transformer-based architecture would not be an essential limitation for real-time inference.

In resource-constrained robotic platforms, transformer-based architectures may hinder real-time inference, as discussed in Appendix A. However, advancements in optimizing transformers, such as sparse attention mechanisms, model pruning, and knowledge distillation, can enhance inference speeds. Since these approaches are orthogonal to our contribution, our method can naturally incorporate them for real-time applications with minimal requirements of the computing resource, although this is not the main focus of this paper.

We respectfully disagree with the reviewer. There have been extensive community efforts in building foundation models for continuous control [11, 12, 13, 14], given their significant potential to create versatile robotic agents that can be utilized in various real-world applications. Since generalization capability inherently stems from the scale of the model and training data, most (if not all) approaches are built upon large-scale Transformer backbones. Although these models may not be advantageous for real-time operation compared to lightweight models, we believe that the value of these efforts should not be underestimated simply due to computational cost, especially when considering their capabilities and significance for general-purpose robots.

Additionally, there have always been parallel efforts in the machine learning community to develop more capable models and to accelerate their computation. As discussed in our rebuttal, there is an extensive body of work aimed at reducing the computational demands of Transformers, including techniques such as quantization[15], pruning[16], distillation[17], and linear attention[18], as well as methods to enhance hardware utilization, such as FlashAttention[19, 20, 21]. These approaches are orthogonal to our method and generally applicable to any Transformer backbone. Moreover, considering ongoing developments in hardware that rapidly reduce the cost of computation and enhance the applicability of large-scale models, we believe it is important to refrain from judging the real-time applicability based solely on current computational limitations.

[11] Janner et al., “Offline Reinforcement Learning as One Big Sequence Modeling Problem”, NeurIPS, 2021

[12] Sun et al., “SMART: SELF-SUPERVISED MULTI-TASK PRETRAINING WITH CONTROL TRANSFORMERS”, ICLR, 2023

[13] Chen et al., “Decision Transformer: Reinforcement Learning via Sequence Modeling”, NeurIPS, 2021

[14] Liu et al., “Masked Autoencoding for Scalable and Generalizable Decision Making”, NeurIPS, 2022

[15] Liu et al., “Post-Training Quantization for Vision Transformer”, NeurIPS, 2021

[16] Kim, Sehoon, et al. "Learned token pruning for transformers." Proceedings of the 28th ACM SIGKDD Conference on Knowledge Discovery and Data Mining. 2022.

[17] Sanh, V. "DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter." Proceedings of Thirty-third Conference on Neural Information Processing Systems (NIPS2019). 2019.

[18] Choromanski, Krzysztof Marcin, et al. "Rethinking Attention with Performers." International Conference on Learning Representations. 2020.

[19] Dao, Tri, et al. "Flashattention: Fast and memory-efficient exact attention with io-awareness." Advances in Neural Information Processing Systems 35 (2022): 16344-16359.

[20] Dao, Tri. "FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning." The Twelfth International Conference on Learning Representations.

[21] Shah, Jay, et al. "Flashattention-3: Fast and accurate attention with asynchrony and low-precision." arXiv preprint arXiv:2407.08608 (2024).

Your response shows your unprofessionalism and lack of seriousness.

• The third and fourth paragraphs in the response appear to be rephrasings of the first two paragraphs.
• The mentioned works [11, 12, 13, 14] primarily aim to address general decision-making problems. In contrast, the focus of this submission is on the generalization of robotic control systems across different embodiments and tasks. The starting point of different problems implies distinct requirements. While solving general decision-making problems may prioritize model effectiveness on wide-range tasks over computational cost and real-time performance, the same leniency cannot be applied to robotic control systems. Ignoring the real-time demands and limited computational resources during system deployment is unprofessional for real-world applications that require precision and reliability. Related works[a, b, c, d] focusing on the similar problem settings have verified the transferability of the model on real robots. Why can't your work do it?
• For robotic control systems, the need for real-time performance in real-world robots is paramount. The paper's oversight in addressing the computational requirements and real-time capabilities necessary for practical deployment is a significant shortcoming. This is especially critical given the diverse hardware systems that entail varying control frequencies and dynamic responses, which should be integral to model design considerations. The techniques[15~21] are potential solutions to reduce the computational demands of Transformers but have not been verified on similar embodied control tasks yet.

In summary, the methodology presented in this paper appears to be under-considered and lacks depth in addressing the complexities of the problem at hand. It does not seem to have delved into the nuances of robotic control systems and their requirements for continuous control. Given the above concerns, I believe that the paper does not meet the rigorous standards required for publication in NeurIPS. It is essential for robot learning research to not only present new ideas but also to thoroughly consider and address the practical implications and challenges associated with real-world applications.

Ref:

[a] Mirage: Cross-Embodiment Zero-Shot Policy Transfer with Cross-Painting

[b] Cross-Embodiment Robot Manipulation Skill Transfer using Latent Space Alignment

[c] Polybot: Training One Policy Across Robots While Embracing Variability

[d] RoboDuet: A Framework Affording Mobile-Manipulation and Cross-Embodiment

Q1. In the case of robot arm manipulation, what are the benefits of utilizing the proposed structure-motion encoder over the direct application of end-effector-based control?

Compared to end-effector-based control, our method eliminates the need for separate low-level controllers manually tuned for each robot and simplifies the control process by learning a unified representation adaptable to various tasks and embodiments. By leveraging joint-level representations and the structure-motion state encoder, our framework seamlessly adapts to different robot morphologies and dynamics, ensuring robust performance across a wide range of embodiments. This is crucial for developing versatile and adaptive robotic systems.

Q2. Could you clarify what is meant by the "Embodiment-specific positional embedding" as depicted in Figure 2?

As explained in Section 3.2.1 (line 132-147), the "Embodiment-specific positional embedding" refers to a mechanism to incorporate specific information about the robot's physical configuration into the model. This embedding accounts for the unique positional relationships and characteristics of each joint within a given robot embodiment. Similar to positional encoding in transformers, it is implemented by adding a learnable vector to each joint's state token.

Q3. Why not consider the state features from all joints to simultaneously generate motions for each joint?

The architecture design of our state encoder allows the structure encoder to contextualize the joint representation considering all joints, while the motion encoder captures their temporal dynamics. By applying these two encoders in sequence, we can encode motion features that incorporate both spatial and temporal information, in a more computationally efficient way than using full spatio-temporal interactions. This approach is similar to axial attention techniques [6] used in video models, which are widely adopted for their efficiency and effectiveness in handling spatio-temporal data.

Q4. What necessitates the use of an action encoder and decoder? W2. The proposed method appears to share conceptual connections to some retrieval-based techniques.

While matching is effective for few-shot learning, applying it directly to the raw action space imposes a strong constraint: the output action must be a convex combination of the support actions. To alleviate this constraint, we introduce an action encoder and decoder, allowing matching in the latent space rather than the raw label space. This enhances the expressiveness of the policy network, enabling adaptation to various unseen tasks with non-convex relationships between states and actions. Additionally, by encoding actions along the temporal axis, the model can construct a pool of transferrable action features related to local motor skills, which facilitates efficient transfer to unseen tasks that share modular skills but involve different skill combinations.

We note that retrieval-based techniques suggested by the reviewer are similar to our matching-based policy unless they interpolate (or hardly select) the raw actions. To validate matching in the latent space, we conducted an ablation study by removing action encoder and decoder from our framework. As shown in Table R.1 in the author rebuttal, removing the action encoder and decoder decreases performance, supporting their effectiveness.

W1. The proposed method assumes the number of joints is equal to the action dimension.

We clarify that our method can address the cases where the number of joints is not equal to the action dimension. As described in Section 3.1 (line 118-119), we assign zero values for free (non-actuable or passive) joints, and simply discard their tokens after encoding via the state encoder. This is because the states of passive joints provide useful information about the morphology of the embodiment during the encoding process, but there are no actions to predict for these joints. Note that most of the embodiments used in our experiment (e.g., hopper, walker, cheetah, acrobot, cartpole, wolf) contains passive joints.

W3. Limited real-world validation.

We clarify that our experiments were conducted exclusively within the DeepMind Control Suite (DMC) because, to our knowledge, it is the only benchmark currently available that includes both diverse embodiments and tasks.

To address concerns about the transferability of our approach to real-world robotic systems, we have extended our experiments to simulate more realistic conditions. Specifically, we introduced varying levels of noise to the control actions within the DMC environments, thereby making the transition dynamics both stochastic and noisy. As illustrated in Figure R.5 of the author rebuttal, the performance of our Meta-Controller does not significantly drops despite the presence of random noise. This result underscores the potential of our method to operate effectively under realistic conditions that involve complex dynamics and noise.

W4. Limited failure case discussion.

To address the reviewer’s concerns, we conducted additional analyses on the failure cases of the Meta-Controller. We present visualizations of scenarios where the Meta-Controller performs less effectively and cumulative rewards over time for each scenario in Figures R.3 and R.4 of the author rebuttal, respectively.

In these failure cases, we observed that agents struggle to obtain rewards until they reach a specific posture. Once they achieve this posture (highlighted by the red box in Figure R.3), they begin to solve the task effectively. This pattern is also reflected in Figure R.4, where the rewards remain near zero until a certain timestep, after which they rise consistently. This result implies that encouraging the agent to reach states similar to those in the demonstrations (e.g., via exploration strategies) would improve performance in challenging few-shot scenarios.