Merge branch 'main' into direct_optimizer_python

Author: thowell

mjpc/tasks/cube/solve.cc

@@ -115,6 +115,11 @@ void CubeSolve::ResidualFn::Residual(const mjModel* model, const mjData* data,
   mju_copy(residual + counter, data->qvel + 97, 24);
   counter += 24;
 
+  // ---------- Residual (6) ----------
+  residual[counter++] =
+      goal_index_ * 12;  // each face has ~12 cost to unscramble based on
+                         // current weights, settings, etc.
+
   // sensor dim sanity check
   CheckSensorDim(model, counter);
 }
@@ -215,6 +220,11 @@ void CubeSolve::TransitionLocked(mjModel* model, mjData* data) {
     mju_zero(data->qvel, 6);
   }
 
+  // check goal index
+  if (residual_.goal_index_ != goal_index_) {
+    residual_.goal_index_ = goal_index_;
+  }
+
   // check for mode change
   if (residual_.current_mode_ != mode) {
     // update mode for residual

mjpc/tasks/cube/solve.h

@@ -32,14 +32,18 @@ class CubeSolve : public Task {
 
   class ResidualFn : public BaseResidualFn {
    public:
-    explicit ResidualFn(const CubeSolve* task, int current_mode = 0)
-        : BaseResidualFn(task), current_mode_(current_mode) {}
+    explicit ResidualFn(const CubeSolve* task, int current_mode = 0,
+                        int goal_index = 0)
+        : BaseResidualFn(task),
+          current_mode_(current_mode),
+          goal_index_(goal_index) {}
     void Residual(const mjModel* model, const mjData* data,
                   double* residual) const override;
 
    private:
     friend class CubeSolve;
     int current_mode_ = 0;
+    int goal_index_ = 0;
   };
 
   CubeSolve();
@@ -57,7 +61,8 @@ class CubeSolve : public Task {
 
  protected:
   std::unique_ptr<mjpc::ResidualFn> ResidualLocked() const override {
-    return std::make_unique<ResidualFn>(this, residual_.current_mode_);
+    return std::make_unique<ResidualFn>(this, residual_.current_mode_,
+                                        residual_.goal_index_);
   }
   ResidualFn* InternalResidual() override { return &residual_; }
 

mjpc/tasks/cube/task.xml

@@ -1,9 +1,11 @@
 <mujoco model="Cube Solving Manipulation">
   <include file="../common.xml"/>
 
+  <statistic extent="0.4" center="0.3 0.1 0"/>
   <size memory="1M"/>
 
   <custom>
+    <!-- agent/planner settings -->
     <numeric name="agent_planner" data="0" />
     <numeric name="agent_horizon" data="0.25" />
     <numeric name="agent_timestep" data="0.01" />
@@ -12,19 +14,19 @@
     <numeric name="sampling_exploration" data="0.1" />
     <numeric name="sampling_trajectories" data="20" />
     <numeric name="sampling_representation" data="0" />
+    <!-- manual face goals -->
     <numeric name="residual_Red" data="0 -3.14 3.14"/>
     <numeric name="residual_Orange" data="0 -3.14 3.14"/>
     <numeric name="residual_Blue" data="0 -3.14 3.14"/>
     <numeric name="residual_Green" data="0 -3.14 3.14"/>
     <numeric name="residual_White" data="0 -3.14 3.14"/>
     <numeric name="residual_Yellow" data="0 -3.14 3.14"/>
+    <!-- GUI elements -->
     <text name="task_transition" data="Scramble|Solve|Wait|Manual"/>
     <numeric name="residual_select_Scramble" data="3"/>
     <text name="residual_list_Scramble" data="1|2|3|4|5|6|7|8|9|10"/>
   </custom>
 
-  <statistic extent="0.4" center="0.3 0.1 0"/>
-
   <visual>
     <quality shadowsize="8192"/>
     <global azimuth="180" elevation="-30"/>
@@ -45,6 +47,7 @@
   </worldbody>
 
   <sensor>
+    <!-- cost terms -->
     <user name="In Hand" dim="3" user="1 50 0 100 0.02 2"/>
     <user name="Orientation" dim="3" user="0 0 0 10" />
     <user name="Cube Vel." dim="3" user="0 2.5 0 20" />
@@ -57,6 +60,8 @@
     <user name="Yellow" dim="1" user="0 10.0 0.0 25.0"/>
     <user name="Grasp" dim="24" user="0 1.0 0.0 10.0" />
     <user name="Joint Vel." dim="24" user="0 0.0 0.0 1.0e-1" />
+    <user name="Remaining" dim="1" user="-1 1.0 0.0 1.0" />
+    <!-- sensors -->
     <framepos name="palm_position" objtype="site" objname="grasp_site"/>
     <framequat name="cube_goal_orientation" objtype="body" objname="goal"/>
   </sensor>