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

Thanks a lot for your review, especially the praise for PEAC's contributions to embodied intelligence. Below we will address all your concerns.

W1: About baselines utilizing embodiment information $e$.

A: Thank you for recognizing that we have included various standard and SOTA baselines. Following your valuable suggestions of adding baselines incorporating embodiment information, we have supplemented experiments of LBS-Context and DIAYN-Context, which inherit LBS and DIAYN by incorporating embodiment information $e$, to better highlight the paper's contribution. Our results are below:

As shown in these two tables, LBS-Context and DIAYN-Context own superior performance compared with LBS and DIAYN respectively, and PEAC still significantly outperforms these baselines. Especially, in state-based DMC, the improvement in the performance of LBS-Context and DIAYN-Context is more remarkable than in image-based DMC. It may be because, in image-based DMC, we take DreamerV2 as the RL backbone, which has encoded historical information for making the decision and is possible to distinguish different embodiments through historical information to some degree. But in state-based DMC, we take DDPG as the RL backbones following previous works [1], which are Markovian policies and thus utilizing embodiment information during the pre-training is helpful (Lines 210-216). Consequently, our supplemented results highlight that both the embodiment discriminator and cross-embodiment intrinsic rewards $\mathcal{R}_{\text{CE}}$ are effective for handling CEURL (more details, including figures with aggregate metrics of these methods, are in the global response and the corresponding PDF).

W2: About experimental settings that can better highlight cross-embodiment.

A: Thank you for your recognition of our experiments on Robosuite. In DMC and Isaacgym, we mainly modify the property parameters or introduce joint torque failures because embodiments with different properties are one of the most basic cross-embodiment settings, which is worth paying attention to especially when we first research CEURL. Moreover, although these embodiments only modify property parameters, their ability to handle tasks may differ a lot, e.g. in our real-world experiments, when the front foot joint of the A1 robot fails, it can only drag its legs forward, while when the rear foot joint fails, it will lift its rear legs to move forward (Fig. 7 and video in the supplementary material).

In experiments of DMC, besides settings with different properties like Walker-mass, we have also included settings with different morphology in Appendix B.10: Walker-Cheetah, where Walker and Cheetah both own two legs but their morphology differs a lot. Experiments in Appendix B.10 demonstrate that PEAC can also significantly outperform baselines and reveal PEAC's potential in handling exactly different embodiments.

Moreover, following your constructive suggestions, we have also supplemented extensive experiments in DMC considering settings with greater differences across embodiments, especially differing in morphology:

• Walker-Humanoid: Following your constructive suggestions, we consider the cross-embodiment setting that includes Walker robots and Humanoid robots. Their robot properties, robot shapes, and action spaces are all different.
• Walker-length and Cheetah-torsolength [2]: The former includes walker robots with different foot lengths while the second one includes cheetah robots with different torso lengths. Thus robots' properties and morphologies are different (Figs of these embodiments are in the attached PDF of the global response).

The results in these settings are below:

As shown in these tables, PEAC can achieve much greater performance compared to baselines. These experiments indicate that PEAC has powerful abilities to handle various kinds of embodiment differences, including different morphologies. Also, designing more effective methods for handling complicated embodiments like Humanoid is a promising future direction.

Reference:

[1] URLB: Unsupervised Reinforcement Learning Benchmark

[2] Learning Robust State Abstractions for Hidden-Parameter Block MDPs

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

Thanks a lot for the supportive review and constructive suggestions. Below we address the detailed comments.

W1: About the embodiment context encoder in line 68.

A: The embodiment context encoder is the embodiment discriminator in the main text of the paper, and we will polish our paper to make it consistent in context. As for the training of the discriminator, we utilize cross-entropy loss of the one-hot embodiment vector during the pre-training stage, following previous works [1]. We will clarify it in the revised versions.

To better clarify the impact of the embodiment discriminator and the cross-embodiment intrinsic reward $R_{\text{CE}}$, we have supplemented ablation studies of LBS-Context and DIAYN-Context, i.e., combining LBS and DIAYN with our embodiment discriminator in PEAC but without $R_{\text{CE}}$. Our results are as below:

