Modify solver settings for model derivative finite difference computations and warmstart with accelerations.

thowell · copybara-github · commit 0c76c9581ee9 · 2025-03-04T09:41:54.000-08:00
PiperOrigin-RevId: 733362111
Change-Id: I85cc7d7fae3178d8ccc4acfa8171810302d9a045
diff --git a/mjpc/agent.cc b/mjpc/agent.cc
@@ -87,7 +87,8 @@ void Agent::Initialize(const mjModel* model) {
   }
 
   // planner
-  planner_ = GetNumberOrDefault(0, model, "agent_planner");
+  planner_ =
+      static_cast<PlannerType>(GetNumberOrDefault(0, model, "agent_planner"));
 
   // estimator
   estimator_ =
diff --git a/mjpc/agent.h b/mjpc/agent.h
@@ -213,7 +213,7 @@ class Agent {
 
   // planners
   std::vector<std::unique_ptr<mjpc::Planner>> planners_;
-  int planner_;
+  PlannerType planner_;
 
   // estimators
   std::vector<std::unique_ptr<mjpc::Estimator>> estimators_;
diff --git a/mjpc/planners/gradient/planner.cc b/mjpc/planners/gradient/planner.cc
@@ -203,9 +203,9 @@ void GradientPlanner::OptimizePolicy(int horizon, ThreadPool& pool) {
     // compute model and sensor Jacobians
     model_derivative.Compute(
         model, data_, trajectory[0].states.data(), trajectory[0].actions.data(),
-        trajectory[0].times.data(), dim_state, dim_state_derivative, dim_action,
-        dim_sensor, horizon, settings.fd_tolerance, settings.fd_mode, pool,
-        skip);
+        trajectory[0].accelerations.data(), trajectory[0].times.data(),
+        dim_state, dim_state_derivative, dim_action, dim_sensor, horizon,
+        settings.fd_tolerance, settings.fd_mode, pool, skip);
 
     // stop timer
     model_derivative_time += GetDuration(model_derivative_start);
diff --git a/mjpc/planners/ilqg/planner.cc b/mjpc/planners/ilqg/planner.cc
@@ -395,6 +395,7 @@ void iLQGPlanner::Iteration(int horizon, ThreadPool& pool) {
   model_derivative.Compute(
       model, data_, candidate_policy[0].trajectory.states.data(),
       candidate_policy[0].trajectory.actions.data(),
+      candidate_policy[0].trajectory.accelerations.data(),
       candidate_policy[0].trajectory.times.data(), dim_state,
       dim_state_derivative, dim_action, dim_sensor, horizon,
       settings.fd_tolerance, settings.fd_mode, pool, derivative_skip_);
diff --git a/mjpc/planners/model_derivatives.cc b/mjpc/planners/model_derivatives.cc
@@ -42,10 +42,10 @@ void ModelDerivatives::Reset(int dim_state_derivative, int dim_action,
 }
 
 // compute derivatives at all time steps
-void ModelDerivatives::Compute(const mjModel* m,
+void ModelDerivatives::Compute(mjModel* m,
                                const std::vector<UniqueMjData>& data,
                                const double* x, const double* u,
-                               const double* h, int dim_state,
+                               const double* a, const double* h, int dim_state,
                                int dim_state_derivative, int dim_action,
                                int dim_sensor, int T, double tol, int mode,
                                ThreadPool& pool, int skip) {
@@ -71,10 +71,23 @@ void ModelDerivatives::Compute(const mjModel* m,
     }
   }
 
+  // warmstart
+  int saved_flags = m->opt.disableflags;
+  m->opt.disableflags &= ~mjDSBL_WARMSTART;
+
+  // solver settings
+  int saved_iterations = m->opt.iterations;
+  mjtNum saved_tolerance = m->opt.tolerance;
+  if (m->opt.solver == mjSOL_NEWTON) {
+    m->opt.iterations = 1;
+    m->opt.tolerance = 0.0;
+  }
+  // TODO(taylorhowell): settings for CG and PGS
+
   // evaluate derivatives
   int count_before = pool.GetCount();
   for (int t : evaluate_) {
-    pool.Schedule([&m, &data, &A = A, &B = B, &C = C, &D = D, &x, &u, &h,
+    pool.Schedule([&m, &data, &A = A, &B = B, &C = C, &D = D, &x, &u, &a, &h,
                    dim_state, dim_state_derivative, dim_action, dim_sensor, tol,
                    mode, t, T]() {
       mjData* d = data[ThreadPool::WorkerId()].get();
@@ -85,6 +98,9 @@ void ModelDerivatives::Compute(const mjModel* m,
       // set action
       mju_copy(d->ctrl, u + t * dim_action, dim_action);
 
+      // set acceleration
+      mju_copy(d->qacc_warmstart, a + t * m->nv, m->nv);
+
       // Jacobians
       if (t == T - 1) {
         // Jacobians
@@ -105,6 +121,11 @@ void ModelDerivatives::Compute(const mjModel* m,
   pool.WaitCount(count_before + evaluate_.size());
   pool.ResetCount();
 
+  // restore settings
+  m->opt.tolerance = saved_tolerance;
+  m->opt.iterations = saved_iterations;
+  m->opt.disableflags = saved_flags;
+
   // interpolate derivatives
   count_before = pool.GetCount();
   for (int t : interpolate_) {
diff --git a/mjpc/planners/model_derivatives.h b/mjpc/planners/model_derivatives.h
@@ -42,10 +42,11 @@ class ModelDerivatives {
   void Reset(int dim_state_derivative, int dim_action, int dim_sensor, int T);
 
   // compute derivatives at all time steps
-  void Compute(const mjModel* m, const std::vector<UniqueMjData>& data,
-               const double* x, const double* u, const double* h, int dim_state,
-               int dim_state_derivative, int dim_action, int dim_sensor, int T,
-               double tol, int mode, ThreadPool& pool, int skip = 0);
+  void Compute(mjModel* m, const std::vector<UniqueMjData>& data,
+               const double* x, const double* u, const double* a,
+               const double* h, int dim_state, int dim_state_derivative,
+               int dim_action, int dim_sensor, int T, double tol, int mode,
+               ThreadPool& pool, int skip = 0);
 
   // Jacobians
   std::vector<double> A;  // model Jacobians wrt state
diff --git a/mjpc/test/agent/agent_test.cc b/mjpc/test/agent/agent_test.cc
@@ -21,6 +21,7 @@
 
 #include "gtest/gtest.h"
 #include <mujoco/mujoco.h>
+#include "mjpc/planners/include.h"
 #include "mjpc/planners/ilqs/planner.h"
 #include "mjpc/planners/sampling/planner.h"
 #include "mjpc/task.h"
@@ -71,7 +72,7 @@ class AgentTest : public ::testing::Test {
     // test
     EXPECT_EQ(agent->integrator_, 0);
     EXPECT_NEAR(agent->timestep_, 0.1, 1.0e-5);
-    EXPECT_EQ(agent->planner_, 0);
+    EXPECT_EQ(agent->planner_, kSamplingPlanner);
     EXPECT_NEAR(agent->horizon_, 1, 1.0e-5);
     EXPECT_EQ(agent->steps_, 11);
     EXPECT_FALSE(agent->plan_enabled);
@@ -152,7 +153,7 @@ class AgentTest : public ::testing::Test {
         0.0, 1.0e-1);
 
     // ----- switch to iLQG planner ----- //
-    agent->planner_ = 2;
+    agent->planner_ = kILQGPlanner;
     agent->Allocate();
     agent->Reset();
     exitrequest.store(false);
@@ -215,7 +216,7 @@ class AgentTest : public ::testing::Test {
     agent->plan_enabled = true;
 
     bool success = false;
-    agent->planner_ = 0;  // sampling
+    agent->planner_ = kSamplingPlanner;
     reinterpret_cast<SamplingPlanner*>(&agent->ActivePlanner())
         ->num_trajectory_ = 128;
 
@@ -278,7 +279,7 @@ class AgentTest : public ::testing::Test {
     agent->Reset();
     agent->plan_enabled = true;
 
-    agent->planner_ = 2;  // iLQG
+    agent->planner_ = kILQGPlanner;
 
     agent->Reset();
     data->qpos[0] = 0;
@@ -328,7 +329,7 @@ class AgentTest : public ::testing::Test {
     agent->Reset();
     agent->plan_enabled = true;
 
-    agent->planner_ = 3;  // iLQS
+    agent->planner_ = kILQSPlanner;
     iLQSPlanner* planner =
         reinterpret_cast<iLQSPlanner*>(&agent->ActivePlanner());
 
diff --git a/mjpc/trajectory.cc b/mjpc/trajectory.cc
@@ -59,6 +59,9 @@ void Trajectory::Allocate(int T) {
 
   // traces
   trace.resize(dim_trace * T);
+
+  // accelerations
+  accelerations.resize(dim_state * T); // TODO(taylorhowell): allocate nv x T
 }
 
 // reset memory to zeros
@@ -86,6 +89,9 @@ void Trajectory::Reset(int T, const double* initial_repeated_action) {
 
   // traces
   std::fill(trace.begin(), trace.begin() + dim_trace * T, 0.0);
+
+  // accelerations
+  std::fill(accelerations.begin(), accelerations.begin() + dim_state * T, 0.0);
 }
 
 // simulate model forward in time with continuous-time indexed policy
@@ -177,6 +183,9 @@ void Trajectory::NoisyRollout(
     mju_copy(DataAt(states, (t + 1) * dim_state + nq), data->qvel, nv);
     mju_copy(DataAt(states, (t + 1) * dim_state + nq + nv), data->act, na);
     times[t + 1] = data->time;
+
+    // record acceleration
+    mju_copy(DataAt(accelerations, t * nv), data->qacc, nv);
   }
 
   // check for step warnings
@@ -205,6 +214,9 @@ void Trajectory::NoisyRollout(
   GetTraces(DataAt(trace, (horizon - 1) * 3 * task->num_trace), model, data,
             task->num_trace);
 
+  // final acceleration
+  mju_copy(DataAt(accelerations, (horizon - 1) * nv), data->qacc, nv);
+
   // compute return
   UpdateReturn(task);
 }
@@ -276,6 +288,9 @@ void Trajectory::RolloutDiscrete(
     mju_copy(DataAt(states, (t + 1) * dim_state + nq), data->qvel, nv);
     mju_copy(DataAt(states, (t + 1) * dim_state + nq + nv), data->act, na);
     times[t + 1] = data->time;
+
+    // record acceleration
+    mju_copy(DataAt(accelerations, t * nv), data->qacc, nv);
   }
 
   // check for step warnings
@@ -304,6 +319,9 @@ void Trajectory::RolloutDiscrete(
   GetTraces(DataAt(trace, (horizon - 1) * 3 * task->num_trace), model, data,
             task->num_trace);
 
+  // final acceleration
+  mju_copy(DataAt(accelerations, (horizon - 1) * nv), data->qacc, nv);
+
   // compute return
   UpdateReturn(task);
 }
diff --git a/mjpc/trajectory.h b/mjpc/trajectory.h
@@ -71,19 +71,20 @@ class Trajectory {
       double time, const double* mocap, const double* userdata, int steps);
 
   // ----- members ----- //
-  int horizon;                   // trajectory length
-  int dim_state;                 // states dimension
-  int dim_action;                // actions dimension
-  int dim_residual;              // residual dimension
-  int dim_trace;                 // traces dimension
-  std::vector<double> states;    // (horizon   x nq + nv + na)
-  std::vector<double> actions;   // (horizon-1 x num_action)
-  std::vector<double> times;     // horizon
-  std::vector<double> residual;  // (horizon   x num_residual)
-  std::vector<double> costs;     // horizon
-  std::vector<double> trace;     // (horizon   x 3)
-  double total_return;           // (1)
-  bool failure;                  // true if last rollout had a warning
+  int horizon;                        // trajectory length
+  int dim_state;                      // states dimension
+  int dim_action;                     // actions dimension
+  int dim_residual;                   // residual dimension
+  int dim_trace;                      // traces dimension
+  std::vector<double> states;         // (horizon   x nq + nv + na)
+  std::vector<double> actions;        // (horizon-1 x num_action)
+  std::vector<double> times;          // horizon
+  std::vector<double> residual;       // (horizon   x num_residual)
+  std::vector<double> costs;          // horizon
+  std::vector<double> trace;          // (horizon   x 3)
+  double total_return;                // (1)
+  bool failure;                       // true if last rollout had a warning
+  std::vector<double> accelerations;  // (horizon x nv)
 
  private:
   // calculates total_return and costs

Original file line number	Diff line number	Diff line change
`@@ -87,7 +87,8 @@ void Agent::Initialize(const mjModel* model) {`
`87`	`87`	`}`
`88`	`88`
`89`	`89`	`// planner`
`90`		`- planner_ = GetNumberOrDefault(0, model, "agent_planner");`
	`90`	`+ planner_ =`
	`91`	`+ static_cast<PlannerType>(GetNumberOrDefault(0, model, "agent_planner"));`
`91`	`92`
`92`	`93`	`// estimator`
`93`	`94`	`estimator_ =`