PEAC: Unsupervised Pre-training for Cross-Embodiment Reinforcement Learning

Designing generalizable agents capable of adapting to diverse embodiments has achieved significant attention in Reinforcement Learning (RL), which is critical for deploying RL agents in various real-world applications. Previous Cross-Embodiment RL approaches have focused on transferring knowledge across embodiments within specific tasks. These methods often result in knowledge tightly coupled with those tasks and fail to adequately capture the distinct characteristics of different embodiments. To address this limitation, we introduce the notion of Cross-Embodiment Unsupervised RL (CEURL), which leverages unsupervised learning to enable agents to acquire embodiment-aware and task-agnostic knowledge through online interactions within reward-free environments. We formulate CEURL as a novel Controlled Embodiment Markov Decision Process (CE-MDP) and systematically analyze CEURL's pre-training objectives under CE-MDP. Based on these analyses, we develop a novel algorithm Pre-trained Embodiment-Aware Control (PEAC) for handling CEURL, incorporating an intrinsic reward function specifically designed for cross-embodiment pre-training. PEAC not only provides an intuitive optimization strategy for cross-embodiment pre-training but also can integrate flexibly with existing unsupervised RL methods, facilitating cross-embodiment exploration and skill discovery. Extensive experiments in both simulated (e.g., DMC and Robosuite) and real-world environments (e.g., legged locomotion) demonstrate that PEAC significantly improves adaptation performance and cross-embodiment generalization, demonstrating its effectiveness in overcoming the unique challenges of CEURL. The project page and code are in https://yingchengyang.github.io/ceurl.

Paper

Similar papers

Peer review

Reviewer RG3H6/10 · confidence 4/52024-07-10

Summary

This paper presents a new unsupervised reinforcement learning (URL) algorithm called Pre-trained Embodiment-Aware Control (PEAC). PEAC is designed to tackle cross-embodiment tasks, explicitly considering the influence of different embodiments to facilitate exploration and skill discovery across embodiments. Experimental results demonstrate that PEAC notably enhances adaptation performance.

Strengths

1. The paper is well-written and easy to follow. The intuitive figures and logically structured text make PEAC clear and understandable. 2. The author(s) conduct comprehensive experiments and present a solid work. 3. As embodied intelligence gains increasing attention, cross-embodiment research is indeed an important direction. PEAC undoubtedly contributes to the community in this regard.

Weaknesses

1. Although the author(s) compare various standard and SOTA baselines in the experiments, they lack strong persuasiveness. Since the author(s) focus on cross-embodiment tasks, it is apparent that the baselines, which do not utilize embodiment information $e$, exhibit poorer performance. Therefore, it is unclear whether PEAC's high performance is due to the use of embodiment information $e$ or the utilization of cross-embodiment intrinsic rewards $\mathcal R_{\text{CE}}$. To highlight the paper's contribution, I suggest that the author(s) modify some baselines to naively incorporate embodiment information and then compare with them. For example, conditioning the policy and reward of DIAYN on embodiment $e$, etc. 2. Experiments conducted on Robosuite effectively showcase "cross-embodiment," while other environments may not do so as prominently. By modifying parameters like "mass" and "damping" in DMC tasks or introducing joint torque failures in Unitree A1, the setting becomes more like a Meta RL, where the robot's morphology remains relatively unchanged. In locomotion tasks, "cross-embodiment" could be exemplified by robots like Gym-MuJoCo's Walker2d, DMC's Walker, and Humanoid, where they all use two legs for alternating walking and aim to enhance mobility. However, they differ in morphology. Demonstrating whether PEAC can learn the same shared knowledge of bipedal walking from these diverse morphologies may better illustrate the topic of "cross-embodiment."

Questions

Please see Weaknesses.

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

My primary concern is that the baselines in the experiments do not utilize embodiment information $e$. Please refer to the first point in Weaknesses.

Authorsrebuttal2024-08-08

Thanks a lot for your feedback

Dear reviewer RG3H, Thank you very much for your valuable suggestions and increasing the score. We will try our best to further improve our paper. best regards, Authors

Reviewer xr646/10 · confidence 3/52024-07-12

Summary

