Merge branch 'main' into direct_optimizer_python

thowell · web-flow · commit cf0fade4f1ba · 2024-02-03T20:16:04.000-08:00
diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml
@@ -68,7 +68,7 @@ jobs:
         $cmake_extra_args
     - name: Build MuJoCo MPC
       working-directory: build
-      run: cmake --build . --config=Release ${{ matrix.cmake_build_args }} --target mjpc agent_test cost_derivatives_test norm_test rollout_test threadpool_test trajectory_test utilities_test direct_force_test direct_optimize_test direct_parameter_test direct_sensor_test direct_trajectory_test direct_utilities_test batch_filter_test batch_prior_test kalman_test unscented_test ${{ matrix.additional_targets }}
+      run: cmake --build . --config=Release ${{ matrix.cmake_build_args }} --target mjpc agent_test agent_utilities_test cost_derivatives_test norm_test rollout_test threadpool_test trajectory_test direct_force_test direct_optimize_test direct_parameter_test direct_sensor_test direct_trajectory_test direct_utilities_test batch_filter_test batch_prior_test kalman_test unscented_test cubic_test gradient_planner_test gradient_test linear_test zero_test backward_pass_test ilqg_test robust_planner_test sampling_planner_test state_test task_test utilities_test ${{ matrix.additional_targets }}
     - name: Test MuJoCo MPC
       working-directory: build
       run: ctest -C Release --output-on-failure .
diff --git a/mjpc/CMakeLists.txt b/mjpc/CMakeLists.txt
@@ -175,6 +175,28 @@ if(APPLE)
   target_link_libraries(mjpc "-framework Cocoa")
 endif()
 
+add_executable(
+  testspeed
+  testspeed_app.cc
+  testspeed.h
+  testspeed.cc
+)
+target_link_libraries(
+  testspeed
+  absl::flags
+  absl::flags_parse
+  absl::random_random
+  absl::strings
+  libmjpc
+  mujoco::mujoco
+  threadpool
+  Threads::Threads
+)
+target_include_directories(testspeed PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/..)
+target_compile_options(testspeed PUBLIC ${MJPC_COMPILE_OPTIONS})
+target_link_options(testspeed PRIVATE ${MJPC_LINK_OPTIONS})
+target_compile_definitions(testspeed PRIVATE MJSIMULATE_STATIC)
+
 add_subdirectory(tasks)
 
 if(BUILD_TESTING AND MJPC_BUILD_TESTS)
