Merge of 14adceb

PiperOrigin-RevId: 602664992
Change-Id: I12999d88e09e267fcc357e12a51f21a641e80f06

Author: copybara-github

mjpc/planners/cross_entropy/planner.cc

@@ -241,7 +241,8 @@ void CrossEntropyPlanner::OptimizePolicy(int horizon, ThreadPool& pool) {
   // copy first elite trajectory
   mju_copy(elite_avg.actions.data(), trajectory[idx].actions.data(),
            model->nu * (horizon - 1));
-  mju_copy(elite_avg.trace.data(), trajectory[idx].trace.data(), 3 * horizon);
+  mju_copy(elite_avg.trace.data(), trajectory[idx].trace.data(),
+           trajectory[idx].dim_trace * horizon);
   mju_copy(elite_avg.residual.data(), trajectory[idx].residual.data(),
            elite_avg.dim_residual * horizon);
   mju_copy(elite_avg.costs.data(), trajectory[idx].costs.data(), horizon);
@@ -260,7 +261,7 @@ void CrossEntropyPlanner::OptimizePolicy(int horizon, ThreadPool& pool) {
     mju_addTo(elite_avg.actions.data(), trajectory[idx].actions.data(),
               model->nu * (horizon - 1));
     mju_addTo(elite_avg.trace.data(), trajectory[idx].trace.data(),
-              3 * horizon);
+              trajectory[idx].dim_trace * horizon);
     mju_addTo(elite_avg.residual.data(), trajectory[idx].residual.data(),
               elite_avg.dim_residual * horizon);
     mju_addTo(elite_avg.costs.data(), trajectory[idx].costs.data(), horizon);
@@ -273,7 +274,7 @@ void CrossEntropyPlanner::OptimizePolicy(int horizon, ThreadPool& pool) {
   mju_scl(elite_avg.actions.data(), elite_avg.actions.data(), 1.0 / n_elite,
           model->nu * (horizon - 1));
   mju_scl(elite_avg.trace.data(), elite_avg.trace.data(), 1.0 / n_elite,
-          3 * horizon);
+          elite_avg.dim_trace * horizon);
   mju_scl(elite_avg.residual.data(), elite_avg.residual.data(), 1.0 / n_elite,
           elite_avg.dim_residual * horizon);
   mju_scl(elite_avg.costs.data(), elite_avg.costs.data(), 1.0 / n_elite,

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>