Merge pull request #302 from thowell/estimator_update

Add mode input to Estimator::Update
google-deepmind · Mar 3, 2024 · a01e100 · a01e100
2 parents c1b15e5 + 15221bd
commit a01e100
Show file tree

Hide file tree

Showing 10 changed files with 20 additions and 12 deletions.
diff --git a/mjpc/estimators/batch.cc b/mjpc/estimators/batch.cc
@@ -282,7 +282,7 @@ void Batch::Reset(const mjData* data) {
 }
 
 // update
-void Batch::Update(const double* ctrl, const double* sensor) {
+void Batch::Update(const double* ctrl, const double* sensor, int mode) {
   // start timer
   auto start = std::chrono::steady_clock::now();
 

diff --git a/mjpc/estimators/batch.h b/mjpc/estimators/batch.h
@@ -58,7 +58,7 @@ class Batch : public Direct, public Estimator {
   void Reset(const mjData* data = nullptr) override;
 
   // update
-  void Update(const double* ctrl, const double* sensor) override;
+  void Update(const double* ctrl, const double* sensor, int mode = 0) override;
 
   // get state
   double* State() override { return state.data(); };

diff --git a/mjpc/estimators/estimator.h b/mjpc/estimators/estimator.h
@@ -43,7 +43,8 @@ class Estimator {
 
   // TODO(etom): time input
   // update
-  virtual void Update(const double* ctrl, const double* sensor) = 0;
+  virtual void Update(const double* ctrl, const double* sensor,
+                      int mode = 0) = 0;
 
   // get state
   virtual double* State() = 0;
@@ -197,7 +198,7 @@ class GroundTruth : public Estimator {
   }
 
   // update
-  void Update(const double* ctrl, const double* sensor) override {
+  void Update(const double* ctrl, const double* sensor, int mode = 0) override {
     mju_copy(data_->qpos, state.data(), model->nq);
     mju_copy(data_->qvel, state.data() + model->nq, model->nv);
     mju_copy(data_->act, state.data() + model->nq + model->nv, model->na);

diff --git a/mjpc/estimators/kalman.h b/mjpc/estimators/kalman.h
@@ -54,12 +54,12 @@ class Kalman : public Estimator {
   void UpdatePrediction();
 
   // update
-  void Update(const double* ctrl, const double* sensor) override {
+  void Update(const double* ctrl, const double* sensor, int mode = 0) override {
     // correct state with latest measurement
-    UpdateMeasurement(ctrl, sensor);
+    if (mode == 0 || mode == 1) UpdateMeasurement(ctrl, sensor);
 
     // propagate state forward in time with model
-    UpdatePrediction();
+    if (mode == 0 || mode == 2) UpdatePrediction();
 
     // set time
     time = data_->time;

diff --git a/mjpc/estimators/unscented.cc b/mjpc/estimators/unscented.cc
@@ -481,7 +481,7 @@ void Unscented::SigmaCovariances() {
 }
 
 // unscented filter update
-void Unscented::Update(const double* ctrl, const double* sensor) {
+void Unscented::Update(const double* ctrl, const double* sensor, int mode) {
   // start timer
   auto start = std::chrono::steady_clock::now();
 

diff --git a/mjpc/estimators/unscented.h b/mjpc/estimators/unscented.h
@@ -61,7 +61,7 @@ class Unscented : public Estimator {
   void SigmaCovariances();
 
   // update
-  void Update(const double* ctrl, const double* sensor) override;
+  void Update(const double* ctrl, const double* sensor, int mode = 0) override;
 
   // quaternion means
   void QuaternionMeans();

diff --git a/mjpc/grpc/filter.proto b/mjpc/grpc/filter.proto
@@ -49,6 +49,8 @@ message ResetResponse {}
 message UpdateRequest {
   repeated double ctrl = 1 [packed = true];
   repeated double sensor = 2 [packed = true];
+  optional int32 mode = 3;
+
 }
 
 message UpdateResponse {}

diff --git a/mjpc/grpc/filter_service.cc b/mjpc/grpc/filter_service.cc
@@ -128,8 +128,9 @@ grpc::Status FilterService::Update(grpc::ServerContext* context,
     return {grpc::StatusCode::FAILED_PRECONDITION, "Init not called."};
   }
 
-  // measurement update
-  filters_[filter_]->Update(request->ctrl().data(), request->sensor().data());
+  // update
+  filters_[filter_]->Update(request->ctrl().data(), request->sensor().data(),
+                            request->mode());
 
   return grpc::Status::OK;
 }

diff --git a/python/mujoco_mpc/filter.py b/python/mujoco_mpc/filter.py
@@ -158,11 +158,13 @@ def update(
       self,
       ctrl: Optional[npt.ArrayLike] = [],
       sensor: Optional[npt.ArrayLike] = [],
+      mode: Optional[int] = 0,
   ):
     # request
     request = filter_pb2.UpdateRequest(
         ctrl=ctrl,
         sensor=sensor,
+        mode=mode,
     )
 
     # response

diff --git a/python/mujoco_mpc/filter_test.py b/python/mujoco_mpc/filter_test.py
@@ -62,10 +62,12 @@ def test_updates(self):
     self.assertLess(np.linalg.norm(noise["process"] - process), 1.0e-5)
     self.assertLess(np.linalg.norm(noise["sensor"] - sensor), 1.0e-5)
 
-    # measurement update
+    # update
     ctrl = np.random.normal(scale=1.0, size=model.nu)
     sensor = np.random.normal(scale=1.0, size=model.nsensordata)
     filter.update(ctrl=ctrl, sensor=sensor)
+    filter.update(ctrl=ctrl, sensor=sensor, mode=0)
+    filter.update(ctrl=ctrl, sensor=sensor, mode=1)
 
     # TODO(etom): more tests