Merge deepmind branch into main

Deepmind merge

CMakeLists.txt

@@ -56,12 +56,12 @@ set(MUJOCO_BUILD_TESTS OFF)
 set(MUJOCO_TEST_PYTHON_UTIL OFF)
 
 set(MUJOCO_MPC_MUJOCO_GIT_TAG
-    893c4042302e5a5738dbc462466e746f2bbc1ea5
+    3.1.1
     CACHE STRING "Git revision for MuJoCo."
 )
 
 set(MUJOCO_MPC_MENAGERIE_GIT_TAG
-    94ea114fa8c60a0fd542c8e1ffeb997204acbea2
+    8ef01e87fffaa8ec634a4826c5b2092733b2f3c8
     CACHE STRING "Git revision for MuJoCo Menagerie."
 )
 

mjpc/CMakeLists.txt

@@ -44,12 +44,8 @@ add_library(
   tasks/bimanual/bimanual.h
   tasks/cartpole/cartpole.cc
   tasks/cartpole/cartpole.h
-  tasks/cube/solve.cc
-  tasks/cube/solve.h
   tasks/fingers/fingers.cc
   tasks/fingers/fingers.h
-  tasks/hand/hand.cc
-  tasks/hand/hand.h
   tasks/humanoid/stand/stand.cc
   tasks/humanoid/stand/stand.h
   tasks/humanoid/tracking/tracking.cc
@@ -66,6 +62,10 @@ add_library(
   tasks/panda/panda.h
   tasks/particle/particle.cc
   tasks/particle/particle.h
+  tasks/rubik/solve.cc
+  tasks/rubik/solve.h
+  tasks/shadow_reorient/hand.cc
+  tasks/shadow_reorient/hand.h
   tasks/quadrotor/quadrotor.cc
   tasks/quadrotor/quadrotor.h
   tasks/quadruped/quadruped.cc

mjpc/app.cc

@@ -40,7 +40,8 @@
 #include "mjpc/threadpool.h"
 #include "mjpc/utilities.h"
 
-ABSL_FLAG(std::string, task, "", "Which model to load on startup.");
+ABSL_FLAG(std::string, task, "Quadruped Flat",
+          "Which model to load on startup.");
 ABSL_FLAG(bool, planner_enabled, false,
           "If true, the planner will run on startup");
 ABSL_FLAG(float, sim_percent_realtime, 100,
@@ -392,7 +393,7 @@ void PhysicsLoop(mj::Simulate& sim) {
 
 namespace mjpc {
 
-MjpcApp::MjpcApp(std::vector<std::shared_ptr<mjpc::Task>> tasks, int task_id) {
+MjpcApp::MjpcApp(std::vector<std::shared_ptr<mjpc::Task>> tasks) {
   // MJPC
   printf("MuJoCo MPC (MJPC)\n");
 
@@ -415,8 +416,8 @@ MjpcApp::MjpcApp(std::vector<std::shared_ptr<mjpc::Task>> tasks, int task_id) {
 
   sim->agent->SetTaskList(std::move(tasks));
   std::string task_name = absl::GetFlag(FLAGS_task);
-  if (task_name.empty()) {
-    sim->agent->gui_task_id = task_id;
+  if (task_name.empty()) {  // shouldn't happen, flag has a default value
+    sim->agent->gui_task_id = 0;
   } else {
     sim->agent->gui_task_id = sim->agent->GetTaskIdByName(task_name);
     if (sim->agent->gui_task_id == -1) {
@@ -519,8 +520,8 @@ mj::Simulate* MjpcApp::Sim() {
   return sim.get();
 }
 
-void StartApp(std::vector<std::shared_ptr<mjpc::Task>> tasks, int task_id) {
-  MjpcApp app(std::move(tasks), task_id);
+void StartApp(std::vector<std::shared_ptr<mjpc::Task>> tasks) {
+  MjpcApp app(std::move(tasks));
   app.Start();
 }
 

mjpc/app.h

@@ -24,7 +24,7 @@
 namespace mjpc {
 class MjpcApp {
  public:
-  MjpcApp(std::vector<std::shared_ptr<mjpc::Task>> tasks, int task_id = 0);
+  explicit MjpcApp(std::vector<std::shared_ptr<mjpc::Task>> tasks);
   MjpcApp(const MjpcApp&) = delete;
   MjpcApp& operator=(const MjpcApp&) = delete;
   ~MjpcApp();
@@ -34,7 +34,7 @@ class MjpcApp {
   ::mujoco::Simulate* Sim();
 };
 
-void StartApp(std::vector<std::shared_ptr<mjpc::Task>> tasks, int task_id = 0);
+void StartApp(std::vector<std::shared_ptr<mjpc::Task>> tasks);
 
 }  // namespace mjpc
 

mjpc/main.cc

@@ -47,6 +47,6 @@ int main(int argc, char** argv) {
 #endif
   absl::ParseCommandLine(argc, argv);
 
-  mjpc::StartApp(mjpc::GetTasks(), 11);  // start with quadruped flat
+  mjpc::StartApp(mjpc::GetTasks());
   return 0;
 }

mjpc/planners/cross_entropy/planner.cc

@@ -322,7 +322,7 @@ void CrossEntropyPlanner::ResamplePolicy(int horizon) {
 
   LinearRange(resampled_policy.times.data(), time_shift,
               resampled_policy.times[0], num_spline_points);
-  
+
   resampled_policy.representation = policy.representation;
 }
 

mjpc/planners/sampling/planner.cc

@@ -267,8 +267,13 @@ void SamplingPlanner::AddNoiseToPolicy(int i) {
   int shift = i * (model->nu * kMaxTrajectoryHorizon);
 
   // sample noise
-  for (int k = 0; k < num_parameters; k++) {
-    noise[k + shift] = absl::Gaussian<double>(gen_, 0.0, noise_exploration);
+  for (int t = 0; t < num_spline_points; t++) {
+    for (int k = 0; k < model->nu; k++) {
+      double scale = 0.5 * (model->actuator_ctrlrange[2 * k + 1] -
+                            model->actuator_ctrlrange[2 * k]);
+      noise[shift + t * model->nu + k] =
+          absl::Gaussian<double>(gen_, 0.0, scale * noise_exploration);
+    }
   }
 
   // add noise

mjpc/simulate.cc

@@ -1163,10 +1163,7 @@ void UiEvent(mjuiState* state) {
         sim->screenshotrequest.store(true);
         break;
       }
-    }
-
-    // option section
-    else if (it && it->sectionid==SECT_OPTION) {
+    } else if (it && it->sectionid == SECT_OPTION) {
       if (it->pdata == &sim->spacing) {
         sim->ui0.spacing = mjui_themeSpacing(sim->spacing);
         sim->ui1.spacing = mjui_themeSpacing(sim->spacing);
@@ -1184,37 +1181,16 @@ void UiEvent(mjuiState* state) {
       // modify UI
       UiModify(&sim->ui0, state, &sim->platform_ui->mjr_context());
       UiModify(&sim->ui1, state, &sim->platform_ui->mjr_context());
-    }
 
-    // simulation section
-    else if (it && it->sectionid==SECT_SIMULATION) {
+    } else if (it && it->sectionid == SECT_SIMULATION) {
       switch (it->itemid) {
       case 1:             // Reset
         if (m) {
-          mj_resetData(m, d);
+          mj_resetDataKeyframe(m, d, mj_name2id(m, mjOBJ_KEY, "home"));
           mj_forward(m, d);
           UpdateProfiler(sim);
           UpdateSensor(sim);
           UpdateSettings(sim);
-          // set initial qpos via keyframe
-          double* key_qpos = mjpc::KeyQPosByName(sim->mnew, sim->dnew,
-                                                 "home");
-          if (key_qpos) {
-            mju_copy(sim->dnew->qpos, key_qpos, sim->mnew->nq);
-          }
-          // set initial qvel via keyframe
-          double* key_qvel = mjpc::KeyQVelByName(sim->mnew, sim->dnew,
-                                                 "home");
-          if (key_qvel) {
-            mju_copy(sim->dnew->qvel, key_qvel, sim->mnew->nv);
-          }
-          // set initial act via keyframe
-          double* act = mjpc::KeyActByName(sim->mnew, sim->dnew,
-                                           "home");
-          if (act) {
-            mju_copy(sim->dnew->act, act, sim->mnew->na);
-          }
-
           sim->agent->PlotReset();
         }
         break;

mjpc/tasks/CMakeLists.txt

@@ -64,19 +64,21 @@ add_custom_target(
   ## Menagerie models
   COMMAND ${CMAKE_COMMAND} -E copy
           ${menagerie_SOURCE_DIR}/shadow_hand/right_hand.xml
-          ${CMAKE_CURRENT_BINARY_DIR}/hand/right_hand.xml
+          ${CMAKE_CURRENT_BINARY_DIR}/shadow_reorient/right_hand.xml
   COMMAND ${CMAKE_COMMAND} -E copy
           ${menagerie_SOURCE_DIR}/wonik_allegro/right_hand.xml
           ${CMAKE_CURRENT_BINARY_DIR}/allegro/right_hand.xml
+  COMMAND ${CMAKE_COMMAND} -E copy_directory
+          ${menagerie_SOURCE_DIR}/shadow_hand/assets
+          ${CMAKE_CURRENT_BINARY_DIR}/shadow_reorient/assets
+
   COMMAND ${CMAKE_COMMAND} -E copy_directory
           ${menagerie_SOURCE_DIR}/wonik_allegro/assets
           ${CMAKE_CURRENT_BINARY_DIR}/allegro/assets
   COMMAND patch -o ${CMAKE_CURRENT_BINARY_DIR}/allegro/right_hand_modified.xml
           ${CMAKE_CURRENT_BINARY_DIR}/allegro/right_hand.xml
           <${CMAKE_CURRENT_SOURCE_DIR}/allegro/right_hand.xml.patch
-  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
@@ -86,6 +88,7 @@ add_custom_target(
   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
@@ -95,12 +98,14 @@ add_custom_target(
   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
@@ -110,31 +115,32 @@ add_custom_target(
   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
           ${mujoco_SOURCE_DIR}/model/cube/cube_3x3x3.xml
-          ${CMAKE_CURRENT_BINARY_DIR}/cube/cube_3x3x3.xml
+          ${CMAKE_CURRENT_BINARY_DIR}/rubik/cube_3x3x3.xml
   # copy cube assets from MuJoCo
   COMMAND ${CMAKE_COMMAND} -E copy_directory
           ${mujoco_SOURCE_DIR}/model/cube/assets
-          ${CMAKE_CURRENT_BINARY_DIR}/cube/assets
+          ${CMAKE_CURRENT_BINARY_DIR}/rubik/assets
   # modified cube model for task
-  COMMAND patch -o ${CMAKE_CURRENT_BINARY_DIR}/cube/cube_3x3x3_modified.xml
-          ${CMAKE_CURRENT_BINARY_DIR}/cube/cube_3x3x3.xml
-          <${CMAKE_CURRENT_SOURCE_DIR}/cube/cube_3x3x3.xml.patch
+  COMMAND patch -o ${CMAKE_CURRENT_BINARY_DIR}/rubik/cube_3x3x3_modified.xml
+          ${CMAKE_CURRENT_BINARY_DIR}/rubik/cube_3x3x3.xml
+          <${CMAKE_CURRENT_SOURCE_DIR}/rubik/cube_3x3x3.xml.patch
   # modified cube model to transition model for scramble mode
-  COMMAND patch -o ${CMAKE_CURRENT_BINARY_DIR}/cube/transition_model.xml
-          ${CMAKE_CURRENT_BINARY_DIR}/cube/cube_3x3x3.xml
-          <${CMAKE_CURRENT_SOURCE_DIR}/cube/transition_model.xml.patch
+  COMMAND patch -o ${CMAKE_CURRENT_BINARY_DIR}/rubik/transition_model.xml
+          ${CMAKE_CURRENT_BINARY_DIR}/rubik/cube_3x3x3.xml
+          <${CMAKE_CURRENT_SOURCE_DIR}/rubik/transition_model.xml.patch
   # copy hand model from Menagerie
   COMMAND ${CMAKE_COMMAND} -E copy
           ${menagerie_SOURCE_DIR}/shadow_hand/right_hand.xml
-          ${CMAKE_CURRENT_BINARY_DIR}/cube/right_hand.xml
+          ${CMAKE_CURRENT_BINARY_DIR}/rubik/right_hand.xml
   # copy hand assets from Menagerie
   COMMAND ${CMAKE_COMMAND} -E copy_directory
           ${menagerie_SOURCE_DIR}/shadow_hand/assets
-          ${CMAKE_CURRENT_BINARY_DIR}/cube/assets
+          ${CMAKE_CURRENT_BINARY_DIR}/rubik/assets
 
   # ALOHA
   COMMAND ${CMAKE_COMMAND} -E copy_directory

mjpc/tasks/allegro/task.xml

@@ -58,6 +58,7 @@
     <!-- Measurements we want to use -->
     <framepos name="cube_goal_position" objtype="site" objname="grasp_site"/>
     <framequat name="cube_goal_orientation" objtype="body" objname="goal"/>
+    <framepos name="trace0" objtype="body" objname="cube"/>
   </sensor>
 
   <include file="../common_assets/reorientation_cube.xml"/>

mjpc/tasks/bimanual/bimanual.cc

@@ -48,9 +48,6 @@ void Bimanual::ResidualFn::Residual(const mjModel* model, const mjData* data,
 }
 
 void mjpc::Bimanual::TransitionLocked(mjModel* model, mjData* data) {
-  data->mocap_pos[0] = -0.4;
-  data->mocap_pos[1] = -0.2;
-  data->mocap_pos[2] = 0.3;
 }
 
 }  // namespace mjpc

mjpc/tasks/bimanual/task.xml

@@ -15,7 +15,7 @@
     <numeric name="agent_sample_width" data="0.0025" />
     <numeric name="agent_policy_width" data="0.0035" />
     <numeric name="sampling_exploration" data="0.5" />
-    <numeric name="sampling_trajectories" data="100"/>
+    <numeric name="sampling_trajectories" data="120"/>
     <numeric name="sampling_spline_points" data="4" />
     <numeric name="gradient_spline_points" data="6" />
   </custom>
@@ -36,30 +36,31 @@
     <geom name="floor" pos="0 0 -0.75" size="0 0 0.05" type="plane" material="groundplane"/>
     <body name="table" pos="0 0 -0.75">
       <geom name="table" pos="0 0 0.6509" size="0.61 0.37 0.1" type="box" class="collision"/>
+      <geom name="table_visual" pos="0 0 0.6509" size="0.61 0.37 0.1" type="box" rgba="0.4 0.4 0.4 1" conaffinity="0" contype="0"/>
     </body>
     <body mocap="true" name="box_goal_mocap">
-      <geom group="5" size="0.02" rgba="1 0.4 0.4 1" conaffinity="0" contype="0"/>
+      <geom group="2" size="0.02" rgba="1 0.4 0.4 1" conaffinity="0" contype="0"/>
     </body>
     <body name="box">
       <freejoint/>
-      <geom name="box" type="box" size="0.025 0.025 0.025" priority="1" condim="6"
+      <geom name="box" type="box" size="0.015 0.015 0.015" priority="1" condim="6"
         friction="1.5 .03 .003" rgba="0 1 0 1"/>
     </body>
   </worldbody>
   <keyframe>
     <key name="home" qpos=
     "0 -0.96 1.16 0 -0.3 0 0.002 0.002
     0 -0.96 1.16 0 -0.3 0 0.002 0.002
-    0.1 0.2 0.3 1 0 0 0"
-    ctrl="-0.1 0 0 0 0 0 0.01 0.1 0 0 0 0 0 0.01"/>
+    -0.25 0.2 0.3 1 0 0 0"
+    mpos="0.4 -0.25 0.3"
+    act= "-0.1 0 0 0 0 0 0.03 0.1 0 0 0 0 0 0.03"
+    ctrl="-0.1 0 0 0 0 0 0.03 0.1 0 0 0 0 0 0.03"/>
   </keyframe>
 
   <sensor>
-    <user name="Reach L" dim="3" user="2 0.1 0 5 0.01"/>
-    <user name="Reach R" dim="3" user="2 0.1 0 5 0.01"/>
+    <user name="Reach L" dim="3" user="2 0.1 0 .5 0.005"/>
+    <user name="Reach R" dim="3" user="2 0.1 0 .5 0.005"/>
     <user name="Bring" dim="3" user="2 1 0 1 0.003"/>
-    <!-- <user name="arm_table_contact" dim="2" user="2 0.5 0 1 0.01"/>
-    <user name="arm_arm_contact" dim="1" user="2 0.5 0 1 0.01"/> -->
     <framepos name="left/gripper" objtype="site" objname="left/gripper"/>
     <framepos name="right/gripper" objtype="site" objname="right/gripper"/>
     <framepos name="box" objtype="body" objname="box"/>

mjpc/tasks/common_assets/reorientation_cube.xml

@@ -26,7 +26,6 @@
   </worldbody>
 
    <sensor>
-        <framepos name="trace0" objtype="body" objname="cube"/>
         <framepos name="cube_position" objtype="body" objname="cube"/>
         <framequat name="cube_orientation" objtype="body" objname="cube"/>
         <framelinvel name="cube_linear_velocity" objtype="body" objname="cube"/>

mjpc/tasks/manipulation/manipulation.cc

@@ -16,19 +16,17 @@
 
 #include <string>
 
-#include <absl/container/flat_hash_map.h>
 #include <absl/random/random.h>
 #include <mujoco/mjmodel.h>
 #include <mujoco/mujoco.h>
-#include "mjpc/task.h"
 #include "mjpc/tasks/manipulation/common.h"
 #include "mjpc/utilities.h"
 
 namespace mjpc {
 std::string manipulation::Bring::XmlPath() const {
   return GetModelPath("manipulation/task_panda_bring.xml");
 }
-std::string manipulation::Bring::Name() const { return "Panda Robotiq Bring"; }
+std::string manipulation::Bring::Name() const { return "PickAndPlace"; }
 
 void manipulation::Bring::ResidualFn::Residual(const mjModel* model,
                                                const mjData* data,

mjpc/tasks/panda/panda.cc

@@ -25,7 +25,7 @@ namespace mjpc {
 std::string Panda::XmlPath() const {
   return GetModelPath("panda/task.xml");
 }
-std::string Panda::Name() const { return "Panda"; }
+std::string Panda::Name() const { return "Pick"; }
 
 // ---------- Residuals for in-panda manipulation task ---------
 //   Number of residuals: 5