Repository: vincekurtz/hydrax
Pull request #48: Optimization iterations
Author: login: Zolkin1
Created: 2025-04-18T01:00:45Z | Merged: 2025-04-29T01:28:07Z

I added an argument that lets the user choose the number of optimization iterations to make at each step.

The idea is that you can get better performance by performing the optimization update multiple times at the same state.

The iteration argument gets passed into the controller parameters.


Diff:
diff --git a/.gitignore b/.gitignore
index dc1fa724..5e58a68c 100644
--- a/.gitignore
+++ b/.gitignore
@@ -6,4 +6,5 @@ __pycache__
 *.pkl
 *.egg-info
 MUJOCO_LOG.TXT
-hydrax/recordings/
\ No newline at end of file
+hydrax/recordings/
+.idea
diff --git a/hydrax/alg_base.py b/hydrax/alg_base.py
index a56e2351..d7647a59 100644
--- a/hydrax/alg_base.py
+++ b/hydrax/alg_base.py
@@ -63,6 +63,7 @@ def __init__(
         plan_horizon: float,
         spline_type: Literal["zero", "linear", "cubic"] = "zero",
         num_knots: int = 4,
+        iterations: int = 1,
     ) -> None:
         """Initialize the MPC controller.
 
@@ -75,6 +76,7 @@ def __init__(
             spline_type: The type of spline used for control interpolation.
                          Defaults to "zero" (zero-order hold).
             num_knots: The number of knots in the control spline.
+            iterations: The number of optimization iterations to perform.
         """
         self.task = task
         self.num_randomizations = max(num_randomizations, 1)