This papaer introduces a novel setting, cross-emodiment unsupervised RL, which deals with pre-training good policies in reward free environments across different embodiments in order to perform well on downstream tasks on unseen embodiments. The authors propose the algorithm PEAC for unsupervised learning in this setting. Their method proposes a novel intrinsic reward which is derived from the minmax objective of maximizing the improvement in the downstream fine-tuned policy over the pre-trained policy minus a regularizing KL term, which ensures that the fine-tuned policy stays close to the pre-trained policy, under the worst-case extrinsic reward. From this objective, they derive an intrinsic reward for the unsupervised cross-embodiment setting with maximizes the (fixed) log prior over embodiments minus the (learned) log posterior over embodiments, which encourages the agent to seek areas of the state space that are common accross embodiments, or underexplored. They conduct extensive experiments in both simulated and real environments, including both state-based and image-based DMC tasks, and compare against a large number of pre-existing skill-based and exploratoin based unsupervised RL methods. Their results are impressive, showing improvement in almost all cases over the baselines. Moreover, they conduct real-world experiments in which their algorithm similarly outperforms.

Strengths

- Introduces a novel setting: cross-embodiment unsupervised RL - Convincing motivation - the idea that most cross-embodiment training which focuses on a single task may not learn generalizable embodiment-specific skills rather than task-specific skills - Extensive experiments in both simulation and real-world settings with very strong results compared to baselines

Weaknesses

- The paper is hard to follow and key details on the implementation are not fully explained. For example in line 68 a "embodiment context encoder" is mentioned but I do not see this explained any where else in the paper. How is the discriminator trained? - The extension of the information geometry analyses in Sec 3 do not seem particularly related to the main claims of the paper. This section serves primarly to highlight that policies may be stochastic, however it does not seem that the this information is utilized in the algorithm presented - The stated goal of the paper is to better learn embodiement-aware skills, but the policy is optimized to maximize the expectation over all embodiments. Would this encourage embodiment-aware skills or rather embodiment-agnostic skills?

Questions

1. It is not clear to me how the cross-embodiment policy handles different state spaces. For example in the cheetah-walker experiments in the appendix, how is the policy architecture designed such that it can take in different size state vectors? 2. What does super script c in the MDP M represent? 3. Can you please explain the equality in the last line of Eq 29? 4. It is not clear what is unique about the cross-embodiment setting for unsupervised RL versus the more generic contextual MDP setting. Why do embodiments needed to be treated differently? 5. If log(p(e)) is fixed, does it need to be in the reward?

Rating

6

Confidence

3

Soundness

3

Presentation

2

Contribution

3

Limitations

- As the authors note, their experiments are limited to settings where the embodiments are very similar. It is not clear how their method extends to dissimilar embodiments, particularly with different sized state vectors.

Area Chair YC1z2024-08-12

Reviewer Discussion Needed

Dear Reviewer, The discussion time is coming to an end soon. Please engage in the discussion process which is important to ensure a smooth and fruitful review process. Give notes on what parts of the reviewers responses that have and have not addressed your concerns.

Reviewer xr642024-08-12

Dear authors, thank you for the thorough response and for addressing my questions. The additional experiments on more complicated cross-embodiment settings have alleviated some of my concerns. However, I believe the setting is still very similar to multi-task or meta-RL and it is not clear to me why the problem must be treated as distinct except in the treatment of different state and action spaces and, in this respect, the method of zero-padding is slightly underwhelming. That said, the new experiments certainly make the paper's claims more convincing. I will raise my rating by 1 point.

Authorsrebuttal2024-08-12

Thanks a lot for your feedback and supportive comments

Dear Reviewer xr64: Thank you very much for your constructive feedback and for increasing the score. Below we will answer your questions, especially about the difference between CEURL and multi-task/meta RL. First, we briefly introduce multi-task/meta RL and CEURL: - **Multi-task RL**: train an agent to handle several different tasks, represented by MDPs **with rewards**, at the same time. - **Meta RL**: pre-train an agent in several training tasks, represented by MDPs **with rewards**. The pre-trained agent is required to fast adapt to testing tasks. Here training tasks and testing tasks are sampled from the same distribution [1,2]. - **CEURL**: pre-train an agent with several embodiments **without rewards**. The pre-trained agent is required to fast adapt to any testing tasks. Here we have no prior knowledge of the testing task. (lines 137-146) Consequently, from a modeling perspective, the biggest difference is that **CEURL pre-trains in an unsupervised manner**, i.e., without any extrinsic rewards, but **meta-RL pre-trains with extrinsic rewards**, which obeys the same distribution of the testing tasks. Moreover, from the concept perspective, multi-task/meta RL regards the embodiment and the task as a whole and utilizes an MDP to handle them together. Differently, CEURL distinguishes between the concept of embodiments and tasks. Thus CEURL hopes to learn embodiment-aware and task-agnostic knowledge by pre-training agents only with these embodiments without any tasks (lines 125-136). This knowledge is considered very helpful, especially when embodied intelligence requires handling different tasks across embodiments in the real world. Furthermore, we agree that zero-padding is a feasible but not the best method for handling different state/action spaces. And we believe that developing unified action/state embedding spaces is a promising future direction for handling the cross-embodiment setting. Thanks again for your kind comments and we will add these discussions into the paper to further improve it. Also, we are looking forward to hearing from you about any further feedback. best regards, Authors Reference: [1] Model-Agnostic Meta-Learning for Fast Adaptation of Deep Networks [2] Efficient Off-Policy Meta-Reinforcement Learning via Probabilistic Context Variables

