Merge of 20e5030

PiperOrigin-RevId: 606391547
Change-Id: I8d3794f58d67ff36fcd5a46670197e1c07d8a00d

CMakeLists.txt

@@ -60,6 +60,16 @@ set(MUJOCO_MPC_MUJOCO_GIT_TAG
     CACHE STRING "Git revision for MuJoCo."
 )
 
+set(MUJOCO_MPC_MENAGERIE_GIT_TAG
+    aef3ee5c07ea51506e893a62fd832773ff0162c8
+    CACHE STRING "Git revision for MuJoCo Menagerie."
+)
+
+set(MUJOCO_MPC_DM_CONTROL_GIT_TAG
+    774f46182140106e22725914aad3c6299ed91edd
+    CACHE STRING "Git revision for dm_control."
+)
+
 findorfetch(
   USE_SYSTEM_PACKAGE
   OFF
@@ -160,14 +170,25 @@ unset(BUILD_SHARED_LIBS_OLD)
 FetchContent_Declare(
   menagerie
   GIT_REPOSITORY https://github.com/google-deepmind/mujoco_menagerie.git
-  GIT_TAG main
+  GIT_TAG ${MUJOCO_MPC_MENAGERIE_GIT_TAG}
 )
 
 FetchContent_GetProperties(menagerie)
 if(NOT menagerie_POPULATED)
   FetchContent_Populate(menagerie)
 endif()
 
+FetchContent_Declare(
+  dm_control
+  GIT_REPOSITORY https://github.com/google-deepmind/dm_control.git
+  GIT_TAG ${MUJOCO_MPC_DM_CONTROL_GIT_TAG}
+)
+
+FetchContent_GetProperties(dm_control)
+if(NOT dm_control_POPULATED)
+  FetchContent_Populate(dm_control)
+endif()
+
 if(NOT TARGET lodepng)
   FetchContent_Declare(
     lodepng

docs/CONTRIBUTING.md

@@ -32,7 +32,13 @@ This code adheres to the [Google style](https://google.github.io/styleguide/).
 
 ## New Tasks
 
-When submitting a PR for a new task using models from [MuJoCo Menagerie](https://github.com/google-deepmind/mujoco_menagerie), do not include assets directly. Instead, modify the task [CMakeLists](mjpc/tasks/CMakeLists.txt) to copy these assets to the build binary.
+When submitting a PR for a new task that depends on third-party models, including from [MuJoCo Menagerie](https://github.com/google-deepmind/mujoco_menagerie) and [dm_control](https://github.com/google-deepmind/dm_control), do not include the xml model or assets in the task directly. Instead, modify the task [CMakeLists](mjpc/tasks/CMakeLists.txt) to copy the xml model and/or assets to the build binary.
+
+If the xml model needs to be modified, create a patch that is applied in the [CMakeLists](mjpc/tasks/CMakeLists.txt). A [patch](https://github.com/google-deepmind/mujoco_mpc/blob/main/mjpc/tasks/op3/op3.xml.patch) can be generated using the following command:
+```
+diff -u {original}.xml {modified}.xml > {modified}.xml.patch
+```
+The first three lines of the generated patch file will need to be be adapted for your use case. Please see an [example](https://github.com/google-deepmind/mujoco_mpc/blob/main/mjpc/tasks/op3/op3.xml.patch) for a template.
 
 ## Unit Tests
 

mjpc/tasks/CMakeLists.txt

@@ -16,7 +16,55 @@
 # built binary.
 
 add_custom_target(
-  copy_menagerie_resources ALL
+  copy_model_resources ALL
+  ## dm_control models
+  # acrobot
+  COMMAND ${CMAKE_COMMAND} -E copy
+          ${dm_control_SOURCE_DIR}/dm_control/suite/acrobot.xml
+          ${CMAKE_CURRENT_BINARY_DIR}/acrobot/acrobot.xml
+  COMMAND patch -o ${CMAKE_CURRENT_BINARY_DIR}/acrobot/acrobot_modified.xml
+          ${CMAKE_CURRENT_BINARY_DIR}/acrobot/acrobot.xml
+          <${CMAKE_CURRENT_SOURCE_DIR}/acrobot/acrobot.xml.patch
+  # cartpole
+  COMMAND ${CMAKE_COMMAND} -E copy
+          ${dm_control_SOURCE_DIR}/dm_control/suite/cartpole.xml
+          ${CMAKE_CURRENT_BINARY_DIR}/cartpole/cartpole.xml
+  COMMAND patch -o ${CMAKE_CURRENT_BINARY_DIR}/cartpole/cartpole_modified.xml
+          ${CMAKE_CURRENT_BINARY_DIR}/cartpole/cartpole.xml
+          <${CMAKE_CURRENT_SOURCE_DIR}/cartpole/cartpole.xml.patch
+  # humanoid
+  COMMAND ${CMAKE_COMMAND} -E copy
+          ${dm_control_SOURCE_DIR}/dm_control/suite/humanoid.xml
+          ${CMAKE_CURRENT_BINARY_DIR}/humanoid/humanoid.xml
+  COMMAND patch -o ${CMAKE_CURRENT_BINARY_DIR}/humanoid/humanoid_modified.xml
+          ${CMAKE_CURRENT_BINARY_DIR}/humanoid/humanoid.xml
+          <${CMAKE_CURRENT_SOURCE_DIR}/humanoid/humanoid.xml.patch
+  # particle
+  COMMAND ${CMAKE_COMMAND} -E copy
+          ${dm_control_SOURCE_DIR}/dm_control/suite/point_mass.xml
+          ${CMAKE_CURRENT_BINARY_DIR}/particle/particle.xml
+  COMMAND patch -o ${CMAKE_CURRENT_BINARY_DIR}/particle/particle_modified.xml
+          ${CMAKE_CURRENT_BINARY_DIR}/particle/particle.xml
+          <${CMAKE_CURRENT_SOURCE_DIR}/particle/particle.xml.patch
+  # swimmer
+  COMMAND ${CMAKE_COMMAND} -E copy
+          ${dm_control_SOURCE_DIR}/dm_control/suite/swimmer.xml
+          ${CMAKE_CURRENT_BINARY_DIR}/swimmer/swimmer.xml
+  COMMAND patch -o ${CMAKE_CURRENT_BINARY_DIR}/swimmer/swimmer_modified.xml
+          ${CMAKE_CURRENT_BINARY_DIR}/swimmer/swimmer.xml
+          <${CMAKE_CURRENT_SOURCE_DIR}/swimmer/swimmer.xml.patch
+  # walker
+  COMMAND ${CMAKE_COMMAND} -E copy
+          ${dm_control_SOURCE_DIR}/dm_control/suite/walker.xml
+          ${CMAKE_CURRENT_BINARY_DIR}/walker/walker.xml
+  COMMAND patch -o ${CMAKE_CURRENT_BINARY_DIR}/walker/walker_modified.xml
+          ${CMAKE_CURRENT_BINARY_DIR}/walker/walker.xml
+          <${CMAKE_CURRENT_SOURCE_DIR}/walker/walker.xml.patch
+
+  ## Menagerie models
+  COMMAND ${CMAKE_COMMAND} -E copy
+          ${menagerie_SOURCE_DIR}/shadow_hand/right_hand.xml
+          ${CMAKE_CURRENT_BINARY_DIR}/hand/right_hand.xml
   COMMAND ${CMAKE_COMMAND} -E copy
           ${menagerie_SOURCE_DIR}/wonik_allegro/right_hand.xml
           ${CMAKE_CURRENT_BINARY_DIR}/allegro/right_hand.xml
@@ -29,21 +77,39 @@ add_custom_target(
   COMMAND ${CMAKE_COMMAND} -E copy_directory
           ${menagerie_SOURCE_DIR}/shadow_hand/assets
           ${CMAKE_CURRENT_BINARY_DIR}/hand/assets
+  COMMAND ${CMAKE_COMMAND} -E copy
+          ${menagerie_SOURCE_DIR}/franka_emika_panda/panda.xml
+          ${CMAKE_CURRENT_BINARY_DIR}/panda/panda.xml
   COMMAND ${CMAKE_COMMAND} -E copy_directory
           ${menagerie_SOURCE_DIR}/franka_emika_panda/assets
           ${CMAKE_CURRENT_BINARY_DIR}/panda/assets
+  COMMAND patch -o ${CMAKE_CURRENT_BINARY_DIR}/panda/panda_modified.xml
+          ${CMAKE_CURRENT_BINARY_DIR}/panda/panda.xml
+          <${CMAKE_CURRENT_SOURCE_DIR}/panda/panda.xml.patch
+  COMMAND ${CMAKE_COMMAND} -E copy
+          ${menagerie_SOURCE_DIR}/unitree_a1/a1.xml
+          ${CMAKE_CURRENT_BINARY_DIR}/quadruped/a1.xml
   COMMAND ${CMAKE_COMMAND} -E copy_directory
           ${menagerie_SOURCE_DIR}/unitree_a1/assets
           ${CMAKE_CURRENT_BINARY_DIR}/quadruped/assets
+  COMMAND patch -o ${CMAKE_CURRENT_BINARY_DIR}/quadruped/a1_modified.xml
+          ${CMAKE_CURRENT_BINARY_DIR}/quadruped/a1.xml
+          <${CMAKE_CURRENT_SOURCE_DIR}/quadruped/a1.xml.patch
   COMMAND ${CMAKE_COMMAND} -E copy_directory
           ${menagerie_SOURCE_DIR}/franka_emika_panda
           ${CMAKE_CURRENT_BINARY_DIR}/manipulation
   COMMAND ${CMAKE_COMMAND} -E copy_directory
           ${menagerie_SOURCE_DIR}/robotiq_2f85
           ${CMAKE_CURRENT_BINARY_DIR}/manipulation
+  COMMAND ${CMAKE_COMMAND} -E copy
+          ${menagerie_SOURCE_DIR}/skydio_x2/x2.xml
+          ${CMAKE_CURRENT_BINARY_DIR}/quadrotor/quadrotor.xml
   COMMAND ${CMAKE_COMMAND} -E copy_directory
           ${menagerie_SOURCE_DIR}/skydio_x2/assets
           ${CMAKE_CURRENT_BINARY_DIR}/quadrotor/assets
+  COMMAND patch -o ${CMAKE_CURRENT_BINARY_DIR}/quadrotor/quadrotor_modified.xml
+          ${CMAKE_CURRENT_BINARY_DIR}/quadrotor/quadrotor.xml
+          <${CMAKE_CURRENT_SOURCE_DIR}/quadrotor/quadrotor.xml.patch
   ## Cube solve task
   # copy cube model from MuJoCo
   COMMAND ${CMAKE_COMMAND} -E copy
@@ -97,14 +163,14 @@ add_custom_target(
   COMMAND ${Python_EXECUTABLE}
           ${CMAKE_CURRENT_BINARY_DIR}/manipulation/merge_panda_robotiq.py
           ${CMAKE_CURRENT_BINARY_DIR}/manipulation/panda_robotiq.xml
-  COMMENT "Copying menagerie assets into binary directory")
+  COMMENT "Copying Menagerie and dm_control assets into binary directory")
 
 add_custom_target(copy_resources ALL
         COMMAND ${CMAKE_COMMAND} -E copy_directory
         ${CMAKE_CURRENT_SOURCE_DIR}
         ${CMAKE_CURRENT_BINARY_DIR}
         COMMENT "Copying tasks into binary directory")
 
-add_dependencies(copy_menagerie_resources copy_resources)
+add_dependencies(copy_model_resources copy_resources)
 
-add_dependencies(libmjpc copy_menagerie_resources)
+add_dependencies(libmjpc copy_model_resources)

mjpc/tasks/acrobot/acrobot.xml.patch

@@ -0,0 +1,28 @@
+diff --git a/acrobot_modified.xml b/acrobot_modified.xml
+--- a/acrobot_modified.xml
++++ b/acrobot_modified.xml
+@@ -6,22 +6,18 @@
+     IEEE control systems 15, no. 1 (1995): 49-55.
+ -->
+ <mujoco model="acrobot">
+-  <include file="./common/visual.xml"/>
+-  <include file="./common/skybox.xml"/>
+-  <include file="./common/materials.xml"/>
+-
+   <default>
+     <joint damping=".05"/>
+     <geom type="capsule" mass="1"/>
+   </default>
+
+-  <option timestep="0.01" integrator="RK4">
++  <option timestep="0.01">
+     <flag constraint="disable" energy="enable"/>
+   </option>
+
+   <worldbody>
+     <light name="light" pos="0 0 6"/>
+-    <geom name="floor" size="3 3 .2" type="plane" material="grid"/>
++    <geom name="floor" size="3 3 .2" type="plane" material="blue_grid"/>
+     <site name="target" type="sphere" pos="0 0 4" size="0.2" material="target" group="3"/>
+     <camera name="fixed" pos="0 -6 2" zaxis="0 -1 0"/>
+     <camera name="lookat" mode="targetbodycom" target="upper_arm" pos="0 -2 3"/>

mjpc/tasks/acrobot/task.xml

@@ -1,6 +1,7 @@
 <mujoco model="Acrobot Swing-Up">
   <include file="../common.xml"/>
-  <include file="acrobot.xml" />
+  <!-- modified from: https://github.com/google-deepmind/dm_control/blob/main/dm_control/suite/acrobot.xml-->
+  <include file="acrobot_modified.xml" />
 
   <size memory="4K"/>
 

mjpc/tasks/cartpole/cartpole.xml.patch

@@ -0,0 +1,37 @@
+diff --git a/cartpole_modified.xml b/cartpole_modified.xml
+--- a/cartpole_modified.xml
++++ b/cartpole_modified.xml
+@@ -1,10 +1,6 @@
+-<mujoco model="cart-pole">
+-  <include file="./common/skybox.xml"/>
+-  <include file="./common/visual.xml"/>
+-  <include file="./common/materials.xml"/>
+-
+-  <option timestep="0.01" integrator="RK4">
+-    <flag contact="disable" energy="enable"/>
++<mujoco model="Cartpole">
++  <option timestep="0.001">
++    <flag contact="disable"/>
+   </option>
+
+   <default>
+@@ -18,15 +14,16 @@
+     <light name="light" pos="0 0 6"/>
+     <camera name="fixed" pos="0 -4 1" zaxis="0 -1 0"/>
+     <camera name="lookatcart" mode="targetbody" target="cart" pos="0 -2 2"/>
+-    <geom name="floor" pos="0 0 -.05" size="4 4 .2" type="plane" material="grid"/>
++    <geom name="floor" pos="0 0 -.05" size="4 4 .2" type="plane" material="blue_grid"/>
+     <geom name="rail1" type="capsule" pos="0  .07 1" zaxis="1 0 0" size="0.02 2" material="decoration" />
+     <geom name="rail2" type="capsule" pos="0 -.07 1" zaxis="1 0 0" size="0.02 2" material="decoration" />
+     <body name="cart" pos="0 0 1">
+-      <joint name="slider" type="slide" limited="true" axis="1 0 0" range="-1.8 1.8" solreflimit=".08 1" damping="5e-4"/>
++      <joint name="slider" type="slide" limited="true" axis="1 0 0" range="-1.8 1.8" solreflimit=".08 1" damping="1.0e-4"/>
+       <geom name="cart" type="box" size="0.2 0.15 0.1" material="self"  mass="1"/>
+       <body name="pole_1" childclass="pole">
+-        <joint name="hinge_1"/>
++        <joint name="hinge_1" damping="1.0e-4"/>
+         <geom name="pole_1"/>
++        <site name="tip" pos="0 0 1"/>
+       </body>
+     </body>
+   </worldbody>

mjpc/tasks/cartpole/task.xml

@@ -1,6 +1,7 @@
 <mujoco model="Cart-Pole Swing-Up">
   <include file="../common.xml"/>
-  <include file="cartpole.xml" />
+  <!-- modified from: https://github.com/google-deepmind/dm_control/blob/main/dm_control/suite/cartpole.xml -->
+  <include file="cartpole_modified.xml" />
 
   <size memory="4K"/>