As shown in these two tables, LBS-Context and DIAYN-Context own superior performance compared with LBS and DIAYN respectively, and PEAC still significantly outperforms these baselines. Consequently, this ablation study highlights that both the embodiment discriminator and cross-embodiment intrinsic rewards $R_{\text{CE}}$ are effective for handling CEURL (more details and results are in the global response and the attached PDF).

W2: About the information geometry analyses in Sec. 3.

A: Yes, our extension of the information geometry analyses does not directly guide the design of PEAC. As the major contributions of this paper include the novel setting CEURL and our algorithm PEAC, our information geometry analyses are mainly related to CEURL by revealing its difficulty. Moreover, as you have mentioned, our information geometry analyses highlight that policies may be stochastic or history-related, thus our PEAC chooses to encode historical state-action pairs to the embodiment context in practice (Lines 213-216). We will emphasize it in the revised version.

W3: About embodiment-aware or embodiment-agnostic skills?

A: The goal of embodiment-aware skill discovery (Lines 232-252) is to learn similar skills for all these embodiments during the pre-training stage. However, as different embodiments own different properties/morphologies, their actions for achieving similar skills may differ a lot. Thus we name our method embodiment-aware skill discovery. Thanks again for your question, we will discuss the concept more clearly in the revised version.

Q1: How to handle different action/state spaces?

A: When handling state/action spaces with different dimensions, we will take the maximal dimension of these spaces $d$ and pad the states/actions into $d$ dimension by filling in zeros, following previous works [2].

Q2: What does $c$ in the MDP $\mathcal{M}^c$ represent?

A: We use $\mathcal{M}$ to represent original MDP, $\mathcal{M}_e$ represent MDP with embodiment $e$, $\mathcal{M}^c$ represent controlled MDP, i.e., an MDP without rewards in the pre-training stage, and $\mathcal{M}_e^c$ represent controlled MDP with embodiment $e$.

Q3: About the equality in the last line of Eq. 29.

A: The last equality of Eq. 29 can be derived by the property of conditional probability and we abbreviate $\mathcal{M}_e^c$ as $e$:

$\log \frac{p_{\pi}(\tau)}{p_{\pi}(\tau|e)} = \log \frac{p_{\pi}(\tau)}{p_{\pi}(\tau, e) / p_{\pi}(e)} = \log \frac{p_{\pi}(e) }{p_{\pi}(\tau, e) / p_{\pi}(\tau)} = \log \frac{p(e)}{p_{\pi}(e|\tau)}.$

Thanks for your question, we will provide the detailed derivation in the revised versions.

Q4: About the connection between CEURL and contextual MDP.

A: As we have mentioned in Lines 125-126, cross-embodiment RL can be formulated by contextual MDP. Moreover, as the cross-embodiment setting is practical for real-world applications and owns several unique properties, it introduces various new problems like cross-embodiment transfer [3] and has received widespread attention in embodied intelligence [3-5] (as mentioned by Reviewer RG3H). In this work, considering that different embodiments may own similar structures and learn similar skills, we propose unsupervised pre-train across embodiments to learn knowledge only specialized in these embodiments themselves, i.e., CEURL, which is a novel setting that existing contextual MDP can not cover and thus we combine contextual MDP with controlled MDP as our CEMDP to formulate our problem.

Q5: Does $\log p(e)$ used in the reward?

A: No, as we have mentioned in Lines 205-206, our cross-embodiment intrinsic reward eliminates $\log p(e)$ as it is fixed.

Limitation: About experimental settings that can better highlight cross-embodiment

A: We have supplemented more cross-embodiments with different morphologies in DMC: Walker-Humanoid, Walker-length, and Cheetah-torsolength. More details and results are in the global response, where PEAC shows greater performance compared to baselines. The results indicate that PEAC has powerful abilities to handle various kinds of embodiment differences, including different morphologies.

Reference:

[1] Multi-task reinforcement learning with soft modularization

[2] Learning to Modulate pre-trained Models in RL

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

[4] Xirl: Cross-embodiment inverse reinforcement learning

[5] Pushing the limits of cross-embodiment learning for manipulation and navigation

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

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.

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.

Thank you for your supportive review and valuable suggestions. Below we address the detailed comments for all your questions.

W1: About more complex tasks.

