Correctly compute plan variance in cross-entropy planner.

Commit 734bda2 accidentally changed the variance computation from
```
sum((candidate_policy[idx] - avg)**2 for idx in elites) / (len(elites) - 1)
```

To:
```
sum((candidate_policy[0] - avg)**2 for idx in elites) / (len(elites) - 1)
```

I can't visually see a difference in performance on an in-hand reorientation task.
Fixes #350.

PiperOrigin-RevId: 696463663
Change-Id: I8f13e426d883cb01565d8f051ab29bf8a9a67471

mjpc/planners/cross_entropy/planner.cc

@@ -227,12 +227,13 @@ void CrossEntropyPlanner::OptimizePolicy(int horizon, ThreadPool& pool) {
   for (int i = 0; i < n_elite; i++) {
     // ordered trajectory index
     int idx = trajectory_order[i];
+    const TimeSpline& elite_plan = candidate_policy[idx].plan;
 
     // add parameters
-    for (int i = 0; i < num_spline_points; i++) {
-      TimeSpline::Node n = candidate_policy[idx].plan.NodeAt(i);
+    for (int t = 0; t < num_spline_points; t++) {
+      TimeSpline::ConstNode n = elite_plan.NodeAt(t);
       for (int j = 0; j < model->nu; j++) {
-        parameters_scratch[i * model->nu + j] += n.values()[j];
+        parameters_scratch[t * model->nu + j] += n.values()[j];
       }
     }
 
@@ -247,12 +248,15 @@ void CrossEntropyPlanner::OptimizePolicy(int horizon, ThreadPool& pool) {
 
   // loop over elites to compute variance
   std::fill(variance.begin(), variance.end(), 0.0);  // reset variance to zero
-  for (int t = 0; t < num_spline_points; t++) {
-    TimeSpline::Node n = candidate_policy[trajectory_order[0]].plan.NodeAt(t);
-    for (int j = 0; j < model->nu; j++) {
-      // average
-      double p_avg = parameters_scratch[t * model->nu + j];
-      for (int i = 0; i < n_elite; i++) {
+  for (int i = 0; i < n_elite; i++) {
+    int idx = trajectory_order[i];
+    const TimeSpline& elite_plan = candidate_policy[idx].plan;
+    for (int t = 0; t < num_spline_points; t++) {
+      TimeSpline::ConstNode n = elite_plan.NodeAt(t);
+      for (int j = 0; j < model->nu; j++) {
+        // average
+        double p_avg = parameters_scratch[t * model->nu + j];
+
         // candidate parameter
         double pi = n.values()[j];
         double diff = pi - p_avg;