Merge branch 'main' into direct_optimizer_python

Author: thowell

cmake/MujocoLinkOptions.cmake

@@ -23,7 +23,7 @@ function(get_mujoco_extra_link_options OUTPUT_VAR)
     set(EXTRA_LINK_OPTIONS)
 
     if(WIN32)
-      set(CMAKE_REQUIRED_LINK_OPTIONS "-fuse-ld=lld-link")
+      set(CMAKE_REQUIRED_FLAGS "-fuse-ld=lld-link")
       check_c_source_compiles("int main() {}" SUPPORTS_LLD)
       if(SUPPORTS_LLD)
         set(EXTRA_LINK_OPTIONS
@@ -34,24 +34,24 @@ function(get_mujoco_extra_link_options OUTPUT_VAR)
         )
       endif()
     else()
-      set(CMAKE_REQUIRED_LINK_OPTIONS "-fuse-ld=lld")
+      set(CMAKE_REQUIRED_FLAGS "-fuse-ld=lld")
       check_c_source_compiles("int main() {}" SUPPORTS_LLD)
       if(SUPPORTS_LLD)
         set(EXTRA_LINK_OPTIONS ${EXTRA_LINK_OPTIONS} -fuse-ld=lld)
       else()
-        set(CMAKE_REQUIRED_LINK_OPTIONS "-fuse-ld=gold")
+        set(CMAKE_REQUIRED_FLAGS "-fuse-ld=gold")
         check_c_source_compiles("int main() {}" SUPPORTS_GOLD)
         if(SUPPORTS_GOLD)
           set(EXTRA_LINK_OPTIONS ${EXTRA_LINK_OPTIONS} -fuse-ld=gold)
         endif()
       endif()
 
-      set(CMAKE_REQUIRED_LINK_OPTIONS ${EXTRA_LINK_OPTIONS} "-Wl,--gc-sections")
+      set(CMAKE_REQUIRED_FLAGS ${EXTRA_LINK_OPTIONS} "-Wl,--gc-sections")
       check_c_source_compiles("int main() {}" SUPPORTS_GC_SECTIONS)
       if(SUPPORTS_GC_SECTIONS)
         set(EXTRA_LINK_OPTIONS ${EXTRA_LINK_OPTIONS} -Wl,--gc-sections)
       else()
-        set(CMAKE_REQUIRED_LINK_OPTIONS ${EXTRA_LINK_OPTIONS} "-Wl,-dead_strip")
+        set(CMAKE_REQUIRED_FLAGS ${EXTRA_LINK_OPTIONS} "-Wl,-dead_strip")
         check_c_source_compiles("int main() {}" SUPPORTS_DEAD_STRIP)
         if(SUPPORTS_DEAD_STRIP)
           set(EXTRA_LINK_OPTIONS ${EXTRA_LINK_OPTIONS} -Wl,-dead_strip)

mjpc/agent.h

@@ -136,6 +136,7 @@ class Agent {
   mjpc::Planner& ActivePlanner() const { return *planners_[planner_]; }
   mjpc::Estimator& ActiveEstimator() const { return *estimators_[estimator_]; }
   int ActiveEstimatorIndex() const { return estimator_; }
+  double ComputeTime() const { return agent_compute_time_; }
   Task* ActiveTask() const { return tasks_[active_task_id_].get(); }
   // a residual function that can be used from trajectory rollouts. must only
   // be used from trajectory rollout threads (no locking).

mjpc/direct/direct.h

@@ -16,7 +16,6 @@
 #define MJPC_DIRECT_OPTIMIZER_H_
 
 #include <memory>
-#include <mutex>
 #include <string>
 #include <vector>
 

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"
 

mjpc/grpc/agent_service.cc

@@ -118,6 +118,11 @@ grpc::Status AgentService::Init(grpc::ServerContext* context,
   model = mj_copyModel(nullptr, agent_model);
   data_ = mj_makeData(model);
   rollout_data_.reset(mj_makeData(model));
+  int home_id = mj_name2id(model, mjOBJ_KEY, "home");
+  if (home_id >= 0) {
+    mj_resetDataKeyframe(model, data_, home_id);
+    mj_resetDataKeyframe(model, rollout_data_.get(), home_id);
+  }
   mjcb_sensor = residual_sensor_callback;
 
   agent_.SetState(data_);

mjpc/planners/cross_entropy/planner.cc

@@ -14,18 +14,19 @@
 
 #include "mjpc/planners/cross_entropy/planner.h"
 
-#include <absl/random/random.h>
-#include <mujoco/mujoco.h>
-
 #include <algorithm>
 #include <chrono>
 #include <cmath>
-#include <mutex>
 #include <shared_mutex>
 
+#include <absl/random/random.h>
+#include <mujoco/mujoco.h>
 #include "mjpc/array_safety.h"
-#include "mjpc/planners/policy.h"
+#include "mjpc/planners/planner.h"
+#include "mjpc/planners/sampling/planner.h"
 #include "mjpc/states/state.h"
+#include "mjpc/task.h"
+#include "mjpc/threadpool.h"
 #include "mjpc/trajectory.h"
 #include "mjpc/utilities.h"
 
@@ -47,9 +48,6 @@ void CrossEntropyPlanner::Initialize(mjModel* model, const Task& task) {
   // task
   this->task = &task;
 
-  // rollout parameters
-  timestep_power = 1.0;
-
   // sampling noise
   std_initial_ =
       GetNumberOrDefault(0.1, model,
@@ -138,7 +136,7 @@ void CrossEntropyPlanner::Reset(int horizon,
 
   // variance
   double var = std_initial_ * std_initial_;
-  fill(variance.begin(), variance.end(), var);
+  std::fill(variance.begin(), variance.end(), var);
 
   // trajectory samples
   for (int i = 0; i < kMaxTrajectory; i++) {
@@ -362,11 +360,8 @@ void CrossEntropyPlanner::ResamplePolicy(int horizon) {
   mju_copy(resampled_policy.times.data(), times_scratch.data(),
            num_spline_points);
 
-  // time step power scaling
-  PowerSequence(resampled_policy.times.data(), time_shift,
-                resampled_policy.times[0],
-                resampled_policy.times[num_spline_points - 1], timestep_power,
-                num_spline_points);
+  LinearRange(resampled_policy.times.data(), time_shift,
+              resampled_policy.times[0], num_spline_points);
 }
 
 // add random noise to nominal policy
@@ -513,8 +508,6 @@ void CrossEntropyPlanner::GUI(mjUI& ui) {
       {mjITEM_SELECT, "Spline", 2, &policy.representation,
        "Zero\nLinear\nCubic"},
       {mjITEM_SLIDERINT, "Spline Pts", 2, &policy.num_spline_points, "0 1"},
-      // {mjITEM_SLIDERNUM, "Spline Pow. ", 2, &timestep_power, "0 10"},
-      // {mjITEM_SELECT, "Noise type", 2, &noise_type, "Gaussian\nUniform"},
       {mjITEM_SLIDERNUM, "Init. Std", 2, &std_initial_, "0 1"},
       {mjITEM_SLIDERNUM, "Min. Std", 2, &std_min_, "0.01 0.5"},
       {mjITEM_SLIDERINT, "Elite", 2, &n_elite_, "2 128"},

mjpc/planners/cross_entropy/planner.h

@@ -15,16 +15,16 @@
 #ifndef MJPC_PLANNERS_CROSS_ENTROPY_PLANNER_H_
 #define MJPC_PLANNERS_CROSS_ENTROPY_PLANNER_H_
 
-#include <mujoco/mujoco.h>
-
 #include <atomic>
-#include <memory>
 #include <shared_mutex>
 #include <vector>
 
+#include <mujoco/mujoco.h>
 #include "mjpc/planners/planner.h"
-#include "mjpc/planners/sampling/planner.h"
+#include "mjpc/planners/sampling/policy.h"
 #include "mjpc/states/state.h"
+#include "mjpc/task.h"
+#include "mjpc/threadpool.h"
 #include "mjpc/trajectory.h"
 
 namespace mjpc {
@@ -116,9 +116,6 @@ class CrossEntropyPlanner : public Planner {
   // order of indices of rolled out trajectories, ordered by total return
   std::vector<int> trajectory_order;
 
-  // rollout parameters
-  double timestep_power;
-
   // ----- noise ----- //
   double std_initial_;  // standard deviation for sampling normal: N(0,
                         // std)

mjpc/planners/gradient/planner.cc

@@ -49,9 +49,6 @@ void GradientPlanner::Initialize(mjModel* model, const Task& task) {
   // task
   this->task = &task;
 
-  // rollout parameters
-  timestep_power = 1.0;
-
   // dimensions
   dim_state = model->nq + model->nv + model->na;  // state dimension
   dim_state_derivative =
@@ -374,11 +371,8 @@ void GradientPlanner::ResamplePolicy(int horizon) {
   mju_copy(candidate_policy[0].times.data(), times_scratch.data(),
            num_spline_points);
 
-  // time step power scaling
-  PowerSequence(candidate_policy[0].times.data(), time_shift,
-                candidate_policy[0].times[0],
-                candidate_policy[0].times[num_spline_points - 1],
-                timestep_power, num_spline_points);
+  LinearRange(candidate_policy[0].times.data(), time_shift,
+              candidate_policy[0].times[0], num_spline_points);
 }
 
 // compute candidate trajectories
@@ -474,7 +468,6 @@ void GradientPlanner::GUI(mjUI& ui) {
       {mjITEM_SELECT, "Spline", 2, &policy.representation,
        "Zero\nLinear\nCubic"},
       {mjITEM_SLIDERINT, "Spline Pts", 2, &policy.num_spline_points, "0 1"},
-      // {mjITEM_SLIDERNUM, "Spline Pow. ", 2, &timestep_power, "0 10"},
       {mjITEM_END}};
 
   // set number of trajectory slider limits

mjpc/planners/gradient/planner.h

@@ -123,9 +123,6 @@ class GradientPlanner : public Planner {
   Trajectory trajectory[kMaxTrajectory];
   int num_trajectory;
 
-  // rollout parameters
-  double timestep_power;
-
   // model derivatives
   ModelDerivatives model_derivative;
 

mjpc/planners/gradient/spline_mapping.h

@@ -25,11 +25,27 @@ namespace mjpc {
 // ----- spline constants ----- //
 inline constexpr int kMinGradientSplinePoints = 1;
 inline constexpr int kMaxGradientSplinePoints = 25;
-inline constexpr int kMinGradientSplinePower = 1;
-inline constexpr int kMaxGradientSplinePower = 5;
 
 // matrix representation for mapping between spline points and interpolated time
-// series
+// series.
+// A spline is made of num_input points, and each has one associated time, and
+// `dim` associated parameters.
+// The time series has num_output entries, each associated with a time and with
+// dim associated values.
+//
+// For sampling policies, we have
+// dim = model->nu
+// num_input = num_spline_points
+// num_output = trajectory_length
+//
+// The mapping is a matrix, A, of shape (dim*num_output) x (dim*num_input),
+// which can be used to go from spline parameters to the values of the sampled
+// time series, assuming a fixed set of times.
+//
+// Given a vector of containing spline parameters, v (length=dim*num_input),
+// flattened so that the parameters for each spline point are next to each
+// other, A*v gives the corresponding interpolated values, sampled at
+// output_times.
 class SplineMapping {
  public:
   // constructor