A: In our experiments, tasks chosen in DMC and Isaacgym are standard and widely used in unsupervised RL evaluation [1, 2, 3]. These tasks are basic and important to evaluate agents' locomotion ability, and they can also combine to handle much more complex tasks like parkour [2]. Thanks for your suggestion and reference, referring to [4] including tasks of locomotion in complicated terrain like inclines, we have designed more complicated tasks for Walker-mass in DMC to evaluate their locomotion ability in complicated terrain like inclines: Walker-mass-incline. The results are shown below:

As shown in this table, the performance of LBS and PEAC-LBS decreases when locomoting in the incline terrain due to its complexity. PEAC-LBS still significantly outperforms LBS, expressing that our method, especially the cross-embodiment intrinsic rewards, benefits cross-embodiment unsupervised pre-training for handling more complicated tasks.

W2 & Q1: More ablation studies, like only using one of the two parts in Eq. 2.

A: Thanks for your suggestion. There are two terms in Eq.2: the policy improvement part and the policy constraint part, representing that in the fine-tuning stage, we need to maximize the policy performance within limited training steps, respectively (Lines 182-188 in the paper). Thus they need to be considered together and it is unreasonable to use only one of these two parts for this setting. Moreover, we introduce a trade-off parameter $\beta$ in Eq.2 to balance these two terms, which is set as 1.0 in all our experiments. Consequently, we have supplemented ablation of $\beta$, of which the results are:

As shown in this table, when $\beta$ is smaller, the performance of PEAC-LBS decreases a little, and overall PEAC-LBS is stable with different $\beta$.

Besides $\beta$, we have also supplemented more ablation studies, especially about our cross-embodiment intrinsic rewards $R_{\text{CE}}$, we consider LBS-Context and DIAYN-Context i.e., combining LBS and DIAYN with our embodiment discriminator in PEAC but without $R_{\text{CE}}$. Our results are as below:

As shown in these two tables, LBS-Context and DIAYN-Context own superior performance compared with LBS and DIAYN respectively, and PEAC still significantly outperforms these baselines. Consequently, this ablation study highlights that both the embodiment discriminator and cross-embodiment intrinsic rewards $R_{\text{CE}}$ are necessary for handling CEURL (more details of results are in the global response).

W3: About the relation with the BeCL [2].

A: BeCL [2] mainly focuses on encouraging the agent to learn diverse skills related to corresponding behaviors via applying contrastive learning for single-embodiment skill discovery. As you have mentioned, our PEAC considers the cross-embodiment setting and thus is orthogonal to these skill-discovery methods.

We have discussed and compared BeCL as well as other skill discovery methods in the paper. Especially, as shown in Sec.4, we have discussed the relationship between the objective of skill discovery (including BeCL) and our cross-embodiment objective, with a result of a unified cross-embodiment skill-based adaptation objective (Eq. 6-7). Thus our PEAC can naturally combine with these skill-discovery methods, including BeCL, and we propose PEAC-DIAYN as an example. Thanks again for your question, applying contrastive learning methods like BeCL to our PEAC for handling CEURL is indeed an interesting future direction.

Q2: About the performance of PEAC-DIAYN in 100k pre-training steps.

A: Thanks for your question. When the pre-training steps are small (like 100k), it is difficult to learn useful knowledge during the pre-training stage (like world models or our embodiment discriminator), and the uncertainty of results is relatively high, thus current research mainly considers the fine-tuning results after pre-training 2M in DMC [1-2].

Q3: About the performance of PEAC in Robosuite.

A: In Robosuite, the overall performance of PEAC is still 10% higher than all baselines, indicating the effectiveness of PEAC. Compared with DMC, the morphologies in Robosuite vary a lot more, thus baselines may also distinguish different embodiments directly via their morphology observation. Consequently, PEAC's leading advantage may slightly decrease, but still significantly outperforms baselines (We have also supplemented experiments in DMC with varying morphologies in the global rebuttal). Thanks for your question, we will discuss it more in the revised versions.

Reference:

[1] URLB: Unsupervised Reinforcement Learning Benchmark

[2] Behavior contrastive learning for unsupervised skill discovery

[3] Robot Parkour Learning

[4] Embodied intelligence via learning and evolution

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.

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.