Merge of 97ab026

copybara-github · copybara-github · commit 7c421b8439fb · 2024-01-26T02:43:15.000-08:00
PiperOrigin-RevId: 601710890
Change-Id: If7fbe394ca443f015317a2af69e2033b9280e1b8
diff --git a/mjpc/direct/direct.cc b/mjpc/direct/direct.cc
@@ -148,9 +148,6 @@ void Direct::Initialize(const mjModel* model) {
   act.Initialize(na, configuration_length_);
   times.Initialize(1, configuration_length_);
 
-  // ctrl
-  ctrl.Initialize(model->nu, configuration_length_);
-
   // prior
   configuration_previous.Initialize(nq, configuration_length_);
 
@@ -348,7 +345,6 @@ void Direct::Reset(const mjData* data) {
   acceleration.Reset();
   act.Reset();
   times.Reset();
-  ctrl.Reset();
 
   // prior
   configuration_previous.Reset();
@@ -637,8 +633,6 @@ void Direct::SetConfigurationLength(int length) {
   act.SetLength(configuration_length_);
   times.SetLength(configuration_length_);
 
-  ctrl.SetLength(configuration_length_);
-
   configuration_previous.SetLength(configuration_length_);
 
   sensor_measurement.SetLength(configuration_length_);
@@ -1489,8 +1483,7 @@ void Direct::InverseDynamicsPrediction() {
   auto start = std::chrono::steady_clock::now();
 
   // dimension
-  int nq = model->nq, nv = model->nv, na = model->na, nu = model->nu,
-      ns = nsensordata_;
+  int nq = model->nq, nv = model->nv, na = model->na, ns = nsensordata_;
 
   // set parameters
   if (nparam_ > 0) {
@@ -1502,7 +1495,7 @@ void Direct::InverseDynamicsPrediction() {
   int count_before = pool_.GetCount();
 
   // first time step
-  pool_.Schedule([&batch = *this, nq, nv, nu]() {
+  pool_.Schedule([&batch = *this, nq, nv]() {
     // time index
     int t = 0;
 
@@ -1518,7 +1511,6 @@ void Direct::InverseDynamicsPrediction() {
     mju_copy(d->qpos, q0, nq);
     mju_zero(d->qvel, nv);
     mju_zero(d->qacc, nv);
-    mju_zero(d->ctrl, nu);
     d->time = batch.times.Get(t)[0];
 
     // position sensors
@@ -1551,12 +1543,11 @@ void Direct::InverseDynamicsPrediction() {
   // loop over predictions
   for (int t = 1; t < configuration_length_ - 1; t++) {
     // schedule
-    pool_.Schedule([&batch = *this, nq, nv, na, ns, nu, t]() {
+    pool_.Schedule([&batch = *this, nq, nv, na, ns, t]() {
       // terms
       double* qt = batch.configuration.Get(t);
       double* vt = batch.velocity.Get(t);
       double* at = batch.acceleration.Get(t);
-      double* ct = batch.ctrl.Get(t);
 
       // data
       mjData* d = batch.data_[t].get();
@@ -1565,7 +1556,6 @@ void Direct::InverseDynamicsPrediction() {
       mju_copy(d->qpos, qt, nq);
       mju_copy(d->qvel, vt, nv);
       mju_copy(d->qacc, at, nv);
-      mju_copy(d->ctrl, ct, nu);
 
       // inverse dynamics
       mj_inverse(batch.model, d);
@@ -1585,7 +1575,7 @@ void Direct::InverseDynamicsPrediction() {
   }
 
   // last time step
-  pool_.Schedule([&batch = *this, nq, nv, nu]() {
+  pool_.Schedule([&batch = *this, nq, nv]() {
     // time index
     int t = batch.ConfigurationLength() - 1;
 
@@ -1602,7 +1592,6 @@ void Direct::InverseDynamicsPrediction() {
     mju_copy(d->qpos, qT, nq);
     mju_copy(d->qvel, vT, nv);
     mju_zero(d->qacc, nv);
-    mju_zero(d->ctrl, nu);
     d->time = batch.times.Get(t)[0];
 
     // position sensors
@@ -1653,7 +1642,7 @@ void Direct::InverseDynamicsDerivatives() {
   auto start = std::chrono::steady_clock::now();
 
   // dimension
-  int nq = model->nq, nv = model->nv, nu = model->nu;
+  int nq = model->nq, nv = model->nv;
 
   // set parameters
   if (nparam_ > 0) {
@@ -1665,7 +1654,7 @@ void Direct::InverseDynamicsDerivatives() {
   int count_before = pool_.GetCount();
 
   // first time step
-  pool_.Schedule([&batch = *this, nq, nv, nu]() {
+  pool_.Schedule([&batch = *this, nq, nv]() {
     // time index
     int t = 0;
 
@@ -1680,7 +1669,6 @@ void Direct::InverseDynamicsDerivatives() {
     mju_copy(d->qpos, q0, nq);
     mju_zero(d->qvel, nv);
     mju_zero(d->qacc, nv);
-    mju_zero(d->ctrl, nu);
     d->time = batch.times.Get(t)[0];
 
     // finite-difference derivatives
@@ -1725,12 +1713,11 @@ void Direct::InverseDynamicsDerivatives() {
   // loop over predictions
   for (int t = 1; t < configuration_length_ - 1; t++) {
     // schedule
-    pool_.Schedule([&batch = *this, nq, nv, nu, t]() {
+    pool_.Schedule([&batch = *this, nq, nv, t]() {
       // unpack
       double* q = batch.configuration.Get(t);
       double* v = batch.velocity.Get(t);
       double* a = batch.acceleration.Get(t);
-      double* c = batch.ctrl.Get(t);
 
       double* dsdq = batch.block_sensor_configuration_.Get(t);
       double* dsdv = batch.block_sensor_velocity_.Get(t);
@@ -1743,11 +1730,10 @@ void Direct::InverseDynamicsDerivatives() {
       double* dadf = batch.block_force_acceleration_.Get(t);
       mjData* data = batch.data_[t].get();  // TODO(taylor): WorkerID
 
-      // set (state, acceleration) + ctrl
+      // set state, acceleration
       mju_copy(data->qpos, q, nq);
       mju_copy(data->qvel, v, nv);
       mju_copy(data->qacc, a, nv);
-      mju_copy(data->ctrl, c, nu);
 
       // finite-difference derivatives
       mjd_inverseFD(batch.model, data, batch.finite_difference.tolerance,
@@ -1767,7 +1753,7 @@ void Direct::InverseDynamicsDerivatives() {
   }
 
   // last time step
-  pool_.Schedule([&batch = *this, nq, nv, nu]() {
+  pool_.Schedule([&batch = *this, nq, nv]() {
     // time index
     int t = batch.ConfigurationLength() - 1;
 
@@ -1784,7 +1770,6 @@ void Direct::InverseDynamicsDerivatives() {
     mju_copy(d->qpos, qT, nq);
     mju_copy(d->qvel, vT, nv);
     mju_zero(d->qacc, nv);
-    mju_zero(d->ctrl, nu);
     d->time = batch.times.Get(t)[0];
 
     // finite-difference derivatives
@@ -2061,7 +2046,7 @@ double Direct::Cost(double* gradient, double* hessian) {
     // set dense rows in band Hessian
     if (hessian) {
       mju_copy(hessian + nvel_ * nband_, dense_parameter_.data(),
-             nparam_ * ntotal_);
+               nparam_ * ntotal_);
     }
   }
 
diff --git a/mjpc/direct/direct.h b/mjpc/direct/direct.h
@@ -16,7 +16,6 @@
 #define MJPC_DIRECT_OPTIMIZER_H_
 
 #include <memory>
-#include <mutex>
 #include <string>
 #include <vector>
 
@@ -161,7 +160,6 @@ class Direct {
   DirectTrajectory<double> acceleration;            // nv x T
   DirectTrajectory<double> act;                     // na x T
   DirectTrajectory<double> times;                   //  1 x T
-  DirectTrajectory<double> ctrl;                    // nu x T
   DirectTrajectory<double> sensor_measurement;      // ns x T
   DirectTrajectory<double> sensor_prediction;       // ns x T
   DirectTrajectory<int> sensor_mask;                // num_sensor x T
diff --git a/mjpc/estimators/batch.cc b/mjpc/estimators/batch.cc
@@ -24,7 +24,6 @@
 #include "mjpc/array_safety.h"
 #include "mjpc/estimators/estimator.h"
 #include "mjpc/direct/direct.h"
-#include "mjpc/norm.h"
 #include "mjpc/threadpool.h"
 #include "mjpc/utilities.h"
 
@@ -119,9 +118,6 @@ void Batch::Initialize(const mjModel* model) {
   act_cache_.Initialize(na, max_history_);
   times_cache_.Initialize(1, max_history_);
 
-  // ctrl
-  ctrl_cache_.Initialize(model->nu, max_history_);
-
   // prior
   configuration_previous_cache_.Initialize(nq, max_history_);
 
@@ -241,7 +237,6 @@ void Batch::Reset(const mjData* data) {
   acceleration_cache_.Reset();
   act_cache_.Reset();
   times_cache_.Reset();
-  ctrl_cache_.Reset();
 
   // prior
   configuration_previous_cache_.Reset();
@@ -325,9 +320,6 @@ void Batch::Update(const double* ctrl, const double* sensor) {
   // set next time
   times.Set(&d->time, t + 1);
 
-  // set ctrl
-  this->ctrl.Set(ctrl, t);
-
   // set sensor
   sensor_measurement.Set(sensor + sensor_start_index_, t);
 
@@ -527,7 +519,6 @@ void Batch::Shift(int shift) {
   acceleration.Shift(shift);
   act.Shift(shift);
   times.Shift(shift);
-  ctrl.Shift(shift);
 
   configuration_previous.Shift(shift);
 
@@ -815,7 +806,6 @@ void Batch::ShiftResizeTrajectory(int new_head, int new_length) {
   acceleration_cache_.Reset();
   act_cache_.Reset();
   times_cache_.Reset();
-  ctrl_cache_.Reset();
   sensor_measurement_cache_.Reset();
   sensor_prediction_cache_.Reset();
   sensor_mask_cache_.Reset();
@@ -831,7 +821,6 @@ void Batch::ShiftResizeTrajectory(int new_head, int new_length) {
   acceleration_cache_.SetLength(length);
   act_cache_.SetLength(length);
   times_cache_.SetLength(length);
-  ctrl_cache_.SetLength(length);
   sensor_measurement_cache_.SetLength(length);
   sensor_prediction_cache_.SetLength(length);
   sensor_mask_cache_.SetLength(length);
@@ -846,7 +835,6 @@ void Batch::ShiftResizeTrajectory(int new_head, int new_length) {
     acceleration_cache_.Set(acceleration.Get(i), i);
     act_cache_.Set(act.Get(i), i);
     times_cache_.Set(times.Get(i), i);
-    ctrl_cache_.Set(ctrl.Get(i), i);
     sensor_measurement_cache_.Set(sensor_measurement.Get(i), i);
     sensor_prediction_cache_.Set(sensor_prediction.Get(i), i);
     sensor_mask_cache_.Set(sensor_mask.Get(i), i);
@@ -867,7 +855,6 @@ void Batch::ShiftResizeTrajectory(int new_head, int new_length) {
     acceleration.Set(acceleration_cache_.Get(new_head + i), i);
     act.Set(act_cache_.Get(new_head + i), i);
     times.Set(times_cache_.Get(new_head + i), i);
-    ctrl.Set(ctrl_cache_.Get(new_head + i), i);
     sensor_measurement.Set(sensor_measurement_cache_.Get(new_head + i), i);
     sensor_prediction.Set(sensor_prediction_cache_.Get(new_head + i), i);
     sensor_mask.Set(sensor_mask_cache_.Get(new_head + i), i);
diff --git a/mjpc/estimators/batch.h b/mjpc/estimators/batch.h
@@ -32,7 +32,7 @@
 namespace mjpc {
 
 // max filter history
-inline constexpr int kMaxFilterHistory = 128;
+inline constexpr int kMaxFilterHistory = 64;
 
 // ----- batch estimator ----- //
 // based on: "Physically-Consistent Sensor Fusion in Contact-Rich Behaviors"
@@ -239,7 +239,6 @@ class Batch : public Direct, public Estimator {
   DirectTrajectory<double> acceleration_cache_;            // nv x T
   DirectTrajectory<double> act_cache_;                     // na x T
   DirectTrajectory<double> times_cache_;                   //  1 x T
-  DirectTrajectory<double> ctrl_cache_;                    // nu x T
   DirectTrajectory<double> sensor_measurement_cache_;      // ns x T
   DirectTrajectory<double> sensor_prediction_cache_;       // ns x T
   DirectTrajectory<int> sensor_mask_cache_;                // num_sensor x T
diff --git a/mjpc/grpc/direct.proto b/mjpc/grpc/direct.proto
@@ -55,15 +55,14 @@ message Data {
   repeated double velocity = 2 [packed = true];
   repeated double acceleration = 3 [packed = true];
   repeated double time = 4 [packed = true];
-  repeated double ctrl = 5 [packed = true];
-  repeated double configuration_previous = 6 [packed = true];
-  repeated double sensor_measurement = 7 [packed = true];
-  repeated double sensor_prediction = 8 [packed = true];
-  repeated int32 sensor_mask = 9 [packed = true];
-  repeated double force_measurement = 10 [packed = true];
-  repeated double force_prediction = 11 [packed = true];
-  repeated double parameters = 12 [packed = true];
-  repeated double parameters_previous = 13 [packed = true];
+  repeated double configuration_previous = 5 [packed = true];
+  repeated double sensor_measurement = 6 [packed = true];
+  repeated double sensor_prediction = 7 [packed = true];
+  repeated int32 sensor_mask = 8 [packed = true];
+  repeated double force_measurement = 9 [packed = true];
+  repeated double force_prediction = 10 [packed = true];
+  repeated double parameters = 11 [packed = true];
+  repeated double parameters_previous = 12 [packed = true];
 }
 
 message DataRequest {
diff --git a/mjpc/grpc/direct_service.cc b/mjpc/grpc/direct_service.cc
@@ -176,19 +176,6 @@ grpc::Status DirectService::Data(grpc::ServerContext* context,
   double* time = optimizer_.times.Get(index);
   output->add_time(time[0]);
 
-  // set ctrl
-  int nu = optimizer_.model->nu;
-  if (input.ctrl_size() > 0) {
-    CHECK_SIZE("ctrl", nu, input.ctrl_size());
-    optimizer_.ctrl.Set(input.ctrl().data(), index);
-  }
-
-  // get ctrl
-  double* ctrl = optimizer_.ctrl.Get(index);
-  for (int i = 0; i < nu; i++) {
-    output->add_ctrl(ctrl[i]);
-  }
-
   // set previous configuration
   if (input.configuration_previous_size() > 0) {
     CHECK_SIZE("configuration_previous", nq,
diff --git a/mjpc/test/direct/direct_sensor_test.cc b/mjpc/test/direct/direct_sensor_test.cc
@@ -58,7 +58,6 @@ TEST(SensorCost, Particle) {
   // copy configuration, measurement
   mju_copy(optimizer.configuration.Data(), sim.qpos.Data(), nq * T);
   mju_copy(optimizer.sensor_measurement.Data(), sim.sensor.Data(), ns * T);
-  mju_copy(optimizer.ctrl.Data(), sim.ctrl.Data(), model->nu * T);
 
   // corrupt configurations
   absl::BitGen gen_;
diff --git a/python/mujoco_mpc/demos/direct/particle_smoother.py b/python/mujoco_mpc/demos/direct/particle_smoother.py
@@ -12,13 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import mujoco
 import matplotlib.pyplot as plt
 import mediapy as media
+import mujoco
+from mujoco_mpc import direct as direct_lib
 import numpy as np
 
-# set current directory to mjpc/python/mujoco_mpc
-from mujoco_mpc import direct as direct_lib
 # %%
 # 1D Particle Model
 xml = """
@@ -167,7 +166,6 @@
   optimizer.data(
       t,
       configuration=qinit[:, t],
-      ctrl=ctrl[:, t],
       sensor_measurement=noisy_sensor[:, t],
       force_measurement=qfrc[:, t],
       time=np.array([time[t]]),
diff --git a/python/mujoco_mpc/demos/direct/particle_trajopt.py b/python/mujoco_mpc/demos/direct/particle_trajopt.py
@@ -12,13 +12,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import mujoco
-import numpy as np
 import matplotlib.pyplot as plt
 import mediapy as media
-
-# set current directory to mjpc/python/mujoco_mpc
+import mujoco
 from mujoco_mpc import direct as direct_lib
+import numpy as np
+
+
 # %%
 # 2D Particle Model
 xml = """
@@ -163,7 +163,6 @@
   data_ = optimizer.data(
       t,
       configuration=qt,
-      ctrl=ct,
       sensor_measurement=st,
       sensor_mask=mt,
       force_measurement=ft,
diff --git a/python/mujoco_mpc/demos/filter/particle_drop.py b/python/mujoco_mpc/demos/filter/particle_drop.py
@@ -12,13 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import mujoco
-import numpy as np
 import matplotlib.pyplot as plt
 import mediapy as media
-
-# set current directory to mjpc/python/mujoco_mpc
+import mujoco
 from mujoco_mpc import filter as filter_lib
+import numpy as np
 
 # %%
 xml = """
diff --git a/python/mujoco_mpc/demos/predictive_sampling/cube_orientation.py b/python/mujoco_mpc/demos/predictive_sampling/cube_orientation.py
diff --git a/python/mujoco_mpc/demos/predictive_sampling/particle.py b/python/mujoco_mpc/demos/predictive_sampling/particle.py
diff --git a/python/mujoco_mpc/demos/predictive_sampling/predictive_sampling.py b/python/mujoco_mpc/demos/predictive_sampling/predictive_sampling.py
diff --git a/python/mujoco_mpc/direct.py b/python/mujoco_mpc/direct.py
diff --git a/python/mujoco_mpc/direct_test.py b/python/mujoco_mpc/direct_test.py