Reviewer 3BHi5/10 · confidence 3/52024-07-12

Summary

This paper addresses the challenge of designing generalizable agents capable of adapting to diverse embodiments. The authors propose the CEURL setting as a novel framework for this problem and introduce the PEAC algorithm to address it. Recognizing that CEURL requires minimizing across different downstream tasks while maximizing the fine-tuned policy, PEAC tackles the issue through policy improvement and policy constraint mechanisms. PEAC can integrate with existing unsupervised reinforcement learning methods designed for single-embodiment settings, such as LBS and DIAYN. Through experiments conducted in both simulation and real-world settings, the effectiveness of PEAC is demonstrated.

Strengths

Originality: This paper introduces a novel problem formulation, CEURL, and presents a new algorithm, PEAC. Quality: 1. The authors thoroughly discuss the challenges inherent in the CEURL problem definition, supported by mathematical proofs. This elevates CEURL from a simple definition to a high-quality problem formulation. 2.Convincing experiments involving real robots are presented. Clarity: The paper is well-written and easy to follow. Significance: This paper has the potential to significantly impact cross-embodiment control through its introduction of a new problem formulation and an innovative unsupervised pretraining framework.

Weaknesses

1. Among all the experiments presented, tasks for DMC and Issacgym are relatively limited, with only simple movements. I recommend adding more complex tasks like [1]. 2. In ablation study, the authors only show the results of different training timesteps. More ablation studies can help demonstrate how the algorithm works. 3. From my perspective, this paper is closely related to [2], which proposes an unsupervised skill discovery method through contrastive learning. While [2] focuses on single-embodiment scenarios, PEAC extends this approach to multiple embodiments, addressing the additional challenges posed by varying state and action spaces via introducing additional embodiment prior and posterior probabilities. [1] Gupta, Agrim, et al. "Embodied intelligence via learning and evolution." Nature communications 12.1 (2021): 5721. [2] Yang, Rushuai, et al. "Behavior contrastive learning for unsupervised skill discovery." International Conference on Machine Learning. PMLR, 2023.

Questions

1. For objective function (2), can you provide an ablation study using only the Policy Improvement part and only the Policy Constraint part? 2. In Figure 6 (right), PEAC-DIAYN shows the worst performance at 100k training steps. What is the cause? 3. In Appendix B.6, the performance of PEAC in Robosuite is fair comparing to that in DMC or Issacgym. Does PEAC work better for simple tasks (like walk, stand) than the tasks in Robosuite?

Rating

5

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

Yes

Area Chair YC1z2024-08-12

Reviewer Discussion Needed

Dear Reviewer, The discussion time is coming to an end soon. Please engage in the discussion process which is important to ensure a smooth and fruitful review process. Give notes on what parts of the reviewers responses that have and have not addressed your concerns.

Authorsrebuttal2024-08-14

Thank you again for your supportive comments at the end of the discussion

Dear Reviewer 3BHi: Thanks a lot for your supportive and constructive comments. As the rebuttal phase is nearing its end, we sincerely hope that you could take a moment to review our rebuttal. Your feedback is invaluable to us, and we truly appreciate your time and efforts in this review process. Thank you again for your time and effort in helping us improve our paper. Best, Authors.

Reviewer RG3H2024-08-08

Thank you very much for the thorough response. The additional experiments have completely addressed my concerns, and I will raise the rating by 2 points.

Authorsrebuttal2024-08-12

Look forward to further feedback

Dear Reviewer 3BHi: We sincerely thank you again for your valuable and constructive comments. We have tried our best to answer the concerns raised, especially including more ablation studies and more complex tasks, and are happy to clarify/discuss any further questions. We hope you may find our response satisfactory and raise your rating accordingly. We are looking forward to hearing from you about any further feedback. Best, Authors.

Authorsrebuttal2024-08-12

Look forward to further feedback

Dear Reviewer xr64: We sincerely thank you again for your constructive feedback and suggestions. We have tried our best to answer the concerns raised, especially including more ablation studies to explain the embodiment context encoder, and are happy to clarify/discuss any further questions. We hope you may find our response satisfactory and raise your rating accordingly. We are looking forward to hearing from you about any further feedback. Best, Authors.

Program Chairsdecision2024-09-25

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC