nimrod comments

google-deepmind · Feb 9, 2024 · 79db2bc · 79db2bc
1 parent c392c2a
commit 79db2bc
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 11 deletions.
diff --git a/mjpc/planners/gradient/planner.cc b/mjpc/planners/gradient/planner.cc
@@ -194,6 +194,7 @@ void GradientPlanner::OptimizePolicy(int horizon, ThreadPool& pool) {
 
   // update policy
   double c_best = c_prev;
+  int skip = derivative_skip_;
   for (int i = 0; i < settings.max_rollout; i++) {
     // ----- model derivatives ----- //
     // start timer
@@ -204,7 +205,7 @@ void GradientPlanner::OptimizePolicy(int horizon, ThreadPool& pool) {
         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,
-        derivative_skip_);
+        skip);
 
     // stop timer
     model_derivative_time += GetDuration(model_derivative_start);

diff --git a/mjpc/planners/model_derivatives.cc b/mjpc/planners/model_derivatives.cc
@@ -63,9 +63,11 @@ void ModelDerivatives::Compute(const mjModel* m,
   evaluate_.push_back(T - 1);
 
   // interpolate indices
-  for (int t = 0; t < T; t++) {
-    if (std::find(evaluate_.begin(), evaluate_.end(), t) == evaluate_.end()) {
+  for (int t = 0, e = 0; t < T; t++) {
+    if (e == evaluate_.size() || evaluate_[e] > t) {
       interpolate_.push_back(t);
+    } else {
+      e++;
     }
   }
 
@@ -123,35 +125,35 @@ void ModelDerivatives::Compute(const mjModel* m,
       // A
       int nA = dim_state_derivative * dim_state_derivative;
       double* Ai = DataAt(A, t * nA);
-      double* AL = DataAt(A, e0 * nA);
-      double* AU = DataAt(A, e1 * nA);
+      const double* AL = DataAt(A, e0 * nA);
+      const double* AU = DataAt(A, e1 * nA);
 
       mju_scl(Ai, AL, 1.0 - tt, nA);
       mju_addToScl(Ai, AU, tt, nA);
 
       // B
       int nB = dim_state_derivative * dim_action;
       double* Bi = DataAt(B, t * nB);
-      double* BL = DataAt(B, e0 * nB);
-      double* BU = DataAt(B, e1 * nB);
+      const double* BL = DataAt(B, e0 * nB);
+      const double* BU = DataAt(B, e1 * nB);
 
       mju_scl(Bi, BL, 1.0 - tt, nB);
       mju_addToScl(Bi, BU, tt, nB);
 
       // C
       int nC = dim_sensor * dim_state_derivative;
       double* Ci = DataAt(C, t * nC);
-      double* CL = DataAt(C, e0 * nC);
-      double* CU = DataAt(C, e1 * nC);
+      const double* CL = DataAt(C, e0 * nC);
+      const double* CU = DataAt(C, e1 * nC);
 
       mju_scl(Ci, CL, 1.0 - tt, nC);
       mju_addToScl(Ci, CU, tt, nC);
 
       // D
       int nD = dim_sensor * dim_action;
       double* Di = DataAt(D, t * nD);
-      double* DL = DataAt(D, e0 * nD);
-      double* DU = DataAt(D, e1 * nD);
+      const double* DL = DataAt(D, e0 * nD);
+      const double* DU = DataAt(D, e1 * nD);
 
       mju_scl(Di, DL, 1.0 - tt, nD);
       mju_addToScl(Di, DU, tt, nD);