@@ -101,6 +103,12 @@ def __init__(
         self.model = task.model
         self.randomized_axes = None
 
+        # Number of optimization iterations
+        if iterations < 1:
+            raise ValueError("iterations must be greater than 0!")
+
+        self.iterations = iterations
+
         if self.num_randomizations > 1:
             # Make domain randomized models
             rng = jax.random.key(seed)
@@ -135,23 +143,33 @@ def optimize(self, state: mjx.Data, params: Any) -> Tuple[Any, Trajectory]:
         new_mean = self.interp_func(new_tk, tk, params.mean[None, ...])[0]
         params = params.replace(tk=new_tk, mean=new_mean)
 
-        # Sample random control sequences from spline knots
-        knots, params = self.sample_knots(params)
-        knots = jnp.clip(
-            knots, self.task.u_min, self.task.u_max
-        )  # (num_rollouts, num_knots, nu)
-
-        # Roll out the control sequences, applying domain randomizations and
-        # combining costs using self.risk_strategy.
-        rng, dr_rng = jax.random.split(params.rng)
-        rollouts = self.rollout_with_randomizations(
-            state, new_tk, knots, dr_rng
+        def _optimize_scan_body(params: Any, _: Any):
+            # Sample random control sequences from spline knots
+            knots, params = self.sample_knots(params)
+            knots = jnp.clip(
+                knots, self.task.u_min, self.task.u_max
+            )  # (num_rollouts, num_knots, nu)
+
+            # Roll out the control sequences, applying domain randomizations and
+            # combining costs using self.risk_strategy.
+            rng, dr_rng = jax.random.split(params.rng)
+            rollouts = self.rollout_with_randomizations(
+                state, new_tk, knots, dr_rng
+            )
+            params = params.replace(rng=rng)
+
+            # Update the policy parameters based on the combined costs
+            params = self.update_params(params, rollouts)
+
+            return params, rollouts
+
+        params, rollouts = jax.lax.scan(
+            f=_optimize_scan_body, init=params, xs=jnp.arange(self.iterations)
         )
-        params = params.replace(rng=rng)
 
-        # Update the policy parameters based on the combined costs
-        params = self.update_params(params, rollouts)
-        return params, rollouts
+        rollouts_final = jax.tree.map(lambda x: x[-1], rollouts)
+
+        return params, rollouts_final
 
     def rollout_with_randomizations(
         self,
diff --git a/hydrax/algs/cem.py b/hydrax/algs/cem.py
index 0b44103a..cdb8f7eb 100644
--- a/hydrax/algs/cem.py
+++ b/hydrax/algs/cem.py
@@ -40,6 +40,7 @@ def __init__(
         plan_horizon: float = 1.0,
         spline_type: Literal["zero", "linear", "cubic"] = "zero",
         num_knots: int = 4,
+        iterations: int = 1,
     ) -> None:
         """Initialize the controller.
 
@@ -58,10 +59,12 @@ def __init__(
             spline_type: The type of spline used for control interpolation.
                          Defaults to "zero" (zero-order hold).
             num_knots: The number of knots in the control spline.
+            iterations: The number of optimization iterations to perform.
         """
         if not 0 <= explore_fraction <= 1:
             raise ValueError(
-                f"explore_fraction must be between 0 and 1, got {explore_fraction}"
+                f"explore_fraction must be between 0 and 1, got "
+                f"{explore_fraction}"
             )
         super().__init__(
             task,
@@ -71,6 +74,7 @@ def __init__(
             plan_horizon=plan_horizon,
             spline_type=spline_type,
             num_knots=num_knots,
+            iterations=iterations,
         )
         self.num_samples = num_samples
         self.sigma_min = sigma_min
diff --git a/hydrax/algs/evosax.py b/hydrax/algs/evosax.py
index 46b3f305..ea7521c4 100644
--- a/hydrax/algs/evosax.py
+++ b/hydrax/algs/evosax.py
@@ -47,6 +47,7 @@ def __init__(
         plan_horizon: float = 1.0,
         spline_type: Literal["zero", "linear", "cubic"] = "zero",
         num_knots: int = 4,
+        iterations: int = 1,
         **kwargs,
     ) -> None:
         """Initialize the controller.
@@ -64,6 +65,7 @@ def __init__(
             spline_type: The type of spline used for control interpolation.
                          Defaults to "zero" (zero-order hold).
             num_knots: The number of knots in the control spline.
+            iterations: The number of optimization iterations to perform.
             **kwargs: Additional keyword arguments for the optimizer.
         """
         super().__init__(
@@ -74,6 +76,7 @@ def __init__(
             plan_horizon=plan_horizon,
             spline_type=spline_type,
             num_knots=num_knots,
+            iterations=iterations,
         )
 
         self.strategy = optimizer(
diff --git a/hydrax/algs/mppi.py b/hydrax/algs/mppi.py
index bc5f4318..2a641f7f 100644
--- a/hydrax/algs/mppi.py
+++ b/hydrax/algs/mppi.py
@@ -43,6 +43,7 @@ def __init__(
         plan_horizon: float = 1.0,
         spline_type: Literal["zero", "linear", "cubic"] = "zero",
         num_knots: int = 4,
+        iterations: int = 1,
     ) -> None:
         """Initialize the controller.
 
@@ -60,6 +61,7 @@ def __init__(
             spline_type: The type of spline used for control interpolation.
                          Defaults to "zero" (zero-order hold).
             num_knots: The number of knots in the control spline.
+            iterations: The number of optimization iterations to perform.
         """
         super().__init__(
             task,
@@ -69,6 +71,7 @@ def __init__(
             plan_horizon=plan_horizon,
             spline_type=spline_type,
             num_knots=num_knots,
+            iterations=iterations,
         )
         self.noise_level = noise_level
         self.num_samples = num_samples
diff --git a/hydrax/algs/predictive_sampling.py b/hydrax/algs/predictive_sampling.py
index 2aa47b3f..1d307c8c 100644
--- a/hydrax/algs/predictive_sampling.py
+++ b/hydrax/algs/predictive_sampling.py
@@ -36,6 +36,7 @@ def __init__(
         plan_horizon: float = 1.0,
         spline_type: Literal["zero", "linear", "cubic"] = "zero",
         num_knots: int = 4,
+        iterations: int = 1,
     ) -> None:
         """Initialize the controller.
 
@@ -51,6 +52,7 @@ def __init__(
             spline_type: The type of spline used for control interpolation.
                          Defaults to "zero" (zero-order hold).
             num_knots: The number of knots in the control spline.
+            iterations: The number of optimization iterations to perform.
         """
         super().__init__(
             task,
@@ -60,6 +62,7 @@ def __init__(
             plan_horizon=plan_horizon,
             spline_type=spline_type,
             num_knots=num_knots,
+            iterations=iterations,
         )
         self.noise_level = noise_level
         self.num_samples = num_samples
diff --git a/tests/test_opt_iterations.py b/tests/test_opt_iterations.py
new file mode 100644
index 00000000..c48fdcea
--- /dev/null
+++ b/tests/test_opt_iterations.py
@@ -0,0 +1,83 @@
+import time
+
+import jax
+import jax.numpy as jnp
+import mujoco
+from mujoco import mjx
+
+from hydrax.algs import MPPI
+from hydrax.tasks.double_cart_pole import DoubleCartPole
+
+
+def test_opt_iterations() -> None:
+    """Test the optimization iterations feature."""
+    task = DoubleCartPole()
+    ctrl = MPPI(
+        task,
+        num_samples=1024,
+        num_randomizations=4,
+        plan_horizon=1.0,
+        noise_level=0.3,
+        temperature=0.1,
+        spline_type="zero",
+        num_knots=10,
+        iterations=1,
+    )
+
+    # Define the model used for simulation
+    mj_model = task.mj_model
+    mj_model.opt.timestep = 0.005
+    mj_model.opt.iterations = 50
+    mj_data = mujoco.MjData(mj_model)
+
+    mjx_data = mjx.put_data(mj_model, mj_data)
+    params = ctrl.init_params()
+
+    # Jit
+    jit_optimize = jax.jit(ctrl.optimize)
+
+    # Start timer
+    start_time = time.perf_counter()
+
+    params, _ = jit_optimize(mjx_data, params)
+    params_1_iter_, _ = jit_optimize(mjx_data, params)
+
+    # End timer
+    end_time = time.perf_counter()
+
+    elapsed_time = end_time - start_time
+    print(f"\nTwo manual iterations time: {elapsed_time:.4f} seconds")
+
+    ctrl = MPPI(
+        task,
+        num_samples=1024,
+        num_randomizations=4,
+        plan_horizon=1.0,
+        noise_level=0.3,
+        temperature=0.1,
+        spline_type="zero",
+        num_knots=10,
+        iterations=2,
+    )
+
+    mjx_data = mjx.put_data(mj_model, mj_data)
+    params = ctrl.init_params()
+
+    # Jit
+    jit_optimize = jax.jit(ctrl.optimize)
+
+    # Start timer
+    start_time = time.perf_counter()
+
+    params, _ = jit_optimize(mjx_data, params)
+
+    # End timer
+    end_time = time.perf_counter()
+
+    elapsed_time = end_time - start_time
+    print(f"\nTwo iterations time: {elapsed_time:.4f} seconds")
+
+    # Verify that the parameters are the same
+    assert jnp.all(params.tk == params_1_iter_.tk)
+    assert jnp.all(params.mean == params_1_iter_.mean)
+    assert jnp.all(params.rng == params_1_iter_.rng)


Reviews:
user: login: vincekurtz
body: Nice! This looks like it could be handy, especially for exploring something like offline trajopt with CMA-ES or some of the other evosax tools.

My only real concerns are:
1) Making sure there isn't a performance regression with `iterations = 1`
2) Adding a unit test. Just something simple. You could run `optimize` twice with `iterations = 1`, and then once with `iterations = 2` (and the same seed) and check that they match. 
state: CHANGES_REQUESTED
author_association: OWNER
submitted_at: 2025-04-18T21:26:20Z
user: login: Zolkin1
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-19T15:07:02Z
user: login: vincekurtz
state: COMMENTED
author_association: OWNER
submitted_at: 2025-04-23T01:14:41Z
user: login: Zolkin1
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-25T19:00:50Z
user: login: vincekurtz
body: Two little comments, otherwise looking good! The `scan` approach is definitely much better. I wasn't able to measure any extra overhead for the `iterations = 1` case on my machine. 
state: CHANGES_REQUESTED
author_association: OWNER
submitted_at: 2025-04-25T20:39:49Z
user: login: vincekurtz
state: APPROVED
author_association: OWNER
submitted_at: 2025-04-25T22:58:36Z

Review comments:
diff_hunk: @@ -6,3 +6,5 @@ __pycache__
 *.pkl
 *.egg-info
 MUJOCO_LOG.TXT
+
+.idea
path: .gitignore
user: login: vincekurtz
body: nit: whitespace

Also what is this file type?
created_at: 2025-04-18T21:14:58Z
author_association: OWNER
diff_hunk: @@ -35,6 +35,7 @@
         spline_type="cubic",
         plan_horizon=1.0,
         num_knots=4,
+        iterations=1,
path: examples/cart_pole.py
user: login: vincekurtz
body: nit: if it's not changing the default behavior, let's not specify parameters that have a reasonable default
created_at: 2025-04-18T21:15:42Z
author_association: OWNER
diff_hunk: @@ -20,6 +20,7 @@
     plan_horizon=0.5,
     spline_type="zero",
     num_knots=6,
+    iterations=3,
path: examples/pusht.py
user: login: vincekurtz
body: Should probably note (not sure where is the best place) that while boosting the number of iterations may make the deterministic sim look better, with async or on hardware the resulting delay could make it not worthwhile
created_at: 2025-04-18T21:17:58Z
author_association: OWNER
diff_hunk: @@ -135,22 +143,33 @@ def optimize(self, state: mjx.Data, params: Any) -> Tuple[Any, Trajectory]:
         new_mean = self.interp_func(new_tk, tk, params.mean[None, ...])[0]
         params = params.replace(tk=new_tk, mean=new_mean)
 
-        # Sample random control sequences from spline knots
-        knots, params = self.sample_knots(params)
-        knots = jnp.clip(
-            knots, self.task.u_min, self.task.u_max
-        )  # (num_rollouts, num_knots, nu)
-
-        # Roll out the control sequences, applying domain randomizations and
-        # combining costs using self.risk_strategy.
-        rng, dr_rng = jax.random.split(params.rng)
-        rollouts = self.rollout_with_randomizations(
-            state, new_tk, knots, dr_rng
+        def _optimize_loop_body(i: int, carry: Any):
path: hydrax/alg_base.py
user: login: vincekurtz
body: Have you run tests to confirm that this doesn't add overhead for the default case of iterations = 1?
created_at: 2025-04-18T21:20:39Z
author_association: OWNER
diff_hunk: @@ -135,22 +143,33 @@ def optimize(self, state: mjx.Data, params: Any) -> Tuple[Any, Trajectory]:
         new_mean = self.interp_func(new_tk, tk, params.mean[None, ...])[0]
         params = params.replace(tk=new_tk, mean=new_mean)
 
-        # Sample random control sequences from spline knots
-        knots, params = self.sample_knots(params)
-        knots = jnp.clip(
-            knots, self.task.u_min, self.task.u_max
-        )  # (num_rollouts, num_knots, nu)
-
-        # Roll out the control sequences, applying domain randomizations and
-        # combining costs using self.risk_strategy.
-        rng, dr_rng = jax.random.split(params.rng)
-        rollouts = self.rollout_with_randomizations(
-            state, new_tk, knots, dr_rng
+        def _optimize_loop_body(i: int, carry: Any):
+            params, _ = carry
+            # Sample random control sequences from spline knots
+            knots, params = self.sample_knots(params)
+            knots = jnp.clip(
+                knots, self.task.u_min, self.task.u_max
+            )  # (num_rollouts, num_knots, nu)
+
+            # Roll out the control sequences, applying domain randomizations and
+            # combining costs using self.risk_strategy.
+            rng, dr_rng = jax.random.split(params.rng)
+            rollouts = self.rollout_with_randomizations(
+                state, new_tk, knots, dr_rng
+            )
+            params = params.replace(rng=rng)
+
+            # Update the policy parameters based on the combined costs
+            params = self.update_params(params, rollouts)
+
+            return (params, rollouts)
+
+        params, rollouts = _optimize_loop_body(0, (params, None))
path: hydrax/alg_base.py
user: login: vincekurtz
body: Why do the first iteration separately, rather than a `fori_loop` starting from `0`?
created_at: 2025-04-18T21:21:44Z
author_association: OWNER
diff_hunk: @@ -6,3 +6,5 @@ __pycache__
 *.pkl
 *.egg-info
 MUJOCO_LOG.TXT
+
+.idea
path: .gitignore
user: login: Zolkin1
body: This is a settings folder that gets generated when working with Jetbrains IDEs.
created_at: 2025-04-19T15:07:01Z
author_association: CONTRIBUTOR
diff_hunk: @@ -135,22 +143,59 @@ def optimize(self, state: mjx.Data, params: Any) -> Tuple[Any, Trajectory]:
         new_mean = self.interp_func(new_tk, tk, params.mean[None, ...])[0]
         params = params.replace(tk=new_tk, mean=new_mean)
 
-        # Sample random control sequences from spline knots
+        def _optimize_loop_body(i: int, carry: Any):
+            params, _ = carry
+            # Sample random control sequences from spline knots
+            knots, params = self.sample_knots(params)
+            knots = jnp.clip(
+                knots, self.task.u_min, self.task.u_max
+            )  # (num_rollouts, num_knots, nu)
+
+            # Roll out the control sequences, applying domain randomizations and
+            # combining costs using self.risk_strategy.
+            rng, dr_rng = jax.random.split(params.rng)
+            rollouts = self.rollout_with_randomizations(
+                state, new_tk, knots, dr_rng
+            )
+            params = params.replace(rng=rng)
+
+            # Update the policy parameters based on the combined costs
+            params = self.update_params(params, rollouts)
+
+            return (params, rollouts)
+
         knots, params = self.sample_knots(params)
-        knots = jnp.clip(
-            knots, self.task.u_min, self.task.u_max
-        )  # (num_rollouts, num_knots, nu)
-
-        # Roll out the control sequences, applying domain randomizations and
-        # combining costs using self.risk_strategy.
-        rng, dr_rng = jax.random.split(params.rng)
-        rollouts = self.rollout_with_randomizations(
-            state, new_tk, knots, dr_rng
+
+        # params, rollouts = _optimize_loop_body(0, (params, None))
+
path: hydrax/alg_base.py
user: login: vincekurtz
body: I think we can avoid all of this pain with pre-allocating `rollouts` of the correct size if we use a `jax.lax.scan` rather than `fori_loop`. In that case, the `carry` in `_optimize_loop_body` should just be `params`. `rollouts` will be the second return argument, which gets stacked. That should be fine, we just need to be sure to return only the rollouts corresponding to the last iteration.

LMK if that's confusing/unclear: I can also try to write it out myself
created_at: 2025-04-23T01:11:20Z
author_association: OWNER
diff_hunk: @@ -0,0 +1,77 @@
+import time
+
+import jax
+import mujoco
+from mujoco import mjx
+
+from hydrax.algs import MPPI
+from hydrax.tasks.double_cart_pole import DoubleCartPole
+
+
+def test_opt_iterations() -> None:
+    """Test the optimization iterations feature."""
+    task = DoubleCartPole()
+    ctrl = MPPI(
+        task,
+        num_samples=1024,
+        num_randomizations=4,
+        plan_horizon=1.0,
+        noise_level=0.3,
+        temperature=0.1,
+        spline_type="zero",
+        num_knots=10,
+        iterations=1,
+    )
+
+    # Define the model used for simulation
+    mj_model = task.mj_model
+    mj_model.opt.timestep = 0.005
+    mj_model.opt.iterations = 50
+    mj_data = mujoco.MjData(mj_model)
+
+    mjx_data = mjx.put_data(mj_model, mj_data)
+    params = ctrl.init_params()
+
+    # Jit
+    jit_optimize = jax.jit(ctrl.optimize)
+
+    # Start timer
+    start_time = time.perf_counter()
+
+    jit_optimize(mjx_data, params)
+    jit_optimize(mjx_data, params)
path: tests/test_opt_iterations.py
user: login: vincekurtz
body: If we want to check that the results match, this should be `params, _ = jit_optimize(mjx_data, params)`, right? Otherwise nice little test!
created_at: 2025-04-23T01:14:09Z
author_association: OWNER
diff_hunk: @@ -135,22 +143,59 @@ def optimize(self, state: mjx.Data, params: Any) -> Tuple[Any, Trajectory]:
         new_mean = self.interp_func(new_tk, tk, params.mean[None, ...])[0]
         params = params.replace(tk=new_tk, mean=new_mean)
 
-        # Sample random control sequences from spline knots
+        def _optimize_loop_body(i: int, carry: Any):
+            params, _ = carry
+            # Sample random control sequences from spline knots
+            knots, params = self.sample_knots(params)
+            knots = jnp.clip(
+                knots, self.task.u_min, self.task.u_max
+            )  # (num_rollouts, num_knots, nu)
+
+            # Roll out the control sequences, applying domain randomizations and
+            # combining costs using self.risk_strategy.
+            rng, dr_rng = jax.random.split(params.rng)
+            rollouts = self.rollout_with_randomizations(
+                state, new_tk, knots, dr_rng
+            )
+            params = params.replace(rng=rng)
+
+            # Update the policy parameters based on the combined costs
+            params = self.update_params(params, rollouts)
+
+            return (params, rollouts)
+
         knots, params = self.sample_knots(params)
-        knots = jnp.clip(
-            knots, self.task.u_min, self.task.u_max
-        )  # (num_rollouts, num_knots, nu)
-
-        # Roll out the control sequences, applying domain randomizations and
-        # combining costs using self.risk_strategy.
-        rng, dr_rng = jax.random.split(params.rng)
-        rollouts = self.rollout_with_randomizations(
-            state, new_tk, knots, dr_rng
+
+        # params, rollouts = _optimize_loop_body(0, (params, None))
+
path: hydrax/alg_base.py
user: login: Zolkin1
body: I like that idea - I've made the change.
created_at: 2025-04-25T19:00:50Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,77 @@
+import time
+
+import jax
+import mujoco
+from mujoco import mjx
+
+from hydrax.algs import MPPI
+from hydrax.tasks.double_cart_pole import DoubleCartPole
+
+
+def test_opt_iterations() -> None:
+    """Test the optimization iterations feature."""
+    task = DoubleCartPole()
+    ctrl = MPPI(
+        task,
+        num_samples=1024,
+        num_randomizations=4,
+        plan_horizon=1.0,
+        noise_level=0.3,
+        temperature=0.1,
+        spline_type="zero",
+        num_knots=10,
+        iterations=1,
+    )
+
+    # Define the model used for simulation
+    mj_model = task.mj_model
+    mj_model.opt.timestep = 0.005
+    mj_model.opt.iterations = 50
+    mj_data = mujoco.MjData(mj_model)
+
+    mjx_data = mjx.put_data(mj_model, mj_data)
+    params = ctrl.init_params()
+
+    # Jit
+    jit_optimize = jax.jit(ctrl.optimize)
+
+    # Start timer
+    start_time = time.perf_counter()
+
+    jit_optimize(mjx_data, params)
+    jit_optimize(mjx_data, params)
path: tests/test_opt_iterations.py
user: login: vincekurtz
body: Let's add the check that using 2 iterations gives the same result as running `jit_optimize` twice.
created_at: 2025-04-25T20:35:21Z
author_association: OWNER
diff_hunk: @@ -135,23 +143,39 @@ def optimize(self, state: mjx.Data, params: Any) -> Tuple[Any, Trajectory]:
         new_mean = self.interp_func(new_tk, tk, params.mean[None, ...])[0]
         params = params.replace(tk=new_tk, mean=new_mean)
 
-        # Sample random control sequences from spline knots
-        knots, params = self.sample_knots(params)
-        knots = jnp.clip(
-            knots, self.task.u_min, self.task.u_max
-        )  # (num_rollouts, num_knots, nu)
-
-        # Roll out the control sequences, applying domain randomizations and
-        # combining costs using self.risk_strategy.
-        rng, dr_rng = jax.random.split(params.rng)
-        rollouts = self.rollout_with_randomizations(
-            state, new_tk, knots, dr_rng
+        def _optimize_scan_body(params: Any, _: Any):
+            # Sample random control sequences from spline knots
+            knots, params = self.sample_knots(params)
+            knots = jnp.clip(
+                knots, self.task.u_min, self.task.u_max
+            )  # (num_rollouts, num_knots, nu)
+
+            # Roll out the control sequences, applying domain randomizations and
+            # combining costs using self.risk_strategy.
+            rng, dr_rng = jax.random.split(params.rng)
+            rollouts = self.rollout_with_randomizations(
+                state, new_tk, knots, dr_rng
+            )
+            params = params.replace(rng=rng)
+
+            # Update the policy parameters based on the combined costs
+            params = self.update_params(params, rollouts)
+
+            return params, rollouts
+
+        params, rollouts = jax.lax.scan(
+            f=_optimize_scan_body, init=params, xs=jnp.arange(self.iterations)
+        )
+
+        # Get the rollouts from the last iteration only
+        rollouts_final = Trajectory(
path: hydrax/alg_base.py
user: login: vincekurtz
body: Let's use `jax.tree.map` rather than constructing a new trajectory. That way we won't need to update this if the definition of `Trajectory` ever changes in the future
created_at: 2025-04-25T20:37:30Z
author_association: OWNER
diff_hunk: @@ -135,23 +143,39 @@ def optimize(self, state: mjx.Data, params: Any) -> Tuple[Any, Trajectory]:
         new_mean = self.interp_func(new_tk, tk, params.mean[None, ...])[0]
         params = params.replace(tk=new_tk, mean=new_mean)
 
-        # Sample random control sequences from spline knots
-        knots, params = self.sample_knots(params)
-        knots = jnp.clip(
-            knots, self.task.u_min, self.task.u_max
-        )  # (num_rollouts, num_knots, nu)
-
-        # Roll out the control sequences, applying domain randomizations and
-        # combining costs using self.risk_strategy.
-        rng, dr_rng = jax.random.split(params.rng)
-        rollouts = self.rollout_with_randomizations(
-            state, new_tk, knots, dr_rng
+        def _optimize_scan_body(params: Any, _: Any):
+            # Sample random control sequences from spline knots
+            knots, params = self.sample_knots(params)
+            knots = jnp.clip(
+                knots, self.task.u_min, self.task.u_max
+            )  # (num_rollouts, num_knots, nu)
+
+            # Roll out the control sequences, applying domain randomizations and
+            # combining costs using self.risk_strategy.
+            rng, dr_rng = jax.random.split(params.rng)
+            rollouts = self.rollout_with_randomizations(
+                state, new_tk, knots, dr_rng
+            )
+            params = params.replace(rng=rng)
+
+            # Update the policy parameters based on the combined costs
+            params = self.update_params(params, rollouts)
+
+            return params, rollouts
+
+        params, rollouts = jax.lax.scan(
+            f=_optimize_scan_body, init=params, xs=jnp.arange(self.iterations)
         )
-        params = params.replace(rng=rng)
 
-        # Update the policy parameters based on the combined costs
-        params = self.update_params(params, rollouts)
-        return params, rollouts
+        def _get_last_element(leaf: Any):
+            if isinstance(leaf, jax.Array):
+                return leaf[-1]
+            else:
+                return leaf
+
+        rollouts_final = jax.tree.map(_get_last_element, rollouts)
path: hydrax/alg_base.py
user: login: vincekurtz
body: The idiomatic way to do this is in one line with a lambda function: `rollouts_final = jax.tree.map(lambda x: x[-1], rollouts)`. Otherwise LGTM!
created_at: 2025-04-25T22:56:21Z
author_association: OWNER

PR comments:
user: login: Tom0Brien
created_at: 2025-04-18T22:29:57Z
body: Nice! iirc I believe this will serve as a step towards potential implementation of a dial-mpc style algorithm (https://arxiv.org/abs/2409.15610)
author_association: CONTRIBUTOR
user: login: Zolkin1
created_at: 2025-04-21T17:16:45Z
body: > Nice! iirc I believe this will serve as a step towards potential implementation of a dial-mpc style algorithm (https://arxiv.org/abs/2409.15610)

Yea, that was one of the reasons I thought about adding this.

@vincekurtz it looks like the `num_samples` parameter is in every control algorithm, but not in the base class. I wasn't sure if this was done on purpose.  I now initialize the rollouts parameter to zero manually, but in order to get the correct dimensions I need to now the num samples amount which I currently get through a call to `sample_knots`. Not sure if there is a better way to get this dimension and/or if we can just move `num_samples` to the base class.

Unit test shows that there is a slight difference in timing, but that is to be expected as there is a bit more overhead when calling the function twice.


author_association: CONTRIBUTOR
user: login: Zolkin1
created_at: 2025-04-25T22:43:43Z
body: Ok I went through and made those changes - thanks for the suggestions with Jax! I'm definitely still learning all the idiomatic ways to do things in Jax. Lmk if there are any other desired changes.

Btw, I considered adding an __eq__ method to the `SamplingParams` class, but it looks like that would need to be implemented for each subclass because they automatically override the superclass in this case, so I just compared the elements manually. 
author_association: CONTRIBUTOR