diff --git a/mjpc/agent.cc b/mjpc/agent.cc
@@ -73,6 +73,19 @@ void Agent::Initialize(const mjModel* model) {
   if (model_) mj_deleteModel(model_);
   model_ = mj_copyModel(nullptr, model);  // agent's copy of model
 
+  // check for limits on all actuators
+  int num_missing = 0;
+  for (int i = 0; i < model_->nu; i++) {
+    if (!model_->actuator_ctrllimited[i]) {
+      num_missing++;
+      printf("%s (actuator %i) missing limits\n",
+             model_->names + model_->name_actuatoradr[i], i);
+    }
+  }
+  if (num_missing > 0) {
+    mju_error("Ctrl limits required for all actuators.\n");
+  }
+
   // planner
   planner_ = GetNumberOrDefault(0, model, "agent_planner");
 
diff --git a/mjpc/direct/direct.h b/mjpc/direct/direct.h
@@ -12,8 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#ifndef MJPC_DIRECT_OPTIMIZER_H_
-#define MJPC_DIRECT_OPTIMIZER_H_
+#ifndef MJPC_DIRECT_DIRECT_H_
+#define MJPC_DIRECT_DIRECT_H_
 
 #include <memory>
 #include <string>
@@ -64,7 +64,7 @@ class Direct {
       : model_parameters_(LoadModelParameters()), pool_(num_threads) {}
 
   // constructor
-  Direct(const mjModel* model, int length = 3, int max_history = 0);
+  explicit Direct(const mjModel* model, int length = 3, int max_history = 0);
 
   // destructor
   virtual ~Direct() {
@@ -509,4 +509,4 @@ std::string StatusString(int code);
 
 }  // namespace mjpc
 
-#endif  // MJPC_DIRECT_OPTIMIZER_H_
+#endif  // MJPC_DIRECT_DIRECT_H_
diff --git a/mjpc/planners/gradient/planner.h b/mjpc/planners/gradient/planner.h
@@ -12,8 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#ifndef MJPC_PLANNERS_GRADIENT_OPTIMIZER_H_
-#define MJPC_PLANNERS_GRADIENT_OPTIMIZER_H_
+#ifndef MJPC_PLANNERS_GRADIENT_PLANNER_H_
+#define MJPC_PLANNERS_GRADIENT_PLANNER_H_
 
 #include <memory>
 #include <shared_mutex>
@@ -164,4 +164,4 @@ class GradientPlanner : public Planner {
 
 }  // namespace mjpc
 
-#endif  // MJPC_PLANNERS_GRADIENT_OPTIMIZER_H_
+#endif  // MJPC_PLANNERS_GRADIENT_PLANNER_H_
diff --git a/mjpc/planners/ilqg/planner.h b/mjpc/planners/ilqg/planner.h
@@ -12,8 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#ifndef MJPC_PLANNERS_ILQG_OPTIMIZER_H_
-#define MJPC_PLANNERS_ILQG_OPTIMIZER_H_
+#ifndef MJPC_PLANNERS_ILQG_PLANNER_H_
+#define MJPC_PLANNERS_ILQG_PLANNER_H_
 
 #include <shared_mutex>
 #include <vector>
@@ -161,4 +161,4 @@ class iLQGPlanner : public Planner {
 
 }  // namespace mjpc
 
-#endif  // MJPC_PLANNERS_ILQG_OPTIMIZER_H_
+#endif  // MJPC_PLANNERS_ILQG_PLANNER_H_
diff --git a/mjpc/planners/ilqs/planner.cc b/mjpc/planners/ilqs/planner.cc
@@ -114,8 +114,8 @@ void iLQSPlanner::OptimizePolicy(int horizon, ThreadPool& pool) {
     if (dim_actions != sampling.model->nu * (horizon - 1) ||
         dim_parameters != sampling.model->nu * num_spline_points) {
       // dimension
-      int dim_parameters = sampling.model->nu * num_spline_points;
-      int dim_actions = sampling.model->nu * (horizon - 1);
+      dim_parameters = sampling.model->nu * num_spline_points;
+      dim_actions = sampling.model->nu * (horizon - 1);
 
       // compute parameter to action mapping
       mappings[sampling.policy.representation]->Compute(
diff --git a/mjpc/planners/ilqs/planner.h b/mjpc/planners/ilqs/planner.h
@@ -12,11 +12,12 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#ifndef MJPC_PLANNERS_ILQS_OPTIMIZER_H_
-#define MJPC_PLANNERS_ILQS_OPTIMIZER_H_
+#ifndef MJPC_PLANNERS_ILQS_PLANNER_H_
+#define MJPC_PLANNERS_ILQS_PLANNER_H_
 
 #include <mujoco/mujoco.h>
 
+#include <memory>
 #include <shared_mutex>
 #include <vector>
 
@@ -110,4 +111,4 @@ class iLQSPlanner : public Planner {
 
 }  // namespace mjpc
 
-#endif  // MJPC_PLANNERS_ILQS_OPTIMIZER_H_
+#endif  // MJPC_PLANNERS_ILQS_PLANNER_H_
diff --git a/mjpc/planners/sampling/planner.cc b/mjpc/planners/sampling/planner.cc
@@ -433,9 +433,10 @@ void SamplingPlanner::Plots(mjvFigure* fig_planner, mjvFigure* fig_timer,
 
   // ----- timer ----- //
 
-  PlotUpdateData(
-      fig_timer, timer_bounds, fig_timer->linedata[0 + timer_shift][0] + 1,
-      1.0e-3 * noise_compute_time * planning, 100, 0 + timer_shift, 0, 1, -100);
+  PlotUpdateData(fig_timer, timer_bounds,
+                 fig_timer->linedata[0 + timer_shift][0] + 1,
+                 1.0e-3 * noise_compute_time * planning, 100,
+                 0 + timer_shift, 0, 1, -100);
 
   PlotUpdateData(fig_timer, timer_bounds,
                  fig_timer->linedata[1 + timer_shift][0] + 1,
diff --git a/mjpc/tasks/cube/cube_3x3x3.xml.patch b/mjpc/tasks/cube/cube_3x3x3.xml.patch
@@ -23,7 +23,7 @@ diff --git a/cube_3x3x3_modified.xml b/cube_3x3x3_modified.xml
      <default class="cubelet">
 -      <joint type="ball" armature="0.0001" damping="0.0005" frictionloss="0.001"/>
 -      <geom type="mesh" condim="1" mesh="cubelet" euler="0 0 90"/>
-+      <joint type="ball" armature="0.00005" damping="0.0001" frictionloss="0.00005"/>
++      <joint type="ball" armature="0.0001" damping="0.0005" frictionloss="0.00005"/>
 +      <geom type="mesh" condim="1" mesh="cubelet" quat="1 0 0 1"/>
      </default>
      <default class="core">
diff --git a/mjpc/tasks/cube/solve.cc b/mjpc/tasks/cube/solve.cc
@@ -15,6 +15,7 @@
 #include "mjpc/tasks/cube/solve.h"
 
 #include <algorithm>
+#include <iostream>
 #include <random>
 #include <string>
 
@@ -108,11 +109,17 @@ void CubeSolve::ResidualFn::Residual(const mjModel* model, const mjData* data,
   counter += 6;
 
   // ---------- Residual (4) ----------
+
+  // The unmodified cube model has 20 ball joints: nq=86, nv=66.
+  // The patch adds a free joint: nq=93, nv=72.
+  // The task adds a ball joint: nq=97, nv=75.
+  // The shadow hand has 24 DoFs: nq=121, nv=99.
+  // The following two residuals apply for the last 24 entries of qpos and qvel:
   mju_sub(residual + counter, data->qpos + 97, model->key_qpos + 97, 24);
   counter += 24;
 
   // ---------- Residual (5) ----------
-  mju_copy(residual + counter, data->qvel + 97, 24);
+  mju_copy(residual + counter, data->qvel + 75, 24);
   counter += 24;
 
   // ---------- Residual (6) ----------
@@ -131,6 +138,7 @@ void CubeSolve::ResidualFn::Residual(const mjModel* model, const mjData* data,
 void CubeSolve::TransitionLocked(mjModel* model, mjData* data) {
   if (transition_model_) {
     if (mode == kModeWait) {
+      weight[11] = .01;  // add penalty on joint movement
       // wait
     } else if (mode == kModeScramble) {  // scramble
       double scramble_param = parameters[6];
@@ -190,9 +198,11 @@ void CubeSolve::TransitionLocked(mjModel* model, mjData* data) {
 
       // set face goal index
       goal_index_ = num_scramble - 1;
+      std::cout << "rotations required: " << num_scramble << "\n";
 
       // set to solve
       mode = kModeSolve;
+      weight[11] = 0;  // remove penalty on joint movement
     } else if (mode == kModeSolve) {  // solve
       // set goal
       mju_copy(parameters.data(), goal_cache_.data() + 6 * goal_index_, 6);
@@ -204,7 +214,9 @@ void CubeSolve::TransitionLocked(mjModel* model, mjData* data) {
       if (mju_norm(error, 6) < 0.085) {
         if (goal_index_ == 0) {
           mode = kModeWait;
+          std::cout << "solved!\n";
         } else {
+          std::cout << "rotations remaining: " << goal_index_ << "\n";
           goal_index_--;
         }
       }
@@ -213,11 +225,10 @@ void CubeSolve::TransitionLocked(mjModel* model, mjData* data) {
 
   // check for drop
   if (data->qpos[6] < kResetHeight) {
-    // reset cube position + orientation
-    mju_copy(data->qpos, model->key_qpos, 7);
+    if (mode != kModeWait) { std::cout << "cube fell\n"; }
 
-    // reset cube velocity
-    mju_zero(data->qvel, 6);
+    // stop optimization
+    mode = kModeWait;
   }
 
   // check goal index
diff --git a/mjpc/tasks/cube/task.xml b/mjpc/tasks/cube/task.xml
@@ -12,8 +12,8 @@
     <numeric name="agent_policy_width" data="0.0035" />
     <numeric name="sampling_spline_points" data="6" />
     <numeric name="sampling_exploration" data="0.1" />
-    <numeric name="sampling_trajectories" data="20" />
-    <numeric name="sampling_representation" data="0" />
+    <numeric name="sampling_trajectories" data="60" />
+    <numeric name="sampling_representation" data="1" />
     <!-- manual face goals -->
     <numeric name="residual_Red" data="0 -3.14 3.14"/>
     <numeric name="residual_Orange" data="0 -3.14 3.14"/>
diff --git a/mjpc/tasks/tasks.cc b/mjpc/tasks/tasks.cc
@@ -17,6 +17,7 @@
 #include <memory>
 #include <vector>
 
+#include "mjpc/task.h"
 #include "mjpc/tasks/acrobot/acrobot.h"
 #include "mjpc/tasks/cube/solve.h"
 #include "mjpc/tasks/cartpole/cartpole.h"
diff --git a/mjpc/test/agent/agent_utilities_test.cc b/mjpc/test/agent/agent_utilities_test.cc
@@ -234,25 +234,6 @@ TEST(AgentUtilitiesTest, Clamp) {
   EXPECT_NEAR(x[2], 0.0, 1.0e-5);
 }
 
-TEST(AgentUtilitiesTest, PowerSequence) {
-  // sequence
-  double sequence[4] = {0.2, 0.3, 0.4, 0.5};
-  int length = 4;
-  double step = 0.1;
-
-  // power
-  double power = 2.0;
-
-  // power sequence
-  PowerSequence(sequence, step, sequence[0], sequence[3], power, length);
-
-  // test
-  EXPECT_NEAR(sequence[0], 0.2, 1.0e-5);
-  EXPECT_NEAR(sequence[1], 0.27142857, 1.0e-5);
-  EXPECT_NEAR(sequence[2], 0.37142857, 1.0e-5);
-  EXPECT_NEAR(sequence[3], 0.5, 1.0e-5);
-}
-
 TEST(AgentUtilitiesTest, FindInterval) {
   // sequence
   std::vector<double> sequence{-1.0, 0.0, 1.0, 2.0};
diff --git a/mjpc/testspeed.cc b/mjpc/testspeed.cc
diff --git a/mjpc/testspeed.h b/mjpc/testspeed.h
diff --git a/mjpc/testspeed_app.cc b/mjpc/testspeed_app.cc
diff --git a/mjpc/testspeed_test.cc b/mjpc/testspeed_test.cc
