Separate planner and estimator model loading

mjpc/agent.cc

@@ -68,11 +68,21 @@ Agent::Agent(const mjModel* model, std::shared_ptr<Task> task)
 }
 
 // initialize data, settings, planners, state
-void Agent::Initialize(const mjModel* model) {
+void Agent::Initialize(const mjModel* model, const mjModel* estimator_model) {
   // ----- model ----- //
+
+  // planner model
   mjModel* old_model = model_;
   model_ = mj_copyModel(nullptr, model);  // agent's copy of model
 
+  // estimator model
+  if (model_estimator_) mj_deleteModel(model_estimator_);
+  if (estimator_model) {
+    model_estimator_ = mj_copyModel(nullptr, estimator_model);
+  } else {
+    model_estimator_ = mj_copyModel(nullptr, model);
+  }
+
   // check for limits on all actuators
   int num_missing = 0;
   for (int i = 0; i < model_->nu; i++) {
@@ -120,7 +130,7 @@ void Agent::Initialize(const mjModel* model) {
   // initialize estimator
   if (reset_estimator && estimator_enabled) {
     for (const auto& estimator : estimators_) {
-      estimator->Initialize(model_);
+      estimator->Initialize(model_estimator_);
       estimator->Reset();
     }
   }

mjpc/agent.h

@@ -54,12 +54,15 @@ class Agent {
   // destructor
   ~Agent() {
     if (model_) mj_deleteModel(model_);  // we made a copy in Initialize
+    if (model_estimator_)
+      mj_deleteModel(model_estimator_);  // we made a copy in Initialize
   }
 
   // ----- methods ----- //
 
   // initialize data, settings, planners, states
-  void Initialize(const mjModel* model);
+  void Initialize(const mjModel* model,
+                  const mjModel* estimator_model = nullptr);
 
   // allocate memory
   void Allocate();
@@ -120,6 +123,8 @@ class Agent {
   std::string GetTaskNames() const { return task_names_; }
   int GetTaskIdByName(std::string_view name) const;
   std::string GetTaskXmlPath(int id) const { return tasks_[id]->XmlPath(); }
+  std::string GetPlannerXmlPath(int id) const { return tasks_[id]->PlannerXmlPath(); }
+  std::string GetEstimatorXmlPath(int id) const { return tasks_[id]->EstimatorXmlPath(); }
 
   // load the latest task model, based on GUI settings
   struct LoadModelResult {
@@ -190,6 +195,7 @@ class Agent {
  private:
   // model
   mjModel* model_ = nullptr;
+  mjModel* model_estimator_ = nullptr;
 
   UniqueMjModel model_override_ = {nullptr, mj_deleteModel};
 

mjpc/app.cc

@@ -219,14 +219,31 @@ void PhysicsLoop(mj::Simulate& sim) {
 
     // ----- task reload ----- //
     if (sim.uiloadrequest.load() == 1) {
+      // get new estimator model
+      sim.filename = sim.agent->GetEstimatorXmlPath(sim.agent->gui_task_id);
+
+      mjModel* mest = nullptr;
+      if (!sim.filename.empty()) {
+        mest = LoadModel(sim.agent.get(), sim);
+      }
+
+      // get new planner model
+      sim.filename = sim.agent->GetPlannerXmlPath(sim.agent->gui_task_id);
+
+      mjModel* mplan = nullptr;
+      if (!sim.filename.empty()) {
+        mplan = LoadModel(sim.agent.get(), sim);
+        sim.agent->Initialize(mplan, mest);
+      }
+
       // get new model + task
       sim.filename = sim.agent->GetTaskXmlPath(sim.agent->gui_task_id);
 
       mjModel* mnew = LoadModel(sim.agent.get(), sim);
       mjData* dnew = nullptr;
       if (mnew) dnew = mj_makeData(mnew);
       if (dnew) {
-        sim.agent->Initialize(mnew);
+        if (mplan == nullptr) sim.agent->Initialize(mnew, mest);
         sim.agent->plot_enabled = absl::GetFlag(FLAGS_show_plot);
         sim.agent->plan_enabled = absl::GetFlag(FLAGS_planner_enabled);
         sim.agent->Allocate();
@@ -414,6 +431,23 @@ MjpcApp::MjpcApp(std::vector<std::shared_ptr<mjpc::Task>> tasks, int task_id) {
   sim->agent->SetTaskList(std::move(tasks));
   sim->agent->gui_task_id = task_id;
 
+  // estimator setup
+  sim->agent->estimator_enabled = absl::GetFlag(FLAGS_estimator_enabled);
+  mjModel* mest = nullptr;
+  sim->filename = sim->agent->GetEstimatorXmlPath(sim->agent->gui_task_id);
+  if (!sim->filename.empty()) {
+    mest = LoadModel(sim->agent.get(), *sim);
+  }
+
+  // load planner model
+  mjModel* mplan = nullptr;
+  sim->filename = sim->agent->GetPlannerXmlPath(sim->agent->gui_task_id);
+  if (!sim->filename.empty()) {
+    mplan = LoadModel(sim->agent.get(), *sim);
+    sim->agent->Initialize(mplan, mest);
+  }
+
+  // load task model
   sim->filename = sim->agent->GetTaskXmlPath(sim->agent->gui_task_id);
   m = LoadModel(sim->agent.get(), *sim);
   if (m) d = mj_makeData(m);
@@ -431,8 +465,7 @@ MjpcApp::MjpcApp(std::vector<std::shared_ptr<mjpc::Task>> tasks, int task_id) {
   mju_zero(ctrlnoise, m->nu);
 
   // agent
-  sim->agent->estimator_enabled = absl::GetFlag(FLAGS_estimator_enabled);
-  sim->agent->Initialize(m);
+  if (mplan == nullptr) sim->agent->Initialize(m, mest);
   sim->agent->Allocate();
   sim->agent->Reset();
   sim->agent->PlotInitialize();

mjpc/task.h

@@ -123,6 +123,8 @@ class Task {
 
   virtual std::string Name() const = 0;
   virtual std::string XmlPath() const = 0;
+  virtual std::string PlannerXmlPath() const { return ""; };
+  virtual std::string EstimatorXmlPath() const { return ""; };
 
   // mode
   int mode;

mjpc/tasks/CMakeLists.txt

@@ -46,6 +46,9 @@ add_custom_target(
   COMMAND patch -o ${CMAKE_CURRENT_BINARY_DIR}/particle/particle_modified.xml
           ${CMAKE_CURRENT_BINARY_DIR}/particle/particle.xml
           <${CMAKE_CURRENT_SOURCE_DIR}/particle/particle.xml.patch
+  COMMAND patch -o ${CMAKE_CURRENT_BINARY_DIR}/particle/particle_modified_red.xml
+          ${CMAKE_CURRENT_BINARY_DIR}/particle/particle.xml
+          <${CMAKE_CURRENT_SOURCE_DIR}/particle/particle_red.xml.patch
   # swimmer
   COMMAND ${CMAKE_COMMAND} -E copy
           ${dm_control_SOURCE_DIR}/dm_control/suite/swimmer.xml

mjpc/tasks/particle/particle.cc

@@ -22,7 +22,10 @@
 namespace mjpc {
 
 std::string Particle::XmlPath() const {
-  return GetModelPath("particle/task_timevarying.xml");
+  return GetModelPath("particle/task_timevarying_task.xml");
+}
+std::string Particle::PlannerXmlPath() const {
+  return GetModelPath("particle/task_timevarying_plan.xml");
 }
 std::string Particle::Name() const { return "Particle"; }
 

mjpc/tasks/particle/particle.h

@@ -25,6 +25,7 @@ class Particle : public Task {
  public:
   std::string Name() const override;
   std::string XmlPath() const override;
+  std::string PlannerXmlPath() const override;
   class ResidualFn : public mjpc::BaseResidualFn {
    public:
     explicit ResidualFn(const Particle* task) : mjpc::BaseResidualFn(task) {}

mjpc/tasks/particle/particle_red.xml.patch

@@ -0,0 +1,57 @@
+diff --git a/particle_modified.xml b/particle_modified.xml
+--- a/particle_modified.xml
++++ b/particle_modified.xml
+@@ -1,9 +1,5 @@
+-<mujoco model="planar point mass">
+-  <include file="./common/skybox.xml"/>
+-  <include file="./common/visual.xml"/>
+-  <include file="./common/materials.xml"/>
+-
+-  <option timestep="0.02">
++<mujoco model="Particle">
++  <option timestep="0.01">
+     <flag contact="disable"/>
+   </option>
+
+@@ -13,9 +9,12 @@
+   </default>
+
+   <worldbody>
++    <body name="goal" mocap="true" pos="0.25 0 0.01" quat="1 0 0 0">
++        <geom type="sphere" size=".01" contype="0" conaffinity="0" rgba="1 0 0 .5"/>
++    </body>
+     <light name="light" pos="0 0 1"/>
+     <camera name="fixed" pos="0 0 .75" quat="1 0 0 0"/>
+-    <geom name="ground" type="plane" pos="0 0 0" size=".3 .3 .1" material="grid"/>
++    <geom name="ground" type="plane" pos="0 0 0" size=".3 .3 .1" material="blue_grid"/>
+     <geom name="wall_x" type="plane" pos="-.3 0 .02" zaxis="1 0 0"  size=".02 .3 .02" material="decoration"/>
+     <geom name="wall_y" type="plane" pos="0 -.3 .02" zaxis="0 1 0"  size=".3 .02 .02" material="decoration"/>
+     <geom name="wall_neg_x" type="plane" pos=".3 0 .02" zaxis="-1 0 0"  size=".02 .3 .02" material="decoration"/>
+@@ -26,24 +25,12 @@
+       <joint name="root_x" type="slide"  pos="0 0 0" axis="1 0 0" />
+       <joint name="root_y" type="slide"  pos="0 0 0" axis="0 1 0" />
+       <geom name="pointmass" type="sphere" size=".01" material="self" mass=".3"/>
++      <site name="tip" pos="0 0 0" size="0.01"/>
+     </body>
+-
+-    <geom name="target" pos="0 0 .01" material="target" type="sphere" size=".015"/>
+   </worldbody>
+
+-  <tendon>
+-    <fixed name="t1">
+-      <joint joint="root_x" coef="1"/>
+-      <joint joint="root_y" coef="0"/>
+-    </fixed>
+-    <fixed name="t2">
+-      <joint joint="root_x" coef="0"/>
+-      <joint joint="root_y" coef="1"/>
+-    </fixed>
+-  </tendon>
+-
+   <actuator>
+-    <motor name="t1" tendon="t1"/>
+-    <motor name="t2" tendon="t2"/>
++    <motor name="x_motor" joint="root_x" gear="1" ctrllimited="true" ctrlrange="-1 1"/>
++    <motor name="y_motor" joint="root_y" gear="1" ctrllimited="true" ctrlrange="-1 1"/>
+   </actuator>
+ </mujoco>

mjpc/tasks/particle/task_timevarying_task.xml

@@ -0,0 +1,46 @@
+<mujoco model="Particle Control">
+  <include file="../common.xml"/>
+  <!-- modified from: https://github.com/google-deepmind/dm_control/blob/main/dm_control/suite/point_mass.xml -->
+  <include file="particle_modified_red.xml" />
+
+  <size memory="10K"/>
+
+  <custom>
+    <!-- agent -->
+    <numeric name="task_risk" data="1" />
+    <numeric name="agent_planner" data="0" />
+    <numeric name="agent_horizon" data="0.5" />
+    <numeric name="agent_timestep" data="0.01" />
+    <numeric name="agent_policy_width" data="0.0015" />
+    <numeric name="sampling_spline_points" data="5" />
+    <numeric name="sampling_exploration" data="0.01" />
+    <numeric name="gradient_spline_points" data="5" />
+
+    <!-- estimator -->
+    <numeric name="estimator" data="0" />
+    <numeric name="estimator_sensor_start" data="3" />
+    <numeric name="estimator_number_sensor" data="3" />
+    <numeric name="estimator_timestep" data="0.01" />
+    <numeric name="batch_configuration_length" data="10" />
+    <numeric name="batch_scale_prior" data="1.0e-3" />
+  </custom>
+
+  <sensor>
+    <!-- cost -->
+    <user name="Position" dim="2" user="0 5.0 0.0 10.0" />
+    <user name="Velocity" dim="2" user="0 0.1 0.0 1.0" />
+    <user name="Control" dim="2" user="0 0.1 0.0 1.0" />
+
+    <!-- estimator -->
+    <jointpos name="x" joint="root_x" />
+    <jointpos name="y" joint="root_y" />
+    <accelerometer name="acc" site="tip" />
+
+    <!-- trace -->
+    <framepos name="trace0" objtype="site" objname="tip"/>
+
+    <!-- residual -->
+    <framepos name="position" objtype="site" objname="tip"/>
+    <framelinvel name="velocity" objtype="site" objname="tip"/>
+  </sensor>
+</mujoco>