From 62e3095074e6f7c021040bb1fe3fe59f432205b2 Mon Sep 17 00:00:00 2001
From: Duy Ta <dta@rai-inst.com>
Date: Fri, 20 Mar 2026 16:02:12 -0400
Subject: [PATCH 01/19] Fix for sumo

---
 judo/app/dora/controller.py    | 17 +++++++++++++----
 judo/app/dora/simulation.py    | 28 ++++++++++++++++++++++------
 judo/app/dora/visualization.py |  3 +++
 judo/simulation/__init__.py    | 19 +++++++++++++------
 judo/tasks/base.py             |  1 +
 judo/visualizers/visualizer.py |  5 +++--
 6 files changed, 55 insertions(+), 18 deletions(-)

diff --git a/judo/app/dora/controller.py b/judo/app/dora/controller.py
index 63ec9b7c..9ed17964 100644
--- a/judo/app/dora/controller.py
+++ b/judo/app/dora/controller.py
@@ -23,10 +23,14 @@ def __init__(
         max_workers: int | None = None,
         task_registration_cfg: DictConfig | None = None,
         optimizer_registration_cfg: DictConfig | None = None,
+        controller_cls: type[Controller] | None = None,
+        make_controller_fn=None,
     ) -> None:
         """Initialize the controller node."""
         super().__init__(node_id=node_id, max_workers=max_workers)
-        self.controller = make_controller(
+        self._controller_cls = controller_cls or Controller
+        self._make_controller_fn = make_controller_fn or make_controller
+        self.controller = self._make_controller_fn(
             init_task=init_task,
             init_optimizer=init_optimizer,
             task_registration_cfg=task_registration_cfg,
@@ -45,9 +49,13 @@ def update_task(self, event: dict) -> None:
             task_cls, _ = task_entry
             with self.lock:
                 task = task_cls()
-                optimizer = self.controller.optimizer_cls(self.controller.optimizer_config_cls(), task.nu)
-                self.controller = Controller(
-                    controller_config=self.controller.controller_cfg,
+                controller_config = type(self.controller.controller_cfg)()
+                controller_config.set_override(new_task)
+                optimizer_config = self.controller.optimizer_config_cls()
+                optimizer_config.set_override(new_task)
+                optimizer = self.controller.optimizer_cls(optimizer_config, task.nu)
+                self.controller = self._controller_cls(
+                    controller_config=controller_config,
                     task=task,
                     optimizer=optimizer,
                 )
@@ -75,6 +83,7 @@ def update_optimizer(self, event: dict) -> None:
         if optimizer_entry is not None:
             optimizer_cls, optimizer_config_cls = optimizer_entry
             optimizer_config = optimizer_config_cls()
+            optimizer_config.set_override(self.controller.task.name)
             optimizer = optimizer_cls(optimizer_config, self.controller.task.nu)
             with self.lock:
                 self.controller.optimizer = optimizer
diff --git a/judo/app/dora/simulation.py b/judo/app/dora/simulation.py
index 054fc595..cbe1b0b6 100644
--- a/judo/app/dora/simulation.py
+++ b/judo/app/dora/simulation.py
@@ -9,7 +9,8 @@
 from omegaconf import DictConfig
 
 from judo.app.structs import SplineData
-from judo.simulation import get_simulation_backend
+from judo.simulation.base import Simulation
+from judo.tasks import get_registered_tasks
 
 
 class SimulationNode(DoraNode):
@@ -22,25 +23,40 @@ def __init__(
         max_workers: int | None = None,
         task_registration_cfg: DictConfig | None = None,
         simulation_backend: str = "mujoco",
+        custom_backends: dict[str, type[Simulation]] | None = None,
     ) -> None:
         """Initialize the simulation node."""
         super().__init__(node_id=node_id, max_workers=max_workers)
         self._simulation_backend = simulation_backend
         self._task_registration_cfg = task_registration_cfg
+        self._custom_backends = custom_backends or {}
         self._init_sim(init_task)
         self.control_spline: Callable | None = None
         self.write_states()
 
+    def _resolve_backend(self, backend: str) -> type:
+        """Resolve a simulation backend class by name, checking custom backends first."""
+        if backend in self._custom_backends:
+            return self._custom_backends[backend]
+        from judo.simulation import get_simulation_backend
+        return get_simulation_backend(backend)
+
     def _init_sim(self, task_name: str) -> None:
         """Initialize simulation, auto-upgrading to policy backend if needed."""
-        backend = self._simulation_backend
-        _sim_backend = get_simulation_backend(backend)
-        self.sim = _sim_backend(init_task=task_name, task_registration_cfg=self._task_registration_cfg)
+        task_entry = get_registered_tasks().get(task_name)
+        if task_entry is None:
+            raise ValueError(f"Task {task_name} not found in task registry.")
+
+        task_cls, _ = task_entry
+        backend = getattr(task_cls, "default_backend", None) or self._simulation_backend
+
+        sim_backend_cls = self._resolve_backend(backend)
+        self.sim = sim_backend_cls(init_task=task_name, task_registration_cfg=self._task_registration_cfg)
 
         # Auto-upgrade to policy backend if task requires locomotion policy
         if backend == "mujoco" and self.sim.task.uses_locomotion_policy:
-            _sim_backend = get_simulation_backend("mujoco_policy")
-            self.sim = _sim_backend(init_task=task_name, task_registration_cfg=self._task_registration_cfg)
+            sim_backend_cls = self._resolve_backend("mujoco_policy")
+            self.sim = sim_backend_cls(init_task=task_name, task_registration_cfg=self._task_registration_cfg)
 
     @on_event("INPUT", "task")
     def update_task(self, event: dict) -> None:
diff --git a/judo/app/dora/visualization.py b/judo/app/dora/visualization.py
index 96381f9e..1eb048ee 100644
--- a/judo/app/dora/visualization.py
+++ b/judo/app/dora/visualization.py
@@ -9,6 +9,7 @@
 from viser import GuiFolderHandle, GuiImageHandle, GuiInputHandle, IcosphereHandle, MeshHandle
 
 from judo.app.structs import MujocoState
+from judo.tasks import Task, TaskConfig
 from judo.visualizers.visualizer import Visualizer
 
 ElementType = GuiImageHandle | GuiInputHandle | GuiFolderHandle | MeshHandle | IcosphereHandle
@@ -29,6 +30,7 @@ def __init__(
         optimizer_override_cfg: DictConfig | None = None,
         sim_pause_button: bool = True,
         geom_exclude_substring: str = "collision",
+        available_tasks: dict[str, tuple[type[Task], type[TaskConfig]]] | None = None,
     ) -> None:
         """Initialize the visualization node."""
         super().__init__(node_id=node_id, max_workers=max_workers)
@@ -41,6 +43,7 @@ def __init__(
             optimizer_override_cfg=optimizer_override_cfg,
             sim_pause_button=sim_pause_button,
             geom_exclude_substring=geom_exclude_substring,
+            available_tasks=available_tasks,
         )
 
     def write_sim_pause(self) -> None:
diff --git a/judo/simulation/__init__.py b/judo/simulation/__init__.py
index f58c801d..2b4edaff 100644
--- a/judo/simulation/__init__.py
+++ b/judo/simulation/__init__.py
@@ -2,21 +2,28 @@
 
 from judo.simulation.base import Simulation
 from judo.simulation.mj_simulation import MJSimulation
-from judo.simulation.policy_mj_simulation import PolicyMJSimulation
 
-simulation_registry = {
-    "mujoco": MJSimulation,
-    "mujoco_policy": PolicyMJSimulation,
+
+def _get_policy_mj_simulation():
+    from judo.simulation.policy_mj_simulation import PolicyMJSimulation
+    return PolicyMJSimulation
+
+
+_simulation_registry = {
+    "mujoco": lambda: MJSimulation,
+    "mujoco_policy": _get_policy_mj_simulation,
 }
 
 
 def get_simulation_backend(simulation_backend: str) -> type:
     """Get the simulation class for a given backend."""
-    return simulation_registry[simulation_backend]
+    if simulation_backend not in _simulation_registry:
+        raise KeyError(f"Unknown simulation backend: {simulation_backend!r}")
+    return _simulation_registry[simulation_backend]()
 
 
 __all__ = [
     "Simulation",
     "MJSimulation",
-    "PolicyMJSimulation",
+    "get_simulation_backend",
 ]
diff --git a/judo/tasks/base.py b/judo/tasks/base.py
index 02d4ff4d..a2676ecc 100644
--- a/judo/tasks/base.py
+++ b/judo/tasks/base.py
@@ -25,6 +25,7 @@ class Task(ABC, Generic[ConfigT]):
     """Task definition."""
 
     config_t: type[ConfigT]
+    name: str
 
     def __init__(self, model_path: Path | str = "", sim_model_path: Path | str | None = None) -> None:
         """Initialize the Mujoco task."""
diff --git a/judo/visualizers/visualizer.py b/judo/visualizers/visualizer.py
index 4562a9bf..69bbf102 100644
--- a/judo/visualizers/visualizer.py
+++ b/judo/visualizers/visualizer.py
@@ -15,7 +15,7 @@
 from judo.controller import ControllerConfig
 from judo.gui import create_gui_elements
 from judo.optimizers import get_registered_optimizers
-from judo.tasks import get_registered_tasks
+from judo.tasks import get_registered_tasks, Task, TaskConfig
 from judo.visualizers.model import ViserMjModel
 
 ElementType = GuiImageHandle | GuiInputHandle | GuiFolderHandle | MeshHandle | IcosphereHandle
@@ -39,6 +39,7 @@ def __init__(
         optimizer_override_cfg: DictConfig | None = None,
         sim_pause_button: bool = True,
         geom_exclude_substring: str = "collision",
+        available_tasks: dict[str, tuple[type[Task], type[TaskConfig]]] | None= None,
     ) -> None:
         """Initialize the visualization node."""
         # handling custom task and optimizer registration
@@ -49,7 +50,7 @@ def __init__(
 
         # starting the server
         self.server = viser.ViserServer()
-        self.available_tasks = get_registered_tasks()
+        self.available_tasks = available_tasks or get_registered_tasks()
         self.available_optimizers = get_registered_optimizers()
         self.geom_exclude_substring = geom_exclude_substring
 

From 89bff5c5b3b52ca2dfe724b9e7b7f5a607c16269 Mon Sep 17 00:00:00 2001
From: Duy Ta <dta@rai-inst.com>
Date: Fri, 20 Mar 2026 20:57:19 -0400
Subject: [PATCH 02/19] extend controller

---
 judo/app/dora/controller.py      | 31 ++++++++++++------------
 judo/controller/controller.py    | 25 +++++++++++++++----
 judo/simulation/base.py          | 17 ++++++++++++-
 judo/simulation/mj_simulation.py | 41 ++------------------------------
 4 files changed, 53 insertions(+), 61 deletions(-)

diff --git a/judo/app/dora/controller.py b/judo/app/dora/controller.py
index 9ed17964..d9c43c10 100644
--- a/judo/app/dora/controller.py
+++ b/judo/app/dora/controller.py
@@ -40,29 +40,28 @@ def __init__(
         self.write_controls()
         self.lock = Lock()
 
+    def _current_optimizer_name(self) -> str:
+        """Look up the name of the current optimizer from the registry."""
+        for name, (cls, _) in self.controller.available_optimizers.items():
+            if isinstance(self.controller.optimizer, cls):
+                return name
+        return "cem"
+
     @on_event("INPUT", "task")
     def update_task(self, event: dict) -> None:
         """Updates the task type."""
         new_task = event["value"].to_numpy(zero_copy_only=False)[0]
         task_entry = self.controller.available_tasks.get(new_task)
-        if task_entry is not None:
-            task_cls, _ = task_entry
-            with self.lock:
-                task = task_cls()
-                controller_config = type(self.controller.controller_cfg)()
-                controller_config.set_override(new_task)
-                optimizer_config = self.controller.optimizer_config_cls()
-                optimizer_config.set_override(new_task)
-                optimizer = self.controller.optimizer_cls(optimizer_config, task.nu)
-                self.controller = self._controller_cls(
-                    controller_config=controller_config,
-                    task=task,
-                    optimizer=optimizer,
-                )
-                self.write_controls()
-        else:
+        if task_entry is None:
             raise ValueError(f"Task {new_task} not found in task registry.")
 
+        with self.lock:
+            self.controller = self._make_controller_fn(
+                init_task=new_task,
+                init_optimizer=self._current_optimizer_name(),
+            )
+            self.write_controls()
+
     @on_event("INPUT", "task_reset")
     def reset_task(self, event: dict) -> None:
         """Resets the task."""
diff --git a/judo/controller/controller.py b/judo/controller/controller.py
index 2c07b1e8..5d283a76 100644
--- a/judo/controller/controller.py
+++ b/judo/controller/controller.py
@@ -51,6 +51,7 @@ def __init__(
         task: Task,
         optimizer: Optimizer,
         rollout_backend: Literal["mujoco"] = "mujoco",
+        custom_rollout_backends: dict[str, type] | None = None,
     ) -> None:
         """Initialize the controller.
 
@@ -59,20 +60,30 @@ def __init__(
             task: The task to use.
             optimizer: The optimizer to use.
             rollout_backend: The backend to use for rollouts. Currently only "mujoco" is supported.
+            custom_rollout_backends: Optional mapping of backend names to backend classes.
+                If the task's ``default_backend`` matches a key, that class is instantiated
+                with ``model`` and ``num_threads`` keyword arguments.
         """
         self._controller_cfg = controller_config
         self.task = task
         self.optimizer = optimizer
+        self._custom_rollout_backends = custom_rollout_backends or {}
 
         self.available_optimizers = get_registered_optimizers()
         self.available_tasks = get_registered_tasks()
 
         self.model = self.task.model
 
-        # Initialize rollout backend (auto-select policy backend if task requires it)
-        if self.task.uses_locomotion_policy:
+        # Initialize rollout backend
+        default_backend = getattr(self.task, "default_backend", None)
+        if default_backend and default_backend in self._custom_rollout_backends:
+            self.rollout_backend: RolloutBackend = self._custom_rollout_backends[default_backend](
+                model=self.model,
+                num_threads=self.optimizer_cfg.num_rollouts,
+            )
+        elif self.task.uses_locomotion_policy:
             assert self.task.locomotion_policy_path is not None
-            self.rollout_backend: RolloutBackend = PolicyMJRolloutBackend(
+            self.rollout_backend = PolicyMJRolloutBackend(
                 model=self.model,
                 num_threads=self.optimizer_cfg.num_rollouts,
                 policy_path=self.task.locomotion_policy_path,
@@ -376,7 +387,7 @@ def _init_action_normalizer(self) -> Normalizer:
             action_normalizer_kwargs["max"] = self.task.actuator_ctrlrange[:, 1]
         elif self.action_normalizer_type == "running":
             action_normalizer_kwargs["init_std"] = 1.0  # TODO(yunhai): make this configurable
-        return make_normalizer(self.action_normalizer_type, self.model.nu, **action_normalizer_kwargs)
+        return make_normalizer(self.action_normalizer_type, self.task.nu, **action_normalizer_kwargs)
 
 
 def make_spline(times: np.ndarray, controls: np.ndarray, spline_order: str) -> interp1d:
@@ -407,6 +418,8 @@ def make_controller(
     task_registration_cfg: DictConfig | None = None,
     optimizer_registration_cfg: DictConfig | None = None,
     rollout_backend: Literal["mujoco"] = "mujoco",
+    controller_cls: type[Controller] | None = None,
+    **controller_kwargs,
 ) -> Controller:
     """Make a controller."""
     available_optimizers = get_registered_optimizers()
@@ -434,9 +447,11 @@ def make_controller(
     controller_cfg = ControllerConfig()
     controller_cfg.set_override(init_task)
 
-    return Controller(
+    cls = controller_cls or Controller
+    return cls(
         controller_config=controller_cfg,
         task=task,
         optimizer=optimizer,
         rollout_backend=rollout_backend,
+        **controller_kwargs,
     )
diff --git a/judo/simulation/base.py b/judo/simulation/base.py
index a32aa20e..19340db6 100644
--- a/judo/simulation/base.py
+++ b/judo/simulation/base.py
@@ -5,6 +5,7 @@
 import numpy as np
 from omegaconf import DictConfig
 
+from judo.app.structs import MujocoState
 from judo.app.utils import register_tasks_from_cfg
 from judo.tasks import get_registered_tasks
 from judo.tasks.base import Task
@@ -54,6 +55,20 @@ def pause(self) -> None:
         self.paused = not self.paused
 
     @property
-    @abstractmethod
+    def sim_state(self) -> MujocoState:
+        """Returns the current simulation state."""
+        return MujocoState(
+            time=self.task.data.time,
+            qpos=self.task.data.qpos,
+            qvel=self.task.data.qvel,
+            xpos=self.task.data.xpos,
+            xquat=self.task.data.xquat,
+            mocap_pos=self.task.data.mocap_pos,
+            mocap_quat=self.task.data.mocap_quat,
+            sim_metadata=self.task.get_sim_metadata(),
+        )
+
+    @property
     def timestep(self) -> float:
         """Timestep the simulation expects to run at."""
+        return self.task.dt
diff --git a/judo/simulation/mj_simulation.py b/judo/simulation/mj_simulation.py
index 6abeb46d..db844f9d 100644
--- a/judo/simulation/mj_simulation.py
+++ b/judo/simulation/mj_simulation.py
@@ -6,7 +6,6 @@
 from mujoco import mj_step
 from omegaconf import DictConfig
 
-from judo.app.structs import MujocoState
 from judo.simulation.base import Simulation
 
 
@@ -22,20 +21,11 @@ def __init__(
         init_task: str = "spot_base",
         task_registration_cfg: DictConfig | None = None,
     ) -> None:
-        """Initialize the MuJoCo simulation.
-
-        Args:
-            init_task: Name of the task to initialize.
-            task_registration_cfg: Optional task registration configuration.
-        """
+        """Initialize the MuJoCo simulation."""
         super().__init__(init_task=init_task, task_registration_cfg=task_registration_cfg)
 
     def step(self, command: np.ndarray) -> None:
-        """Step the simulation forward.
-
-        Args:
-            command: Control array in task format (task.nu dimensions).
-        """
+        """Step the simulation forward."""
         if self.paused:
             return
 
@@ -44,30 +34,3 @@ def step(self, command: np.ndarray) -> None:
         self.task.pre_sim_step()
         mj_step(self.task.sim_model, self.task.data)
         self.task.post_sim_step()
-
-    def set_task(self, task_name: str) -> None:
-        """Set the current task.
-
-        Args:
-            task_name: Name of the task to set.
-        """
-        super().set_task(task_name)
-
-    @property
-    def sim_state(self) -> MujocoState:
-        """Returns the current simulation state."""
-        return MujocoState(
-            time=self.task.data.time,
-            qpos=self.task.data.qpos,
-            qvel=self.task.data.qvel,
-            xpos=self.task.data.xpos,
-            xquat=self.task.data.xquat,
-            mocap_pos=self.task.data.mocap_pos,
-            mocap_quat=self.task.data.mocap_quat,
-            sim_metadata=self.task.get_sim_metadata(),
-        )
-
-    @property
-    def timestep(self) -> float:
-        """Returns the effective simulation timestep (accounting for substeps)."""
-        return self.task.dt

From e2605e510c29e191053c7514018409b3944bcc97 Mon Sep 17 00:00:00 2001
From: Duy Ta <dta@rai-inst.com>
Date: Sun, 22 Mar 2026 02:09:41 -0400
Subject: [PATCH 03/19] refactor make_model_data_pairs

---
 judo/utils/mj_rollout_backend.py | 32 ++++++++++++++------------------
 1 file changed, 14 insertions(+), 18 deletions(-)

diff --git a/judo/utils/mj_rollout_backend.py b/judo/utils/mj_rollout_backend.py
index a4f42c49..e7b1ea8b 100644
--- a/judo/utils/mj_rollout_backend.py
+++ b/judo/utils/mj_rollout_backend.py
@@ -12,6 +12,13 @@
 from judo.utils.rollout_backend import RolloutBackend
 
 
+def make_model_data_pairs(model: MjModel, num_pairs: int) -> tuple[list[MjModel], list[MjData]]:
+    """Create model/data pairs for mujoco threaded rollout."""
+    models = [deepcopy(model) for _ in range(num_pairs)]
+    datas = [MjData(m) for m in models]
+    return models, datas
+
+
 class MJRolloutBackend(RolloutBackend):
     """Backend for conducting multithreaded rollouts using standard MuJoCo.
 
@@ -32,16 +39,9 @@ def __init__(
         self.num_threads = num_threads
         self.model = model
 
-        self._model_data_pairs = self._make_model_data_pairs(model, num_threads)
+        self._models, self._datas = make_model_data_pairs(model, num_threads)
         self._rollout_obj = Rollout(nthread=num_threads)
 
-    @staticmethod
-    def _make_model_data_pairs(model: MjModel, num_pairs: int) -> list[tuple[MjModel, MjData]]:
-        """Create model/data pairs for mujoco threaded rollout."""
-        models = [deepcopy(model) for _ in range(num_pairs)]
-        datas = [MjData(m) for m in models]
-        return list(zip(models, datas, strict=True))
-
     def rollout(
         self,
         x0: np.ndarray,
@@ -64,16 +64,12 @@ def rollout(
         if x0.ndim == 1:
             x0 = np.tile(x0, (self.num_threads, 1))
 
-        ms, ds = zip(*self._model_data_pairs, strict=True)
-        ms = list(ms)
-        ds = list(ds)
-
-        nq = ms[0].nq
-        nv = ms[0].nv
-        nu = ms[0].nu
+        nq = self._models[0].nq
+        nv = self._models[0].nv
+        nu = self._models[0].nu
 
         # Prepend time to batched x0
-        full_states = np.concatenate([time.time() * np.ones((len(ms), 1)), x0], axis=-1)
+        full_states = np.concatenate([time.time() * np.ones((len(self._models), 1)), x0], axis=-1)
 
         assert full_states.shape[-1] == nq + nv + 1
         assert full_states.ndim == 2
@@ -81,7 +77,7 @@ def rollout(
         assert controls.shape[-1] == nu
         assert controls.shape[0] == full_states.shape[0]
 
-        _states, _sensors = self._rollout_obj.rollout(ms, ds, full_states, controls)
+        _states, _sensors = self._rollout_obj.rollout(self._models, self._datas, full_states, controls)
 
         out_states = np.array(_states)[..., 1:]  # Remove time from state
         out_sensors = np.array(_sensors)
@@ -97,5 +93,5 @@ def update(self, num_threads: int) -> None:
         """
         self.num_threads = num_threads
         self._rollout_obj.close()
-        self._model_data_pairs = self._make_model_data_pairs(self.model, num_threads)
+        self._models, self._datas = make_model_data_pairs(self.model, num_threads)
         self._rollout_obj = Rollout(nthread=num_threads)

From 61b67174e50b4100fd4e083bc3ab7595a0a37624 Mon Sep 17 00:00:00 2001
From: Brandon Hung <bhung@theaiinstitute.com>
Date: Tue, 24 Mar 2026 11:51:07 -0400
Subject: [PATCH 04/19] Saving updated formatting

---
 judo/app/dora/controller.py    | 3 ++-
 judo/app/dora/simulation.py    | 1 +
 judo/controller/controller.py  | 4 ++--
 judo/simulation/__init__.py    | 3 ++-
 judo/visualizers/visualizer.py | 4 ++--
 5 files changed, 9 insertions(+), 6 deletions(-)

diff --git a/judo/app/dora/controller.py b/judo/app/dora/controller.py
index d9c43c10..cda4d40b 100644
--- a/judo/app/dora/controller.py
+++ b/judo/app/dora/controller.py
@@ -2,6 +2,7 @@
 
 import time
 from threading import Lock
+from typing import Callable
 
 import pyarrow as pa
 from dora_utils.dataclasses import from_event, to_arrow
@@ -24,7 +25,7 @@ def __init__(
         task_registration_cfg: DictConfig | None = None,
         optimizer_registration_cfg: DictConfig | None = None,
         controller_cls: type[Controller] | None = None,
-        make_controller_fn=None,
+        make_controller_fn: Callable | None = None,
     ) -> None:
         """Initialize the controller node."""
         super().__init__(node_id=node_id, max_workers=max_workers)
diff --git a/judo/app/dora/simulation.py b/judo/app/dora/simulation.py
index cbe1b0b6..ec651237 100644
--- a/judo/app/dora/simulation.py
+++ b/judo/app/dora/simulation.py
@@ -39,6 +39,7 @@ def _resolve_backend(self, backend: str) -> type:
         if backend in self._custom_backends:
             return self._custom_backends[backend]
         from judo.simulation import get_simulation_backend
+
         return get_simulation_backend(backend)
 
     def _init_sim(self, task_name: str) -> None:
diff --git a/judo/controller/controller.py b/judo/controller/controller.py
index 5d283a76..c57ff32f 100644
--- a/judo/controller/controller.py
+++ b/judo/controller/controller.py
@@ -2,7 +2,7 @@
 
 import warnings
 from dataclasses import dataclass
-from typing import Literal
+from typing import Any, Literal
 
 import numpy as np
 from omegaconf import DictConfig
@@ -419,7 +419,7 @@ def make_controller(
     optimizer_registration_cfg: DictConfig | None = None,
     rollout_backend: Literal["mujoco"] = "mujoco",
     controller_cls: type[Controller] | None = None,
-    **controller_kwargs,
+    **controller_kwargs: Any,
 ) -> Controller:
     """Make a controller."""
     available_optimizers = get_registered_optimizers()
diff --git a/judo/simulation/__init__.py b/judo/simulation/__init__.py
index 2b4edaff..f5866598 100644
--- a/judo/simulation/__init__.py
+++ b/judo/simulation/__init__.py
@@ -4,8 +4,9 @@
 from judo.simulation.mj_simulation import MJSimulation
 
 
-def _get_policy_mj_simulation():
+def _get_policy_mj_simulation() -> type[MJSimulation]:
     from judo.simulation.policy_mj_simulation import PolicyMJSimulation
+
     return PolicyMJSimulation
 
 
diff --git a/judo/visualizers/visualizer.py b/judo/visualizers/visualizer.py
index 69bbf102..b058d88a 100644
--- a/judo/visualizers/visualizer.py
+++ b/judo/visualizers/visualizer.py
@@ -15,7 +15,7 @@
 from judo.controller import ControllerConfig
 from judo.gui import create_gui_elements
 from judo.optimizers import get_registered_optimizers
-from judo.tasks import get_registered_tasks, Task, TaskConfig
+from judo.tasks import Task, TaskConfig, get_registered_tasks
 from judo.visualizers.model import ViserMjModel
 
 ElementType = GuiImageHandle | GuiInputHandle | GuiFolderHandle | MeshHandle | IcosphereHandle
@@ -39,7 +39,7 @@ def __init__(
         optimizer_override_cfg: DictConfig | None = None,
         sim_pause_button: bool = True,
         geom_exclude_substring: str = "collision",
-        available_tasks: dict[str, tuple[type[Task], type[TaskConfig]]] | None= None,
+        available_tasks: dict[str, tuple[type[Task], type[TaskConfig]]] | None = None,
     ) -> None:
         """Initialize the visualization node."""
         # handling custom task and optimizer registration

From 4d68bd4aa0d3863fa212c12b1270d6f4009333bc Mon Sep 17 00:00:00 2001
From: Brandon Hung <bhung@theaiinstitute.com>
Date: Tue, 24 Mar 2026 11:54:33 -0400
Subject: [PATCH 05/19] Fixing the import for the tests

---
 judo/simulation/__init__.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/judo/simulation/__init__.py b/judo/simulation/__init__.py
index f5866598..123f0ce6 100644
--- a/judo/simulation/__init__.py
+++ b/judo/simulation/__init__.py
@@ -2,6 +2,7 @@
 
 from judo.simulation.base import Simulation
 from judo.simulation.mj_simulation import MJSimulation
+from judo.simulation.policy_mj_simulation import PolicyMJSimulation
 
 
 def _get_policy_mj_simulation() -> type[MJSimulation]:
@@ -26,5 +27,6 @@ def get_simulation_backend(simulation_backend: str) -> type:
 __all__ = [
     "Simulation",
     "MJSimulation",
+    "PolicyMJSimulation",
     "get_simulation_backend",
 ]

From 642982e32183d2ae62512fa1d37e9338976b4bb8 Mon Sep 17 00:00:00 2001
From: Brandon Hung <bhung@theaiinstitute.com>
Date: Tue, 24 Mar 2026 14:49:59 -0400
Subject: [PATCH 06/19] storing changes back for the testing

---
 judo/simulation/base.py       | 14 +++++++-------
 judo/tasks/spot/spot_base.py  | 13 ++-----------
 tests/test_spot_tasks.py      |  2 +-
 tests/test_tasks/test_spot.py |  3 ++-
 4 files changed, 12 insertions(+), 20 deletions(-)

diff --git a/judo/simulation/base.py b/judo/simulation/base.py
index 19340db6..76b2dcd7 100644
--- a/judo/simulation/base.py
+++ b/judo/simulation/base.py
@@ -58,13 +58,13 @@ def pause(self) -> None:
     def sim_state(self) -> MujocoState:
         """Returns the current simulation state."""
         return MujocoState(
-            time=self.task.data.time,
-            qpos=self.task.data.qpos,
-            qvel=self.task.data.qvel,
-            xpos=self.task.data.xpos,
-            xquat=self.task.data.xquat,
-            mocap_pos=self.task.data.mocap_pos,
-            mocap_quat=self.task.data.mocap_quat,
+            time=self.task.data.time,  # type: ignore
+            qpos=self.task.data.qpos,  # type: ignore
+            qvel=self.task.data.qvel,  # type: ignore
+            xpos=self.task.data.xpos,  # type: ignore
+            xquat=self.task.data.xquat,  # type: ignore
+            mocap_pos=self.task.data.mocap_pos,  # type: ignore
+            mocap_quat=self.task.data.mocap_quat,  # type: ignore
             sim_metadata=self.task.get_sim_metadata(),
         )
 
diff --git a/judo/tasks/spot/spot_base.py b/judo/tasks/spot/spot_base.py
index 87e55630..32ab7925 100644
--- a/judo/tasks/spot/spot_base.py
+++ b/judo/tasks/spot/spot_base.py
@@ -6,7 +6,7 @@
 but adapted for judo's standalone simulation framework.
 """
 
-from dataclasses import dataclass, field
+from dataclasses import dataclass
 from pathlib import Path
 from typing import Any, Generic, TypeVar
 
@@ -43,15 +43,6 @@
 XML_PATH = str(MODEL_PATH / "xml" / "spot_primitive" / "robot.xml")
 
 
-@dataclass
-class GoalPositions:
-    """Goal positions for Spot tasks."""
-
-    origin: np.ndarray = field(default_factory=lambda: np.array([0, 0, 0.0]))
-    blue_cross: np.ndarray = field(default_factory=lambda: np.array([2.77, 0.71, 0.3]))
-    black_cross: np.ndarray = field(default_factory=lambda: np.array([1.5, -1.5, 0.275]))
-
-
 @dataclass
 class SpotBaseConfig(TaskConfig):
     """Base configuration for Spot tasks.
@@ -443,7 +434,7 @@ def reset(self) -> None:
     def get_action_components(self) -> list[str]:
         """Get names of each component in the action command vector.
 
-        Matches starfish/dexterity/tasks/spot_base.py.
+        Matches judo/tasks/spot/spot_base.py.
         """
         action_components = ["spot/base.vx", "spot/base.vy", "spot/base.vtheta"]
         if self.use_arm:
diff --git a/tests/test_spot_tasks.py b/tests/test_spot_tasks.py
index 5ee3784a..90744b90 100644
--- a/tests/test_spot_tasks.py
+++ b/tests/test_spot_tasks.py
@@ -8,12 +8,12 @@
 import pytest
 
 from judo.tasks.spot import (
-    SpotBase,
     SpotBoxPush,
     SpotNavigate,
     SpotTireRoll,
     SpotTireUpright,
 )
+from judo.tasks.spot.spot_base import SpotBase
 
 # (TaskClass, expected_nu)
 # nu depends on use_arm, use_gripper, use_legs, use_torso:
diff --git a/tests/test_tasks/test_spot.py b/tests/test_tasks/test_spot.py
index daccac71..33d96cac 100644
--- a/tests/test_tasks/test_spot.py
+++ b/tests/test_tasks/test_spot.py
@@ -4,7 +4,8 @@
 
 import numpy as np
 
-from judo.tasks.spot import SpotBase, SpotTireUpright
+from judo.tasks.spot import SpotTireUpright
+from judo.tasks.spot.spot_base import SpotBase
 
 
 def test_spot_base_init() -> None:

From 6693e4584213930226dd3c49e9c825e239349c7d Mon Sep 17 00:00:00 2001
From: Brandon Hung <bhung@theaiinstitute.com>
Date: Tue, 24 Mar 2026 14:56:10 -0400
Subject: [PATCH 07/19] Fixing pre-commits

---
 judo/app/dora/simulation.py | 2 +-
 judo/simulation/__init__.py | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/judo/app/dora/simulation.py b/judo/app/dora/simulation.py
index ec651237..d8c0f42c 100644
--- a/judo/app/dora/simulation.py
+++ b/judo/app/dora/simulation.py
@@ -38,7 +38,7 @@ def _resolve_backend(self, backend: str) -> type:
         """Resolve a simulation backend class by name, checking custom backends first."""
         if backend in self._custom_backends:
             return self._custom_backends[backend]
-        from judo.simulation import get_simulation_backend
+        from judo.simulation import get_simulation_backend  # noqa: PLC0415, I001
 
         return get_simulation_backend(backend)
 
diff --git a/judo/simulation/__init__.py b/judo/simulation/__init__.py
index 123f0ce6..3062cbb8 100644
--- a/judo/simulation/__init__.py
+++ b/judo/simulation/__init__.py
@@ -6,7 +6,7 @@
 
 
 def _get_policy_mj_simulation() -> type[MJSimulation]:
-    from judo.simulation.policy_mj_simulation import PolicyMJSimulation
+    from judo.simulation.policy_mj_simulation import PolicyMJSimulation  # noqa: PLC0415, I001
 
     return PolicyMJSimulation
 

From 8545b2aab8a39f509a041dd5cb30c549f190a0f9 Mon Sep 17 00:00:00 2001
From: Brandon Hung <bhung@theaiinstitute.com>
Date: Tue, 24 Mar 2026 16:07:24 -0400
Subject: [PATCH 08/19] addressed minor feedback

---
 judo/simulation/__init__.py      |  9 ++++++++-
 judo/simulation/mj_simulation.py | 13 +++++++++++--
 2 files changed, 19 insertions(+), 3 deletions(-)

diff --git a/judo/simulation/__init__.py b/judo/simulation/__init__.py
index 3062cbb8..fc522e94 100644
--- a/judo/simulation/__init__.py
+++ b/judo/simulation/__init__.py
@@ -18,7 +18,14 @@ def _get_policy_mj_simulation() -> type[MJSimulation]:
 
 
 def get_simulation_backend(simulation_backend: str) -> type:
-    """Get the simulation class for a given backend."""
+    """Get the simulation class for a given backend.
+
+    Args:
+        simulation_backend: Name of the simulation backend to get.
+
+    Returns:
+        The simulation class for the given backend.
+    """
     if simulation_backend not in _simulation_registry:
         raise KeyError(f"Unknown simulation backend: {simulation_backend!r}")
     return _simulation_registry[simulation_backend]()
diff --git a/judo/simulation/mj_simulation.py b/judo/simulation/mj_simulation.py
index db844f9d..4d7164e0 100644
--- a/judo/simulation/mj_simulation.py
+++ b/judo/simulation/mj_simulation.py
@@ -21,11 +21,20 @@ def __init__(
         init_task: str = "spot_base",
         task_registration_cfg: DictConfig | None = None,
     ) -> None:
-        """Initialize the MuJoCo simulation."""
+        """Initialize the MuJoCo simulation.
+
+        Args:
+            init_task: Name of the task to initialize.
+            task_registration_cfg: Optional task registration configuration.
+        """
         super().__init__(init_task=init_task, task_registration_cfg=task_registration_cfg)
 
     def step(self, command: np.ndarray) -> None:
-        """Step the simulation forward."""
+        """Step the simulation forward.
+
+        Args:
+            command: Control command for this timestep.
+        """
         if self.paused:
             return
 

From 7ebdd1b4fa3f13d7740202fd4e77514cacbbd7e7 Mon Sep 17 00:00:00 2001
From: Brandon Hung <bhung@theaiinstitute.com>
Date: Thu, 2 Apr 2026 13:09:01 -0400
Subject: [PATCH 09/19] Separation of visualization and simulation states

---
 judo/app/dora/simulation.py         |  2 ++
 judo/app/dora/visualization.py      | 11 ++++++-----
 judo/app/structs.py                 | 10 ++++++++--
 judo/configs/judo_dora_default.yaml |  5 +++--
 judo/simulation/base.py             | 12 +++++++++---
 5 files changed, 28 insertions(+), 12 deletions(-)

diff --git a/judo/app/dora/simulation.py b/judo/app/dora/simulation.py
index d8c0f42c..7d4849e0 100644
--- a/judo/app/dora/simulation.py
+++ b/judo/app/dora/simulation.py
@@ -102,6 +102,8 @@ def write_states(self) -> None:
         """Reads data from simulation and writes to output topic."""
         arr, metadata = to_arrow(self.sim.sim_state)
         self.node.send_output("states", arr, metadata)
+        arr, metadata = to_arrow(self.sim.world_state)
+        self.node.send_output("world_states", arr, metadata)
 
     @on_event("INPUT", "sim_pause")
     def set_paused_status(self, event: dict) -> None:
diff --git a/judo/app/dora/visualization.py b/judo/app/dora/visualization.py
index 1eb048ee..08d8a2c1 100644
--- a/judo/app/dora/visualization.py
+++ b/judo/app/dora/visualization.py
@@ -8,7 +8,7 @@
 from omegaconf import DictConfig
 from viser import GuiFolderHandle, GuiImageHandle, GuiInputHandle, IcosphereHandle, MeshHandle
 
-from judo.app.structs import MujocoState
+from judo.app.structs import WorldState
 from judo.tasks import Task, TaskConfig
 from judo.visualizers.visualizer import Visualizer
 
@@ -88,7 +88,7 @@ def write_task_config(self) -> None:
             self.node.send_output("task_config", *to_arrow(self.visualizer.task_config))
         self.visualizer.task_config_updated.clear()
 
-    @on_event("INPUT", "states")
+    @on_event("INPUT", "world_states")
     def update_states(self, event: dict) -> None:
         """Callback to update states on receiving a new state measurement."""
         if self.visualizer.controller_config.spline_order == "cubic" and self.visualizer.optimizer_config.num_nodes < 4:
@@ -99,11 +99,12 @@ def update_states(self, event: dict) -> None:
                     break
             self.visualizer.optimizer_config_updated.set()
 
-        state_msg = from_arrow(event["value"], event["metadata"], MujocoState)
+        # TODO: change the mujoco state
+        world_state_msg = from_arrow(event["value"], event["metadata"], WorldState)
         try:
             with self.visualizer.task_lock:
-                self.visualizer.data.xpos[:] = state_msg.xpos
-                self.visualizer.data.xquat[:] = state_msg.xquat
+                self.visualizer.data.xpos[:] = world_state_msg.xpos
+                self.visualizer.data.xquat[:] = world_state_msg.xquat
                 self.visualizer.viser_model.set_data(self.visualizer.data)
         except ValueError:
             # we're switching tasks and the new task has a different number of xpos/xquat
diff --git a/judo/app/structs.py b/judo/app/structs.py
index 1a65eb41..ac143caf 100644
--- a/judo/app/structs.py
+++ b/judo/app/structs.py
@@ -34,8 +34,6 @@ class MujocoState:
     time: float
     qpos: np.ndarray
     qvel: np.ndarray
-    xpos: np.ndarray
-    xquat: np.ndarray
     mocap_pos: np.ndarray
     mocap_quat: np.ndarray
     sim_metadata: dict[str, Any]
@@ -82,3 +80,11 @@ def spline(self) -> interp1d:
             fill_value=fill_value,  # type: ignore
             bounds_error=not self.extrapolate,
         )
+
+
+@dataclass
+class WorldState:
+    """Struct for world states used in visualization."""
+
+    xpos: np.ndarray
+    xquat: np.ndarray
diff --git a/judo/configs/judo_dora_default.yaml b/judo/configs/judo_dora_default.yaml
index f667539d..6487039d 100644
--- a/judo/configs/judo_dora_default.yaml
+++ b/judo/configs/judo_dora_default.yaml
@@ -20,11 +20,12 @@ dataflow:
           queue_size: 1
       outputs:
         - states
+        - world_states
     - id: visualization
       path: dynamic
       inputs:
-        states:
-          source: simulation/states
+        world_states:
+          source: simulation/world_states
           queue_size: 1
         traces:
           source: controller/traces
diff --git a/judo/simulation/base.py b/judo/simulation/base.py
index 76b2dcd7..9575fad9 100644
--- a/judo/simulation/base.py
+++ b/judo/simulation/base.py
@@ -5,7 +5,7 @@
 import numpy as np
 from omegaconf import DictConfig
 
-from judo.app.structs import MujocoState
+from judo.app.structs import MujocoState, WorldState
 from judo.app.utils import register_tasks_from_cfg
 from judo.tasks import get_registered_tasks
 from judo.tasks.base import Task
@@ -61,13 +61,19 @@ def sim_state(self) -> MujocoState:
             time=self.task.data.time,  # type: ignore
             qpos=self.task.data.qpos,  # type: ignore
             qvel=self.task.data.qvel,  # type: ignore
-            xpos=self.task.data.xpos,  # type: ignore
-            xquat=self.task.data.xquat,  # type: ignore
             mocap_pos=self.task.data.mocap_pos,  # type: ignore
             mocap_quat=self.task.data.mocap_quat,  # type: ignore
             sim_metadata=self.task.get_sim_metadata(),
         )
 
+    @property
+    def world_state(self) -> WorldState:
+        """Returns the current global state."""
+        return WorldState(
+            xpos=self.task.data.xpos,  # type: ignore
+            xquat=self.task.data.xquat,  # type: ignore
+        )
+
     @property
     def timestep(self) -> float:
         """Timestep the simulation expects to run at."""

From 97c38ab1540050fa410a54da4a696e55f99cb98f Mon Sep 17 00:00:00 2001
From: Brandon Hung <bhung@theaiinstitute.com>
Date: Fri, 10 Apr 2026 16:46:15 -0400
Subject: [PATCH 10/19] Moved mujoco version

---
 pixi.lock      | 312 ++++++++++++++++++++++++-------------------------
 pyproject.toml |   2 +-
 2 files changed, 157 insertions(+), 157 deletions(-)

diff --git a/pixi.lock b/pixi.lock
index 6d2aa972..fbc2092d 100644
--- a/pixi.lock
+++ b/pixi.lock
@@ -5,8 +5,6 @@ environments:
     - url: https://conda.anaconda.org/conda-forge/
     indexes:
     - https://pypi.org/simple
-    options:
-      pypi-prerelease-mode: if-necessary-or-explicit
     packages:
       linux-64:
       - conda: https://conda.anaconda.org/conda-forge/linux-64/_openmp_mutex-4.5-20_gnu.conda
@@ -93,7 +91,7 @@ environments:
       - pypi: https://files.pythonhosted.org/packages/94/54/e7d793b573f298e1c9013b8c4dade17d481164aa517d1d7148619c2cedbf/markdown_it_py-4.0.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/b3/38/89ba8ad64ae25be8de66a6d463314cf1eb366222074cfda9ee839c56a4b4/mdurl-0.1.2-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/28/83/36557b04cfdc317ed8a525c4993b23e43a8fbcddaddd78619112ca07138c/msgspec-0.20.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl
-      - pypi: https://files.pythonhosted.org/packages/16/47/93c7ac3a9630b49c55d76b0d02aa565543e2f62cecd885f8f574f5c745e7/mujoco-3.5.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl
+      - pypi: https://files.pythonhosted.org/packages/6e/48/c8cd52847d8a973fc606910a5467b8b7b68fa763afbe91f41d87123f957c/mujoco-3.6.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl
       - pypi: https://files.pythonhosted.org/packages/9e/c9/b2622292ea83fbb4ec318f5b9ab867d0a28ab43c5717bb85b0a5f6b3b0a4/networkx-3.6.1-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/88/b2/d0896bdcdc8d28a7fc5717c305f1a861c26e18c05047949fb371034d98bd/nodeenv-1.10.0-py2.py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/04/dc/46066ce18d01645541f0186877377b9371b8fa8017fa8262002b4ef22612/numpy-2.4.2-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl
@@ -195,7 +193,7 @@ environments:
       - pypi: https://files.pythonhosted.org/packages/94/54/e7d793b573f298e1c9013b8c4dade17d481164aa517d1d7148619c2cedbf/markdown_it_py-4.0.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/b3/38/89ba8ad64ae25be8de66a6d463314cf1eb366222074cfda9ee839c56a4b4/mdurl-0.1.2-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/8a/d1/b902d38b6e5ba3bdddbec469bba388d647f960aeed7b5b3623a8debe8a76/msgspec-0.20.0-cp313-cp313-macosx_10_13_x86_64.whl
-      - pypi: https://files.pythonhosted.org/packages/87/2a/371033684e4ddcda47c97661fb6e9617c0e5e3749af082a9b4d5d1bf9f27/mujoco-3.5.0-cp313-cp313-macosx_10_16_x86_64.whl
+      - pypi: https://files.pythonhosted.org/packages/29/0e/f3ea1fc9d1a25f19b173b13c23797805480cdb0a1026d43cf6b37dc2de6e/mujoco-3.6.0-cp313-cp313-macosx_10_16_x86_64.whl
       - pypi: https://files.pythonhosted.org/packages/9e/c9/b2622292ea83fbb4ec318f5b9ab867d0a28ab43c5717bb85b0a5f6b3b0a4/networkx-3.6.1-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/88/b2/d0896bdcdc8d28a7fc5717c305f1a861c26e18c05047949fb371034d98bd/nodeenv-1.10.0-py2.py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/a1/22/815b9fe25d1d7ae7d492152adbc7226d3eff731dffc38fe970589fcaaa38/numpy-2.4.2-cp313-cp313-macosx_10_13_x86_64.whl
@@ -297,7 +295,7 @@ environments:
       - pypi: https://files.pythonhosted.org/packages/94/54/e7d793b573f298e1c9013b8c4dade17d481164aa517d1d7148619c2cedbf/markdown_it_py-4.0.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/b3/38/89ba8ad64ae25be8de66a6d463314cf1eb366222074cfda9ee839c56a4b4/mdurl-0.1.2-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/57/b6/eff0305961a1d9447ec2b02f8c73c8946f22564d302a504185b730c9a761/msgspec-0.20.0-cp313-cp313-macosx_11_0_arm64.whl
-      - pypi: https://files.pythonhosted.org/packages/6f/c9/26bd4979d503d03f7a6ded851c3094a5708cb534cf0dc80b4db6672da2b0/mujoco-3.5.0-cp313-cp313-macosx_11_0_arm64.whl
+      - pypi: https://files.pythonhosted.org/packages/25/f2/2cbeabc6b69110e743100f550ec00ae8c60352b6975cf95470add299ed7a/mujoco-3.6.0-cp313-cp313-macosx_11_0_arm64.whl
       - pypi: https://files.pythonhosted.org/packages/9e/c9/b2622292ea83fbb4ec318f5b9ab867d0a28ab43c5717bb85b0a5f6b3b0a4/networkx-3.6.1-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/88/b2/d0896bdcdc8d28a7fc5717c305f1a861c26e18c05047949fb371034d98bd/nodeenv-1.10.0-py2.py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/09/f0/817d03a03f93ba9c6c8993de509277d84e69f9453601915e4a69554102a1/numpy-2.4.2-cp313-cp313-macosx_11_0_arm64.whl
@@ -395,7 +393,7 @@ environments:
       - pypi: https://files.pythonhosted.org/packages/94/54/e7d793b573f298e1c9013b8c4dade17d481164aa517d1d7148619c2cedbf/markdown_it_py-4.0.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/b3/38/89ba8ad64ae25be8de66a6d463314cf1eb366222074cfda9ee839c56a4b4/mdurl-0.1.2-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/f1/25/5e8080fe0117f799b1b68008dc29a65862077296b92550632de015128579/msgspec-0.20.0-cp313-cp313-win_amd64.whl
-      - pypi: https://files.pythonhosted.org/packages/ab/53/54a0815d43c83e1074cfc7da98a3dea88d7dda48c03edfd225a387a3767b/mujoco-3.5.0-cp313-cp313-win_amd64.whl
+      - pypi: https://files.pythonhosted.org/packages/c9/f4/17e3962681c141182616db9ec556ad902311ec154fffda7e9b35ed9677c2/mujoco-3.6.0-cp313-cp313-win_amd64.whl
       - pypi: https://files.pythonhosted.org/packages/9e/c9/b2622292ea83fbb4ec318f5b9ab867d0a28ab43c5717bb85b0a5f6b3b0a4/networkx-3.6.1-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/88/b2/d0896bdcdc8d28a7fc5717c305f1a861c26e18c05047949fb371034d98bd/nodeenv-1.10.0-py2.py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/43/bc/6352f343522fcb2c04dbaf94cb30cca6fd32c1a750c06ad6231b4293708c/numpy-2.4.2-cp313-cp313-win_amd64.whl
@@ -438,8 +436,6 @@ environments:
     - url: https://conda.anaconda.org/conda-forge/
     indexes:
     - https://pypi.org/simple
-    options:
-      pypi-prerelease-mode: if-necessary-or-explicit
     packages:
       linux-64:
       - conda: https://conda.anaconda.org/conda-forge/linux-64/_openmp_mutex-4.5-20_gnu.conda
@@ -556,7 +552,7 @@ environments:
       - pypi: https://files.pythonhosted.org/packages/09/ec/4b43dae793655b7d8a25f76119624350b4d65eb663459eb9603d7f1f0345/mistune-0.8.4-py2.py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/a4/8e/469e5a4a2f5855992e425f3cb33804cc07bf18d48f2db061aec61ce50270/more_itertools-10.8.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/28/83/36557b04cfdc317ed8a525c4993b23e43a8fbcddaddd78619112ca07138c/msgspec-0.20.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl
-      - pypi: https://files.pythonhosted.org/packages/16/47/93c7ac3a9630b49c55d76b0d02aa565543e2f62cecd885f8f574f5c745e7/mujoco-3.5.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl
+      - pypi: https://files.pythonhosted.org/packages/6e/48/c8cd52847d8a973fc606910a5467b8b7b68fa763afbe91f41d87123f957c/mujoco-3.6.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl
       - pypi: https://files.pythonhosted.org/packages/9e/c9/b2622292ea83fbb4ec318f5b9ab867d0a28ab43c5717bb85b0a5f6b3b0a4/networkx-3.6.1-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/6f/84/c0dc75c7fb596135f999e59a410d9f45bdabb989f1cb911f0016d22b747b/nh3-0.3.3-cp38-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
       - pypi: https://files.pythonhosted.org/packages/88/b2/d0896bdcdc8d28a7fc5717c305f1a861c26e18c05047949fb371034d98bd/nodeenv-1.10.0-py2.py3-none-any.whl
@@ -580,6 +576,7 @@ environments:
       - pypi: https://files.pythonhosted.org/packages/ee/49/1377b49de7d0c1ce41292161ea0f721913fa8722c19fb9c1e3aa0367eecb/pytest_cov-7.0.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/ca/31/d4e37e9e550c2b92a9cbc2e4d0b7420a27224968580b5a447f420847c975/pytest_xdist-3.8.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/ec/57/56b9bcc3c9c6a792fcbaf139543cee77261f3651ca9da0c93f5c1221264b/python_dateutil-2.9.0.post0-py2.py3-none-any.whl
+      - pypi: https://files.pythonhosted.org/packages/d8/db/795879cc3ddfe338599bddea6388cc5100b088db0a4caf6e6c1af1c27e04/python_discovery-1.2.2-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/74/27/e5b8f34d02d9995b80abcef563ea1f8b56d20134d8f4e5e81733b1feceb2/pyyaml-6.0.3-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl
       - pypi: https://files.pythonhosted.org/packages/e1/67/921ec3024056483db83953ae8e48079ad62b92db7880013ca77632921dd0/readme_renderer-44.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/2c/58/ca301544e1fa93ed4f80d724bf5b194f6e4b945841c5bfd555878eea9fcb/referencing-0.37.0-py3-none-any.whl
@@ -618,7 +615,7 @@ environments:
       - pypi: https://files.pythonhosted.org/packages/39/08/aaaad47bc4e9dc8c725e68f9d04865dbcb2052843ff09c97b08904852d84/urllib3-2.6.3-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/ce/de/0638d8aa83c2038c5f820151006b3bf372d109afbfdb2a8fc59187cc3690/uv-0.10.3-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
       - pypi: https://files.pythonhosted.org/packages/b8/b4/07ab1c828bae0eb5c72cd9a4cbe8b0376d374509be3c7055e1a399bf85c3/vhacdx-0.0.10-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl
-      - pypi: https://files.pythonhosted.org/packages/42/4b/6cf85b485be7ec29db837ec2a1d8cd68bc1147b1abf23d8636c5bd65b3cc/virtualenv-20.37.0-py3-none-any.whl
+      - pypi: https://files.pythonhosted.org/packages/20/0e/f083a76cb590e60dff3868779558eefefb8dfb7c9ed020babc7aa014ccbf/virtualenv-21.2.1-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/b9/c4/819f000b3d64f654a054a8b24bf8d0c5e799636393135784d0886d8c808e/viser-1.0.21-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/ff/b2/83a6ddf56cdcbad4e3d841fcc55d6ba7d19aeb89c50f24dd7e859ec0805f/websockets-15.0.1-cp313-cp313-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl
       - pypi: https://files.pythonhosted.org/packages/87/22/b76d483683216dde3d67cba61fb2444be8d5be289bf628c13fc0fd90e5f9/wheel-0.46.3-py3-none-any.whl
@@ -718,7 +715,7 @@ environments:
       - pypi: https://files.pythonhosted.org/packages/09/ec/4b43dae793655b7d8a25f76119624350b4d65eb663459eb9603d7f1f0345/mistune-0.8.4-py2.py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/a4/8e/469e5a4a2f5855992e425f3cb33804cc07bf18d48f2db061aec61ce50270/more_itertools-10.8.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/8a/d1/b902d38b6e5ba3bdddbec469bba388d647f960aeed7b5b3623a8debe8a76/msgspec-0.20.0-cp313-cp313-macosx_10_13_x86_64.whl
-      - pypi: https://files.pythonhosted.org/packages/87/2a/371033684e4ddcda47c97661fb6e9617c0e5e3749af082a9b4d5d1bf9f27/mujoco-3.5.0-cp313-cp313-macosx_10_16_x86_64.whl
+      - pypi: https://files.pythonhosted.org/packages/29/0e/f3ea1fc9d1a25f19b173b13c23797805480cdb0a1026d43cf6b37dc2de6e/mujoco-3.6.0-cp313-cp313-macosx_10_16_x86_64.whl
       - pypi: https://files.pythonhosted.org/packages/9e/c9/b2622292ea83fbb4ec318f5b9ab867d0a28ab43c5717bb85b0a5f6b3b0a4/networkx-3.6.1-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/13/3e/aef8cf8e0419b530c95e96ae93a5078e9b36c1e6613eeb1df03a80d5194e/nh3-0.3.3-cp38-abi3-macosx_10_12_x86_64.macosx_11_0_arm64.macosx_10_12_universal2.whl
       - pypi: https://files.pythonhosted.org/packages/88/b2/d0896bdcdc8d28a7fc5717c305f1a861c26e18c05047949fb371034d98bd/nodeenv-1.10.0-py2.py3-none-any.whl
@@ -742,6 +739,7 @@ environments:
       - pypi: https://files.pythonhosted.org/packages/ee/49/1377b49de7d0c1ce41292161ea0f721913fa8722c19fb9c1e3aa0367eecb/pytest_cov-7.0.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/ca/31/d4e37e9e550c2b92a9cbc2e4d0b7420a27224968580b5a447f420847c975/pytest_xdist-3.8.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/ec/57/56b9bcc3c9c6a792fcbaf139543cee77261f3651ca9da0c93f5c1221264b/python_dateutil-2.9.0.post0-py2.py3-none-any.whl
+      - pypi: https://files.pythonhosted.org/packages/d8/db/795879cc3ddfe338599bddea6388cc5100b088db0a4caf6e6c1af1c27e04/python_discovery-1.2.2-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/d1/11/0fd08f8192109f7169db964b5707a2f1e8b745d4e239b784a5a1dd80d1db/pyyaml-6.0.3-cp313-cp313-macosx_10_13_x86_64.whl
       - pypi: https://files.pythonhosted.org/packages/e1/67/921ec3024056483db83953ae8e48079ad62b92db7880013ca77632921dd0/readme_renderer-44.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/2c/58/ca301544e1fa93ed4f80d724bf5b194f6e4b945841c5bfd555878eea9fcb/referencing-0.37.0-py3-none-any.whl
@@ -779,7 +777,7 @@ environments:
       - pypi: https://files.pythonhosted.org/packages/39/08/aaaad47bc4e9dc8c725e68f9d04865dbcb2052843ff09c97b08904852d84/urllib3-2.6.3-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/40/db/c5729ec39ed557ac673ca993abd71e5b18d71a92f95f1845152e639000ff/uv-0.10.3-py3-none-macosx_10_12_x86_64.whl
       - pypi: https://files.pythonhosted.org/packages/0e/54/c2fc08d9324bbd92735caf9207cbabada3a8dd9d270d6e46ca69eb7f883d/vhacdx-0.0.10-cp313-cp313-macosx_10_13_x86_64.whl
-      - pypi: https://files.pythonhosted.org/packages/42/4b/6cf85b485be7ec29db837ec2a1d8cd68bc1147b1abf23d8636c5bd65b3cc/virtualenv-20.37.0-py3-none-any.whl
+      - pypi: https://files.pythonhosted.org/packages/20/0e/f083a76cb590e60dff3868779558eefefb8dfb7c9ed020babc7aa014ccbf/virtualenv-21.2.1-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/b9/c4/819f000b3d64f654a054a8b24bf8d0c5e799636393135784d0886d8c808e/viser-1.0.21-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/8a/05/aa116ec9943c718905997412c5989f7ed671bc0188ee2ba89520e8765d7b/websockets-15.0.1-cp313-cp313-macosx_10_13_x86_64.whl
       - pypi: https://files.pythonhosted.org/packages/87/22/b76d483683216dde3d67cba61fb2444be8d5be289bf628c13fc0fd90e5f9/wheel-0.46.3-py3-none-any.whl
@@ -879,7 +877,7 @@ environments:
       - pypi: https://files.pythonhosted.org/packages/09/ec/4b43dae793655b7d8a25f76119624350b4d65eb663459eb9603d7f1f0345/mistune-0.8.4-py2.py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/a4/8e/469e5a4a2f5855992e425f3cb33804cc07bf18d48f2db061aec61ce50270/more_itertools-10.8.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/57/b6/eff0305961a1d9447ec2b02f8c73c8946f22564d302a504185b730c9a761/msgspec-0.20.0-cp313-cp313-macosx_11_0_arm64.whl
-      - pypi: https://files.pythonhosted.org/packages/6f/c9/26bd4979d503d03f7a6ded851c3094a5708cb534cf0dc80b4db6672da2b0/mujoco-3.5.0-cp313-cp313-macosx_11_0_arm64.whl
+      - pypi: https://files.pythonhosted.org/packages/25/f2/2cbeabc6b69110e743100f550ec00ae8c60352b6975cf95470add299ed7a/mujoco-3.6.0-cp313-cp313-macosx_11_0_arm64.whl
       - pypi: https://files.pythonhosted.org/packages/9e/c9/b2622292ea83fbb4ec318f5b9ab867d0a28ab43c5717bb85b0a5f6b3b0a4/networkx-3.6.1-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/13/3e/aef8cf8e0419b530c95e96ae93a5078e9b36c1e6613eeb1df03a80d5194e/nh3-0.3.3-cp38-abi3-macosx_10_12_x86_64.macosx_11_0_arm64.macosx_10_12_universal2.whl
       - pypi: https://files.pythonhosted.org/packages/88/b2/d0896bdcdc8d28a7fc5717c305f1a861c26e18c05047949fb371034d98bd/nodeenv-1.10.0-py2.py3-none-any.whl
@@ -903,6 +901,7 @@ environments:
       - pypi: https://files.pythonhosted.org/packages/ee/49/1377b49de7d0c1ce41292161ea0f721913fa8722c19fb9c1e3aa0367eecb/pytest_cov-7.0.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/ca/31/d4e37e9e550c2b92a9cbc2e4d0b7420a27224968580b5a447f420847c975/pytest_xdist-3.8.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/ec/57/56b9bcc3c9c6a792fcbaf139543cee77261f3651ca9da0c93f5c1221264b/python_dateutil-2.9.0.post0-py2.py3-none-any.whl
+      - pypi: https://files.pythonhosted.org/packages/d8/db/795879cc3ddfe338599bddea6388cc5100b088db0a4caf6e6c1af1c27e04/python_discovery-1.2.2-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/b1/16/95309993f1d3748cd644e02e38b75d50cbc0d9561d21f390a76242ce073f/pyyaml-6.0.3-cp313-cp313-macosx_11_0_arm64.whl
       - pypi: https://files.pythonhosted.org/packages/e1/67/921ec3024056483db83953ae8e48079ad62b92db7880013ca77632921dd0/readme_renderer-44.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/2c/58/ca301544e1fa93ed4f80d724bf5b194f6e4b945841c5bfd555878eea9fcb/referencing-0.37.0-py3-none-any.whl
@@ -940,7 +939,7 @@ environments:
       - pypi: https://files.pythonhosted.org/packages/39/08/aaaad47bc4e9dc8c725e68f9d04865dbcb2052843ff09c97b08904852d84/urllib3-2.6.3-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/0b/8c/cce26f6da9f69a9f0df032d642a4cb5603249f145c368e41b0154c463ca9/uv-0.10.3-py3-none-macosx_11_0_arm64.whl
       - pypi: https://files.pythonhosted.org/packages/3b/9e/42adb642a12915acc9cb2bfab21710a6aabf045c26967ba0ff0e08a872d0/vhacdx-0.0.10-cp313-cp313-macosx_11_0_arm64.whl
-      - pypi: https://files.pythonhosted.org/packages/42/4b/6cf85b485be7ec29db837ec2a1d8cd68bc1147b1abf23d8636c5bd65b3cc/virtualenv-20.37.0-py3-none-any.whl
+      - pypi: https://files.pythonhosted.org/packages/20/0e/f083a76cb590e60dff3868779558eefefb8dfb7c9ed020babc7aa014ccbf/virtualenv-21.2.1-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/b9/c4/819f000b3d64f654a054a8b24bf8d0c5e799636393135784d0886d8c808e/viser-1.0.21-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/ff/0b/33cef55ff24f2d92924923c99926dcce78e7bd922d649467f0eda8368923/websockets-15.0.1-cp313-cp313-macosx_11_0_arm64.whl
       - pypi: https://files.pythonhosted.org/packages/87/22/b76d483683216dde3d67cba61fb2444be8d5be289bf628c13fc0fd90e5f9/wheel-0.46.3-py3-none-any.whl
@@ -1036,7 +1035,7 @@ environments:
       - pypi: https://files.pythonhosted.org/packages/09/ec/4b43dae793655b7d8a25f76119624350b4d65eb663459eb9603d7f1f0345/mistune-0.8.4-py2.py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/a4/8e/469e5a4a2f5855992e425f3cb33804cc07bf18d48f2db061aec61ce50270/more_itertools-10.8.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/f1/25/5e8080fe0117f799b1b68008dc29a65862077296b92550632de015128579/msgspec-0.20.0-cp313-cp313-win_amd64.whl
-      - pypi: https://files.pythonhosted.org/packages/ab/53/54a0815d43c83e1074cfc7da98a3dea88d7dda48c03edfd225a387a3767b/mujoco-3.5.0-cp313-cp313-win_amd64.whl
+      - pypi: https://files.pythonhosted.org/packages/c9/f4/17e3962681c141182616db9ec556ad902311ec154fffda7e9b35ed9677c2/mujoco-3.6.0-cp313-cp313-win_amd64.whl
       - pypi: https://files.pythonhosted.org/packages/9e/c9/b2622292ea83fbb4ec318f5b9ab867d0a28ab43c5717bb85b0a5f6b3b0a4/networkx-3.6.1-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/73/88/1ce287ef8649dc51365b5094bd3713b76454838140a32ab4f8349973883c/nh3-0.3.3-cp38-abi3-win_amd64.whl
       - pypi: https://files.pythonhosted.org/packages/88/b2/d0896bdcdc8d28a7fc5717c305f1a861c26e18c05047949fb371034d98bd/nodeenv-1.10.0-py2.py3-none-any.whl
@@ -1060,6 +1059,7 @@ environments:
       - pypi: https://files.pythonhosted.org/packages/ee/49/1377b49de7d0c1ce41292161ea0f721913fa8722c19fb9c1e3aa0367eecb/pytest_cov-7.0.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/ca/31/d4e37e9e550c2b92a9cbc2e4d0b7420a27224968580b5a447f420847c975/pytest_xdist-3.8.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/ec/57/56b9bcc3c9c6a792fcbaf139543cee77261f3651ca9da0c93f5c1221264b/python_dateutil-2.9.0.post0-py2.py3-none-any.whl
+      - pypi: https://files.pythonhosted.org/packages/d8/db/795879cc3ddfe338599bddea6388cc5100b088db0a4caf6e6c1af1c27e04/python_discovery-1.2.2-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/de/3d/8161f7711c017e01ac9f008dfddd9410dff3674334c233bde66e7ba65bbf/pywin32_ctypes-0.2.3-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/97/c9/39d5b874e8b28845e4ec2202b5da735d0199dbe5b8fb85f91398814a9a46/pyyaml-6.0.3-cp313-cp313-win_amd64.whl
       - pypi: https://files.pythonhosted.org/packages/e1/67/921ec3024056483db83953ae8e48079ad62b92db7880013ca77632921dd0/readme_renderer-44.0-py3-none-any.whl
@@ -1098,7 +1098,7 @@ environments:
       - pypi: https://files.pythonhosted.org/packages/39/08/aaaad47bc4e9dc8c725e68f9d04865dbcb2052843ff09c97b08904852d84/urllib3-2.6.3-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/d5/95/89b9caeb4007d950dfa5fa5923c2ea67de541de0e3582598b7b9c7eeac53/uv-0.10.3-py3-none-win_amd64.whl
       - pypi: https://files.pythonhosted.org/packages/32/25/f0e6806824f88d47ab8bc1c9bf6f11634fd7b382d635d0696825f3b5672f/vhacdx-0.0.10-cp313-cp313-win_amd64.whl
-      - pypi: https://files.pythonhosted.org/packages/42/4b/6cf85b485be7ec29db837ec2a1d8cd68bc1147b1abf23d8636c5bd65b3cc/virtualenv-20.37.0-py3-none-any.whl
+      - pypi: https://files.pythonhosted.org/packages/20/0e/f083a76cb590e60dff3868779558eefefb8dfb7c9ed020babc7aa014ccbf/virtualenv-21.2.1-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/b9/c4/819f000b3d64f654a054a8b24bf8d0c5e799636393135784d0886d8c808e/viser-1.0.21-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/1b/6c/c65773d6cab416a64d191d6ee8a8b1c68a09970ea6909d16965d26bfed1e/websockets-15.0.1-cp313-cp313-win_amd64.whl
       - pypi: https://files.pythonhosted.org/packages/87/22/b76d483683216dde3d67cba61fb2444be8d5be289bf628c13fc0fd90e5f9/wheel-0.46.3-py3-none-any.whl
@@ -1112,8 +1112,6 @@ environments:
     - url: https://conda.anaconda.org/conda-forge/
     indexes:
     - https://pypi.org/simple
-    options:
-      pypi-prerelease-mode: if-necessary-or-explicit
     packages:
       linux-64:
       - conda: https://conda.anaconda.org/conda-forge/linux-64/_openmp_mutex-4.5-20_gnu.conda
@@ -1212,7 +1210,7 @@ environments:
       - pypi: https://files.pythonhosted.org/packages/b3/38/89ba8ad64ae25be8de66a6d463314cf1eb366222074cfda9ee839c56a4b4/mdurl-0.1.2-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/09/ec/4b43dae793655b7d8a25f76119624350b4d65eb663459eb9603d7f1f0345/mistune-0.8.4-py2.py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/28/83/36557b04cfdc317ed8a525c4993b23e43a8fbcddaddd78619112ca07138c/msgspec-0.20.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl
-      - pypi: https://files.pythonhosted.org/packages/16/47/93c7ac3a9630b49c55d76b0d02aa565543e2f62cecd885f8f574f5c745e7/mujoco-3.5.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl
+      - pypi: https://files.pythonhosted.org/packages/6e/48/c8cd52847d8a973fc606910a5467b8b7b68fa763afbe91f41d87123f957c/mujoco-3.6.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl
       - pypi: https://files.pythonhosted.org/packages/9e/c9/b2622292ea83fbb4ec318f5b9ab867d0a28ab43c5717bb85b0a5f6b3b0a4/networkx-3.6.1-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/88/b2/d0896bdcdc8d28a7fc5717c305f1a861c26e18c05047949fb371034d98bd/nodeenv-1.10.0-py2.py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/04/dc/46066ce18d01645541f0186877377b9371b8fa8017fa8262002b4ef22612/numpy-2.4.2-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl
@@ -1340,7 +1338,7 @@ environments:
       - pypi: https://files.pythonhosted.org/packages/b3/38/89ba8ad64ae25be8de66a6d463314cf1eb366222074cfda9ee839c56a4b4/mdurl-0.1.2-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/09/ec/4b43dae793655b7d8a25f76119624350b4d65eb663459eb9603d7f1f0345/mistune-0.8.4-py2.py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/8a/d1/b902d38b6e5ba3bdddbec469bba388d647f960aeed7b5b3623a8debe8a76/msgspec-0.20.0-cp313-cp313-macosx_10_13_x86_64.whl
-      - pypi: https://files.pythonhosted.org/packages/87/2a/371033684e4ddcda47c97661fb6e9617c0e5e3749af082a9b4d5d1bf9f27/mujoco-3.5.0-cp313-cp313-macosx_10_16_x86_64.whl
+      - pypi: https://files.pythonhosted.org/packages/29/0e/f3ea1fc9d1a25f19b173b13c23797805480cdb0a1026d43cf6b37dc2de6e/mujoco-3.6.0-cp313-cp313-macosx_10_16_x86_64.whl
       - pypi: https://files.pythonhosted.org/packages/9e/c9/b2622292ea83fbb4ec318f5b9ab867d0a28ab43c5717bb85b0a5f6b3b0a4/networkx-3.6.1-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/88/b2/d0896bdcdc8d28a7fc5717c305f1a861c26e18c05047949fb371034d98bd/nodeenv-1.10.0-py2.py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/a1/22/815b9fe25d1d7ae7d492152adbc7226d3eff731dffc38fe970589fcaaa38/numpy-2.4.2-cp313-cp313-macosx_10_13_x86_64.whl
@@ -1468,7 +1466,7 @@ environments:
       - pypi: https://files.pythonhosted.org/packages/b3/38/89ba8ad64ae25be8de66a6d463314cf1eb366222074cfda9ee839c56a4b4/mdurl-0.1.2-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/09/ec/4b43dae793655b7d8a25f76119624350b4d65eb663459eb9603d7f1f0345/mistune-0.8.4-py2.py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/57/b6/eff0305961a1d9447ec2b02f8c73c8946f22564d302a504185b730c9a761/msgspec-0.20.0-cp313-cp313-macosx_11_0_arm64.whl
-      - pypi: https://files.pythonhosted.org/packages/6f/c9/26bd4979d503d03f7a6ded851c3094a5708cb534cf0dc80b4db6672da2b0/mujoco-3.5.0-cp313-cp313-macosx_11_0_arm64.whl
+      - pypi: https://files.pythonhosted.org/packages/25/f2/2cbeabc6b69110e743100f550ec00ae8c60352b6975cf95470add299ed7a/mujoco-3.6.0-cp313-cp313-macosx_11_0_arm64.whl
       - pypi: https://files.pythonhosted.org/packages/9e/c9/b2622292ea83fbb4ec318f5b9ab867d0a28ab43c5717bb85b0a5f6b3b0a4/networkx-3.6.1-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/88/b2/d0896bdcdc8d28a7fc5717c305f1a861c26e18c05047949fb371034d98bd/nodeenv-1.10.0-py2.py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/09/f0/817d03a03f93ba9c6c8993de509277d84e69f9453601915e4a69554102a1/numpy-2.4.2-cp313-cp313-macosx_11_0_arm64.whl
@@ -1592,7 +1590,7 @@ environments:
       - pypi: https://files.pythonhosted.org/packages/b3/38/89ba8ad64ae25be8de66a6d463314cf1eb366222074cfda9ee839c56a4b4/mdurl-0.1.2-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/09/ec/4b43dae793655b7d8a25f76119624350b4d65eb663459eb9603d7f1f0345/mistune-0.8.4-py2.py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/f1/25/5e8080fe0117f799b1b68008dc29a65862077296b92550632de015128579/msgspec-0.20.0-cp313-cp313-win_amd64.whl
-      - pypi: https://files.pythonhosted.org/packages/ab/53/54a0815d43c83e1074cfc7da98a3dea88d7dda48c03edfd225a387a3767b/mujoco-3.5.0-cp313-cp313-win_amd64.whl
+      - pypi: https://files.pythonhosted.org/packages/c9/f4/17e3962681c141182616db9ec556ad902311ec154fffda7e9b35ed9677c2/mujoco-3.6.0-cp313-cp313-win_amd64.whl
       - pypi: https://files.pythonhosted.org/packages/9e/c9/b2622292ea83fbb4ec318f5b9ab867d0a28ab43c5717bb85b0a5f6b3b0a4/networkx-3.6.1-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/88/b2/d0896bdcdc8d28a7fc5717c305f1a861c26e18c05047949fb371034d98bd/nodeenv-1.10.0-py2.py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/43/bc/6352f343522fcb2c04dbaf94cb30cca6fd32c1a750c06ad6231b4293708c/numpy-2.4.2-cp313-cp313-win_amd64.whl
@@ -1649,8 +1647,6 @@ environments:
     - url: https://conda.anaconda.org/conda-forge/
     indexes:
     - https://pypi.org/simple
-    options:
-      pypi-prerelease-mode: if-necessary-or-explicit
     packages:
       linux-64:
       - conda: https://conda.anaconda.org/conda-forge/linux-64/_openmp_mutex-4.5-20_gnu.conda
@@ -1770,7 +1766,7 @@ environments:
       - pypi: https://files.pythonhosted.org/packages/09/ec/4b43dae793655b7d8a25f76119624350b4d65eb663459eb9603d7f1f0345/mistune-0.8.4-py2.py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/a4/8e/469e5a4a2f5855992e425f3cb33804cc07bf18d48f2db061aec61ce50270/more_itertools-10.8.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/15/2a/48e41d9ef0a24b1c6e67cbd94a676799e0561bfbc163be1aaaff5ca853f5/msgspec-0.20.0-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl
-      - pypi: https://files.pythonhosted.org/packages/ca/d6/a5a7b615b257867b7c97db6b3ce07dec9351d5d9d5a5aca881cbb583d7a3/mujoco-3.5.0-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl
+      - pypi: https://files.pythonhosted.org/packages/94/27/043852c6d46b9f5b092430e334b2102753f2b63b262f4f466c75c0b9daac/mujoco-3.6.0-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl
       - pypi: https://files.pythonhosted.org/packages/b9/54/dd730b32ea14ea797530a4479b2ed46a6fb250f682a9cfb997e968bf0261/networkx-3.4.2-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/6f/84/c0dc75c7fb596135f999e59a410d9f45bdabb989f1cb911f0016d22b747b/nh3-0.3.3-cp38-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
       - pypi: https://files.pythonhosted.org/packages/88/b2/d0896bdcdc8d28a7fc5717c305f1a861c26e18c05047949fb371034d98bd/nodeenv-1.10.0-py2.py3-none-any.whl
@@ -1794,6 +1790,7 @@ environments:
       - pypi: https://files.pythonhosted.org/packages/ee/49/1377b49de7d0c1ce41292161ea0f721913fa8722c19fb9c1e3aa0367eecb/pytest_cov-7.0.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/ca/31/d4e37e9e550c2b92a9cbc2e4d0b7420a27224968580b5a447f420847c975/pytest_xdist-3.8.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/ec/57/56b9bcc3c9c6a792fcbaf139543cee77261f3651ca9da0c93f5c1221264b/python_dateutil-2.9.0.post0-py2.py3-none-any.whl
+      - pypi: https://files.pythonhosted.org/packages/d8/db/795879cc3ddfe338599bddea6388cc5100b088db0a4caf6e6c1af1c27e04/python_discovery-1.2.2-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/7a/1e/7acc4f0e74c4b3d9531e24739e0ab832a5edf40e64fbae1a9c01941cabd7/pyyaml-6.0.3-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl
       - pypi: https://files.pythonhosted.org/packages/e1/67/921ec3024056483db83953ae8e48079ad62b92db7880013ca77632921dd0/readme_renderer-44.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/2c/58/ca301544e1fa93ed4f80d724bf5b194f6e4b945841c5bfd555878eea9fcb/referencing-0.37.0-py3-none-any.whl
@@ -1832,7 +1829,7 @@ environments:
       - pypi: https://files.pythonhosted.org/packages/39/08/aaaad47bc4e9dc8c725e68f9d04865dbcb2052843ff09c97b08904852d84/urllib3-2.6.3-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/ce/de/0638d8aa83c2038c5f820151006b3bf372d109afbfdb2a8fc59187cc3690/uv-0.10.3-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
       - pypi: https://files.pythonhosted.org/packages/bf/84/97e2305f6bd4a4de3d40bb234c38282cbcf2fa30653ff5ae4f7df9d8f3ec/vhacdx-0.0.10-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl
-      - pypi: https://files.pythonhosted.org/packages/42/4b/6cf85b485be7ec29db837ec2a1d8cd68bc1147b1abf23d8636c5bd65b3cc/virtualenv-20.37.0-py3-none-any.whl
+      - pypi: https://files.pythonhosted.org/packages/20/0e/f083a76cb590e60dff3868779558eefefb8dfb7c9ed020babc7aa014ccbf/virtualenv-21.2.1-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/b9/c4/819f000b3d64f654a054a8b24bf8d0c5e799636393135784d0886d8c808e/viser-1.0.21-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/97/3a/5323a6bb94917af13bbb34009fac01e55c51dfde354f63692bf2533ffbc2/websockets-15.0.1-cp310-cp310-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl
       - pypi: https://files.pythonhosted.org/packages/87/22/b76d483683216dde3d67cba61fb2444be8d5be289bf628c13fc0fd90e5f9/wheel-0.46.3-py3-none-any.whl
@@ -1933,7 +1930,7 @@ environments:
       - pypi: https://files.pythonhosted.org/packages/09/ec/4b43dae793655b7d8a25f76119624350b4d65eb663459eb9603d7f1f0345/mistune-0.8.4-py2.py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/a4/8e/469e5a4a2f5855992e425f3cb33804cc07bf18d48f2db061aec61ce50270/more_itertools-10.8.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/e3/5e/151883ba2047cca9db8ed2f86186b054ad200bc231352df15b0c1dd75b1f/msgspec-0.20.0-cp310-cp310-macosx_10_9_x86_64.whl
-      - pypi: https://files.pythonhosted.org/packages/e8/20/9e0595e653543df3e4233bc3ad7e50b371b81dbe48d45ffbc867ed7c379d/mujoco-3.5.0-cp310-cp310-macosx_10_16_x86_64.whl
+      - pypi: https://files.pythonhosted.org/packages/7f/ed/afb86dfbe53ff4aec0eb265713255899db9947539cd3f235450d82b669f7/mujoco-3.6.0-cp310-cp310-macosx_10_16_x86_64.whl
       - pypi: https://files.pythonhosted.org/packages/b9/54/dd730b32ea14ea797530a4479b2ed46a6fb250f682a9cfb997e968bf0261/networkx-3.4.2-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/13/3e/aef8cf8e0419b530c95e96ae93a5078e9b36c1e6613eeb1df03a80d5194e/nh3-0.3.3-cp38-abi3-macosx_10_12_x86_64.macosx_11_0_arm64.macosx_10_12_universal2.whl
       - pypi: https://files.pythonhosted.org/packages/88/b2/d0896bdcdc8d28a7fc5717c305f1a861c26e18c05047949fb371034d98bd/nodeenv-1.10.0-py2.py3-none-any.whl
@@ -1957,6 +1954,7 @@ environments:
       - pypi: https://files.pythonhosted.org/packages/ee/49/1377b49de7d0c1ce41292161ea0f721913fa8722c19fb9c1e3aa0367eecb/pytest_cov-7.0.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/ca/31/d4e37e9e550c2b92a9cbc2e4d0b7420a27224968580b5a447f420847c975/pytest_xdist-3.8.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/ec/57/56b9bcc3c9c6a792fcbaf139543cee77261f3651ca9da0c93f5c1221264b/python_dateutil-2.9.0.post0-py2.py3-none-any.whl
+      - pypi: https://files.pythonhosted.org/packages/d8/db/795879cc3ddfe338599bddea6388cc5100b088db0a4caf6e6c1af1c27e04/python_discovery-1.2.2-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/f4/a0/39350dd17dd6d6c6507025c0e53aef67a9293a6d37d3511f23ea510d5800/pyyaml-6.0.3-cp310-cp310-macosx_10_13_x86_64.whl
       - pypi: https://files.pythonhosted.org/packages/e1/67/921ec3024056483db83953ae8e48079ad62b92db7880013ca77632921dd0/readme_renderer-44.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/2c/58/ca301544e1fa93ed4f80d724bf5b194f6e4b945841c5bfd555878eea9fcb/referencing-0.37.0-py3-none-any.whl
@@ -1994,7 +1992,7 @@ environments:
       - pypi: https://files.pythonhosted.org/packages/39/08/aaaad47bc4e9dc8c725e68f9d04865dbcb2052843ff09c97b08904852d84/urllib3-2.6.3-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/40/db/c5729ec39ed557ac673ca993abd71e5b18d71a92f95f1845152e639000ff/uv-0.10.3-py3-none-macosx_10_12_x86_64.whl
       - pypi: https://files.pythonhosted.org/packages/c4/f4/da308d86daaa9c636851357cbd928715d47963beecd525b3749d2d5c9537/vhacdx-0.0.10-cp310-cp310-macosx_10_9_x86_64.whl
-      - pypi: https://files.pythonhosted.org/packages/42/4b/6cf85b485be7ec29db837ec2a1d8cd68bc1147b1abf23d8636c5bd65b3cc/virtualenv-20.37.0-py3-none-any.whl
+      - pypi: https://files.pythonhosted.org/packages/20/0e/f083a76cb590e60dff3868779558eefefb8dfb7c9ed020babc7aa014ccbf/virtualenv-21.2.1-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/b9/c4/819f000b3d64f654a054a8b24bf8d0c5e799636393135784d0886d8c808e/viser-1.0.21-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/1c/9f/9d11c1a4eb046a9e106483b9ff69bce7ac880443f00e5ce64261b47b07e7/websockets-15.0.1-cp310-cp310-macosx_10_9_x86_64.whl
       - pypi: https://files.pythonhosted.org/packages/87/22/b76d483683216dde3d67cba61fb2444be8d5be289bf628c13fc0fd90e5f9/wheel-0.46.3-py3-none-any.whl
@@ -2095,7 +2093,7 @@ environments:
       - pypi: https://files.pythonhosted.org/packages/09/ec/4b43dae793655b7d8a25f76119624350b4d65eb663459eb9603d7f1f0345/mistune-0.8.4-py2.py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/a4/8e/469e5a4a2f5855992e425f3cb33804cc07bf18d48f2db061aec61ce50270/more_itertools-10.8.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/50/88/a795647672f547c983eff0823b82aaa35db922c767e1b3693e2dcf96678d/msgspec-0.20.0-cp310-cp310-macosx_11_0_arm64.whl
-      - pypi: https://files.pythonhosted.org/packages/8d/6b/fdac8ed97086e12ac930fb44e419eda1626e339010df73678cb1f22527d7/mujoco-3.5.0-cp310-cp310-macosx_11_0_arm64.whl
+      - pypi: https://files.pythonhosted.org/packages/bc/6a/b1b3446f3b99bdbc56b4b63dadf4c5cd47ae08e3bf110f36dc5d3fde7836/mujoco-3.6.0-cp310-cp310-macosx_11_0_arm64.whl
       - pypi: https://files.pythonhosted.org/packages/b9/54/dd730b32ea14ea797530a4479b2ed46a6fb250f682a9cfb997e968bf0261/networkx-3.4.2-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/13/3e/aef8cf8e0419b530c95e96ae93a5078e9b36c1e6613eeb1df03a80d5194e/nh3-0.3.3-cp38-abi3-macosx_10_12_x86_64.macosx_11_0_arm64.macosx_10_12_universal2.whl
       - pypi: https://files.pythonhosted.org/packages/88/b2/d0896bdcdc8d28a7fc5717c305f1a861c26e18c05047949fb371034d98bd/nodeenv-1.10.0-py2.py3-none-any.whl
@@ -2119,6 +2117,7 @@ environments:
       - pypi: https://files.pythonhosted.org/packages/ee/49/1377b49de7d0c1ce41292161ea0f721913fa8722c19fb9c1e3aa0367eecb/pytest_cov-7.0.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/ca/31/d4e37e9e550c2b92a9cbc2e4d0b7420a27224968580b5a447f420847c975/pytest_xdist-3.8.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/ec/57/56b9bcc3c9c6a792fcbaf139543cee77261f3651ca9da0c93f5c1221264b/python_dateutil-2.9.0.post0-py2.py3-none-any.whl
+      - pypi: https://files.pythonhosted.org/packages/d8/db/795879cc3ddfe338599bddea6388cc5100b088db0a4caf6e6c1af1c27e04/python_discovery-1.2.2-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/05/14/52d505b5c59ce73244f59c7a50ecf47093ce4765f116cdb98286a71eeca2/pyyaml-6.0.3-cp310-cp310-macosx_11_0_arm64.whl
       - pypi: https://files.pythonhosted.org/packages/e1/67/921ec3024056483db83953ae8e48079ad62b92db7880013ca77632921dd0/readme_renderer-44.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/2c/58/ca301544e1fa93ed4f80d724bf5b194f6e4b945841c5bfd555878eea9fcb/referencing-0.37.0-py3-none-any.whl
@@ -2156,7 +2155,7 @@ environments:
       - pypi: https://files.pythonhosted.org/packages/39/08/aaaad47bc4e9dc8c725e68f9d04865dbcb2052843ff09c97b08904852d84/urllib3-2.6.3-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/0b/8c/cce26f6da9f69a9f0df032d642a4cb5603249f145c368e41b0154c463ca9/uv-0.10.3-py3-none-macosx_11_0_arm64.whl
       - pypi: https://files.pythonhosted.org/packages/e0/8a/e3462a43ec6712b74d921e4af9d5a2998752378c5554bde9a594dbb0cf0c/vhacdx-0.0.10-cp310-cp310-macosx_11_0_arm64.whl
-      - pypi: https://files.pythonhosted.org/packages/42/4b/6cf85b485be7ec29db837ec2a1d8cd68bc1147b1abf23d8636c5bd65b3cc/virtualenv-20.37.0-py3-none-any.whl
+      - pypi: https://files.pythonhosted.org/packages/20/0e/f083a76cb590e60dff3868779558eefefb8dfb7c9ed020babc7aa014ccbf/virtualenv-21.2.1-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/b9/c4/819f000b3d64f654a054a8b24bf8d0c5e799636393135784d0886d8c808e/viser-1.0.21-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/d5/4f/b462242432d93ea45f297b6179c7333dd0402b855a912a04e7fc61c0d71f/websockets-15.0.1-cp310-cp310-macosx_11_0_arm64.whl
       - pypi: https://files.pythonhosted.org/packages/87/22/b76d483683216dde3d67cba61fb2444be8d5be289bf628c13fc0fd90e5f9/wheel-0.46.3-py3-none-any.whl
@@ -2253,7 +2252,7 @@ environments:
       - pypi: https://files.pythonhosted.org/packages/09/ec/4b43dae793655b7d8a25f76119624350b4d65eb663459eb9603d7f1f0345/mistune-0.8.4-py2.py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/a4/8e/469e5a4a2f5855992e425f3cb33804cc07bf18d48f2db061aec61ce50270/more_itertools-10.8.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/98/bd/5db3c14d675ee12842afb9b70c94c64f2c873f31198c46cbfcd7dffafab0/msgspec-0.20.0-cp310-cp310-win_amd64.whl
-      - pypi: https://files.pythonhosted.org/packages/7e/91/d82dd3c16892e1b0e27a2f537eec8aad54d91d939cb3cd37db2e8c09ecc2/mujoco-3.5.0-cp310-cp310-win_amd64.whl
+      - pypi: https://files.pythonhosted.org/packages/eb/15/ae5df3d9b8fde39f7d75d5143bd9d0a74c1b189b385248667c072136cdea/mujoco-3.6.0-cp310-cp310-win_amd64.whl
       - pypi: https://files.pythonhosted.org/packages/b9/54/dd730b32ea14ea797530a4479b2ed46a6fb250f682a9cfb997e968bf0261/networkx-3.4.2-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/73/88/1ce287ef8649dc51365b5094bd3713b76454838140a32ab4f8349973883c/nh3-0.3.3-cp38-abi3-win_amd64.whl
       - pypi: https://files.pythonhosted.org/packages/88/b2/d0896bdcdc8d28a7fc5717c305f1a861c26e18c05047949fb371034d98bd/nodeenv-1.10.0-py2.py3-none-any.whl
@@ -2277,6 +2276,7 @@ environments:
       - pypi: https://files.pythonhosted.org/packages/ee/49/1377b49de7d0c1ce41292161ea0f721913fa8722c19fb9c1e3aa0367eecb/pytest_cov-7.0.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/ca/31/d4e37e9e550c2b92a9cbc2e4d0b7420a27224968580b5a447f420847c975/pytest_xdist-3.8.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/ec/57/56b9bcc3c9c6a792fcbaf139543cee77261f3651ca9da0c93f5c1221264b/python_dateutil-2.9.0.post0-py2.py3-none-any.whl
+      - pypi: https://files.pythonhosted.org/packages/d8/db/795879cc3ddfe338599bddea6388cc5100b088db0a4caf6e6c1af1c27e04/python_discovery-1.2.2-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/de/3d/8161f7711c017e01ac9f008dfddd9410dff3674334c233bde66e7ba65bbf/pywin32_ctypes-0.2.3-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/2a/fa/926c003379b19fca39dd4634818b00dec6c62d87faf628d1394e137354d4/pyyaml-6.0.3-cp310-cp310-win_amd64.whl
       - pypi: https://files.pythonhosted.org/packages/e1/67/921ec3024056483db83953ae8e48079ad62b92db7880013ca77632921dd0/readme_renderer-44.0-py3-none-any.whl
@@ -2315,7 +2315,7 @@ environments:
       - pypi: https://files.pythonhosted.org/packages/39/08/aaaad47bc4e9dc8c725e68f9d04865dbcb2052843ff09c97b08904852d84/urllib3-2.6.3-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/d5/95/89b9caeb4007d950dfa5fa5923c2ea67de541de0e3582598b7b9c7eeac53/uv-0.10.3-py3-none-win_amd64.whl
       - pypi: https://files.pythonhosted.org/packages/ff/aa/b401565542b927ce3e0a6d5e72acef79343a449ee1a7ad94a5c7266bab26/vhacdx-0.0.10-cp310-cp310-win_amd64.whl
-      - pypi: https://files.pythonhosted.org/packages/42/4b/6cf85b485be7ec29db837ec2a1d8cd68bc1147b1abf23d8636c5bd65b3cc/virtualenv-20.37.0-py3-none-any.whl
+      - pypi: https://files.pythonhosted.org/packages/20/0e/f083a76cb590e60dff3868779558eefefb8dfb7c9ed020babc7aa014ccbf/virtualenv-21.2.1-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/b9/c4/819f000b3d64f654a054a8b24bf8d0c5e799636393135784d0886d8c808e/viser-1.0.21-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/1c/46/aca7082012768bb98e5608f01658ff3ac8437e563eca41cf068bd5849a5e/websockets-15.0.1-cp310-cp310-win_amd64.whl
       - pypi: https://files.pythonhosted.org/packages/87/22/b76d483683216dde3d67cba61fb2444be8d5be289bf628c13fc0fd90e5f9/wheel-0.46.3-py3-none-any.whl
@@ -2329,8 +2329,6 @@ environments:
     - url: https://conda.anaconda.org/conda-forge/
     indexes:
     - https://pypi.org/simple
-    options:
-      pypi-prerelease-mode: if-necessary-or-explicit
     packages:
       linux-64:
       - conda: https://conda.anaconda.org/conda-forge/linux-64/_openmp_mutex-4.5-20_gnu.conda
@@ -2449,7 +2447,7 @@ environments:
       - pypi: https://files.pythonhosted.org/packages/09/ec/4b43dae793655b7d8a25f76119624350b4d65eb663459eb9603d7f1f0345/mistune-0.8.4-py2.py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/a4/8e/469e5a4a2f5855992e425f3cb33804cc07bf18d48f2db061aec61ce50270/more_itertools-10.8.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/6b/96/5c095b940de3aa6b43a71ec76275ac3537b21bd45c7499b5a17a429110fa/msgspec-0.20.0-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl
-      - pypi: https://files.pythonhosted.org/packages/47/14/d5bf98385354318ec2e6c466a8c7cf7fd76f8b711ed6d11d155e2baa81fb/mujoco-3.5.0-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl
+      - pypi: https://files.pythonhosted.org/packages/9b/a2/93ce08c5fcbe04dd997fe207bdc008e856b664ecf69db6442035aacf7fba/mujoco-3.6.0-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl
       - pypi: https://files.pythonhosted.org/packages/9e/c9/b2622292ea83fbb4ec318f5b9ab867d0a28ab43c5717bb85b0a5f6b3b0a4/networkx-3.6.1-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/6f/84/c0dc75c7fb596135f999e59a410d9f45bdabb989f1cb911f0016d22b747b/nh3-0.3.3-cp38-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
       - pypi: https://files.pythonhosted.org/packages/88/b2/d0896bdcdc8d28a7fc5717c305f1a861c26e18c05047949fb371034d98bd/nodeenv-1.10.0-py2.py3-none-any.whl
@@ -2473,6 +2471,7 @@ environments:
       - pypi: https://files.pythonhosted.org/packages/ee/49/1377b49de7d0c1ce41292161ea0f721913fa8722c19fb9c1e3aa0367eecb/pytest_cov-7.0.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/ca/31/d4e37e9e550c2b92a9cbc2e4d0b7420a27224968580b5a447f420847c975/pytest_xdist-3.8.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/ec/57/56b9bcc3c9c6a792fcbaf139543cee77261f3651ca9da0c93f5c1221264b/python_dateutil-2.9.0.post0-py2.py3-none-any.whl
+      - pypi: https://files.pythonhosted.org/packages/d8/db/795879cc3ddfe338599bddea6388cc5100b088db0a4caf6e6c1af1c27e04/python_discovery-1.2.2-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/71/60/917329f640924b18ff085ab889a11c763e0b573da888e8404ff486657602/pyyaml-6.0.3-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl
       - pypi: https://files.pythonhosted.org/packages/e1/67/921ec3024056483db83953ae8e48079ad62b92db7880013ca77632921dd0/readme_renderer-44.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/2c/58/ca301544e1fa93ed4f80d724bf5b194f6e4b945841c5bfd555878eea9fcb/referencing-0.37.0-py3-none-any.whl
@@ -2511,7 +2510,7 @@ environments:
       - pypi: https://files.pythonhosted.org/packages/39/08/aaaad47bc4e9dc8c725e68f9d04865dbcb2052843ff09c97b08904852d84/urllib3-2.6.3-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/ce/de/0638d8aa83c2038c5f820151006b3bf372d109afbfdb2a8fc59187cc3690/uv-0.10.3-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
       - pypi: https://files.pythonhosted.org/packages/a6/f1/464c761dbe24f58d6fc354bf51729342981fb7a621e170e0d3512fadbec8/vhacdx-0.0.10-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl
-      - pypi: https://files.pythonhosted.org/packages/42/4b/6cf85b485be7ec29db837ec2a1d8cd68bc1147b1abf23d8636c5bd65b3cc/virtualenv-20.37.0-py3-none-any.whl
+      - pypi: https://files.pythonhosted.org/packages/20/0e/f083a76cb590e60dff3868779558eefefb8dfb7c9ed020babc7aa014ccbf/virtualenv-21.2.1-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/b9/c4/819f000b3d64f654a054a8b24bf8d0c5e799636393135784d0886d8c808e/viser-1.0.21-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/93/1f/5d6dbf551766308f6f50f8baf8e9860be6182911e8106da7a7f73785f4c4/websockets-15.0.1-cp311-cp311-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl
       - pypi: https://files.pythonhosted.org/packages/87/22/b76d483683216dde3d67cba61fb2444be8d5be289bf628c13fc0fd90e5f9/wheel-0.46.3-py3-none-any.whl
@@ -2611,7 +2610,7 @@ environments:
       - pypi: https://files.pythonhosted.org/packages/09/ec/4b43dae793655b7d8a25f76119624350b4d65eb663459eb9603d7f1f0345/mistune-0.8.4-py2.py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/a4/8e/469e5a4a2f5855992e425f3cb33804cc07bf18d48f2db061aec61ce50270/more_itertools-10.8.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/03/59/fdcb3af72f750a8de2bcf39d62ada70b5eb17b06d7f63860e0a679cb656b/msgspec-0.20.0-cp311-cp311-macosx_10_9_x86_64.whl
-      - pypi: https://files.pythonhosted.org/packages/8b/47/e923589301c197c3ea0776b60cc0d57383b3cc51639ca75e4e4b6c5334d6/mujoco-3.5.0-cp311-cp311-macosx_10_16_x86_64.whl
+      - pypi: https://files.pythonhosted.org/packages/e0/75/d8afb4e98b58a119be3cba8da88b75cf53ff16f83baa9a14d37aa15a426e/mujoco-3.6.0-cp311-cp311-macosx_10_16_x86_64.whl
       - pypi: https://files.pythonhosted.org/packages/9e/c9/b2622292ea83fbb4ec318f5b9ab867d0a28ab43c5717bb85b0a5f6b3b0a4/networkx-3.6.1-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/13/3e/aef8cf8e0419b530c95e96ae93a5078e9b36c1e6613eeb1df03a80d5194e/nh3-0.3.3-cp38-abi3-macosx_10_12_x86_64.macosx_11_0_arm64.macosx_10_12_universal2.whl
       - pypi: https://files.pythonhosted.org/packages/88/b2/d0896bdcdc8d28a7fc5717c305f1a861c26e18c05047949fb371034d98bd/nodeenv-1.10.0-py2.py3-none-any.whl
@@ -2635,6 +2634,7 @@ environments:
       - pypi: https://files.pythonhosted.org/packages/ee/49/1377b49de7d0c1ce41292161ea0f721913fa8722c19fb9c1e3aa0367eecb/pytest_cov-7.0.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/ca/31/d4e37e9e550c2b92a9cbc2e4d0b7420a27224968580b5a447f420847c975/pytest_xdist-3.8.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/ec/57/56b9bcc3c9c6a792fcbaf139543cee77261f3651ca9da0c93f5c1221264b/python_dateutil-2.9.0.post0-py2.py3-none-any.whl
+      - pypi: https://files.pythonhosted.org/packages/d8/db/795879cc3ddfe338599bddea6388cc5100b088db0a4caf6e6c1af1c27e04/python_discovery-1.2.2-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/6d/16/a95b6757765b7b031c9374925bb718d55e0a9ba8a1b6a12d25962ea44347/pyyaml-6.0.3-cp311-cp311-macosx_10_13_x86_64.whl
       - pypi: https://files.pythonhosted.org/packages/e1/67/921ec3024056483db83953ae8e48079ad62b92db7880013ca77632921dd0/readme_renderer-44.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/2c/58/ca301544e1fa93ed4f80d724bf5b194f6e4b945841c5bfd555878eea9fcb/referencing-0.37.0-py3-none-any.whl
@@ -2672,7 +2672,7 @@ environments:
       - pypi: https://files.pythonhosted.org/packages/39/08/aaaad47bc4e9dc8c725e68f9d04865dbcb2052843ff09c97b08904852d84/urllib3-2.6.3-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/40/db/c5729ec39ed557ac673ca993abd71e5b18d71a92f95f1845152e639000ff/uv-0.10.3-py3-none-macosx_10_12_x86_64.whl
       - pypi: https://files.pythonhosted.org/packages/b7/2c/d49df6fec3294cef3c8c88c54784162bd8350c427fecd9b16335772b760f/vhacdx-0.0.10-cp311-cp311-macosx_10_9_x86_64.whl
-      - pypi: https://files.pythonhosted.org/packages/42/4b/6cf85b485be7ec29db837ec2a1d8cd68bc1147b1abf23d8636c5bd65b3cc/virtualenv-20.37.0-py3-none-any.whl
+      - pypi: https://files.pythonhosted.org/packages/20/0e/f083a76cb590e60dff3868779558eefefb8dfb7c9ed020babc7aa014ccbf/virtualenv-21.2.1-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/b9/c4/819f000b3d64f654a054a8b24bf8d0c5e799636393135784d0886d8c808e/viser-1.0.21-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/76/70/ba1ad96b07869275ef42e2ce21f07a5b0148936688c2baf7e4a1f60d5058/websockets-15.0.1-cp311-cp311-macosx_10_9_x86_64.whl
       - pypi: https://files.pythonhosted.org/packages/87/22/b76d483683216dde3d67cba61fb2444be8d5be289bf628c13fc0fd90e5f9/wheel-0.46.3-py3-none-any.whl
@@ -2772,7 +2772,7 @@ environments:
       - pypi: https://files.pythonhosted.org/packages/09/ec/4b43dae793655b7d8a25f76119624350b4d65eb663459eb9603d7f1f0345/mistune-0.8.4-py2.py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/a4/8e/469e5a4a2f5855992e425f3cb33804cc07bf18d48f2db061aec61ce50270/more_itertools-10.8.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/5a/15/3c225610da9f02505d37d69a77f4a2e7daae2a125f99d638df211ba84e59/msgspec-0.20.0-cp311-cp311-macosx_11_0_arm64.whl
-      - pypi: https://files.pythonhosted.org/packages/82/02/aa6057ac4c50fb36558208005d6da19518f9a7857ef9b5fd2ed8f9262fe2/mujoco-3.5.0-cp311-cp311-macosx_11_0_arm64.whl
+      - pypi: https://files.pythonhosted.org/packages/d7/fb/5335c0ba2e88f4b8f8300c15966823dbb96ecc906a61c86bcf9cdac77311/mujoco-3.6.0-cp311-cp311-macosx_11_0_arm64.whl
       - pypi: https://files.pythonhosted.org/packages/9e/c9/b2622292ea83fbb4ec318f5b9ab867d0a28ab43c5717bb85b0a5f6b3b0a4/networkx-3.6.1-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/13/3e/aef8cf8e0419b530c95e96ae93a5078e9b36c1e6613eeb1df03a80d5194e/nh3-0.3.3-cp38-abi3-macosx_10_12_x86_64.macosx_11_0_arm64.macosx_10_12_universal2.whl
       - pypi: https://files.pythonhosted.org/packages/88/b2/d0896bdcdc8d28a7fc5717c305f1a861c26e18c05047949fb371034d98bd/nodeenv-1.10.0-py2.py3-none-any.whl
@@ -2796,6 +2796,7 @@ environments:
       - pypi: https://files.pythonhosted.org/packages/ee/49/1377b49de7d0c1ce41292161ea0f721913fa8722c19fb9c1e3aa0367eecb/pytest_cov-7.0.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/ca/31/d4e37e9e550c2b92a9cbc2e4d0b7420a27224968580b5a447f420847c975/pytest_xdist-3.8.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/ec/57/56b9bcc3c9c6a792fcbaf139543cee77261f3651ca9da0c93f5c1221264b/python_dateutil-2.9.0.post0-py2.py3-none-any.whl
+      - pypi: https://files.pythonhosted.org/packages/d8/db/795879cc3ddfe338599bddea6388cc5100b088db0a4caf6e6c1af1c27e04/python_discovery-1.2.2-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/16/19/13de8e4377ed53079ee996e1ab0a9c33ec2faf808a4647b7b4c0d46dd239/pyyaml-6.0.3-cp311-cp311-macosx_11_0_arm64.whl
       - pypi: https://files.pythonhosted.org/packages/e1/67/921ec3024056483db83953ae8e48079ad62b92db7880013ca77632921dd0/readme_renderer-44.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/2c/58/ca301544e1fa93ed4f80d724bf5b194f6e4b945841c5bfd555878eea9fcb/referencing-0.37.0-py3-none-any.whl
@@ -2833,7 +2834,7 @@ environments:
       - pypi: https://files.pythonhosted.org/packages/39/08/aaaad47bc4e9dc8c725e68f9d04865dbcb2052843ff09c97b08904852d84/urllib3-2.6.3-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/0b/8c/cce26f6da9f69a9f0df032d642a4cb5603249f145c368e41b0154c463ca9/uv-0.10.3-py3-none-macosx_11_0_arm64.whl
       - pypi: https://files.pythonhosted.org/packages/68/1d/bd2456baa6b16977c106adc2386b6e7a34c3e57ade6aeeab68bb61ceb16f/vhacdx-0.0.10-cp311-cp311-macosx_11_0_arm64.whl
-      - pypi: https://files.pythonhosted.org/packages/42/4b/6cf85b485be7ec29db837ec2a1d8cd68bc1147b1abf23d8636c5bd65b3cc/virtualenv-20.37.0-py3-none-any.whl
+      - pypi: https://files.pythonhosted.org/packages/20/0e/f083a76cb590e60dff3868779558eefefb8dfb7c9ed020babc7aa014ccbf/virtualenv-21.2.1-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/b9/c4/819f000b3d64f654a054a8b24bf8d0c5e799636393135784d0886d8c808e/viser-1.0.21-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/86/f2/10b55821dd40eb696ce4704a87d57774696f9451108cff0d2824c97e0f97/websockets-15.0.1-cp311-cp311-macosx_11_0_arm64.whl
       - pypi: https://files.pythonhosted.org/packages/87/22/b76d483683216dde3d67cba61fb2444be8d5be289bf628c13fc0fd90e5f9/wheel-0.46.3-py3-none-any.whl
@@ -2929,7 +2930,7 @@ environments:
       - pypi: https://files.pythonhosted.org/packages/09/ec/4b43dae793655b7d8a25f76119624350b4d65eb663459eb9603d7f1f0345/mistune-0.8.4-py2.py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/a4/8e/469e5a4a2f5855992e425f3cb33804cc07bf18d48f2db061aec61ce50270/more_itertools-10.8.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/89/5e/406b7d578926b68790e390d83a1165a9bfc2d95612a1a9c1c4d5c72ea815/msgspec-0.20.0-cp311-cp311-win_amd64.whl
-      - pypi: https://files.pythonhosted.org/packages/b8/98/c1fac334cc764068e6c5d7eb01d6ed2a3392bab51952c816888b2dfe78c2/mujoco-3.5.0-cp311-cp311-win_amd64.whl
+      - pypi: https://files.pythonhosted.org/packages/3f/a5/a84a9edc6234a2ee29a6b008d432e1c3855795f10a51786ee2260128ed12/mujoco-3.6.0-cp311-cp311-win_amd64.whl
       - pypi: https://files.pythonhosted.org/packages/9e/c9/b2622292ea83fbb4ec318f5b9ab867d0a28ab43c5717bb85b0a5f6b3b0a4/networkx-3.6.1-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/73/88/1ce287ef8649dc51365b5094bd3713b76454838140a32ab4f8349973883c/nh3-0.3.3-cp38-abi3-win_amd64.whl
       - pypi: https://files.pythonhosted.org/packages/88/b2/d0896bdcdc8d28a7fc5717c305f1a861c26e18c05047949fb371034d98bd/nodeenv-1.10.0-py2.py3-none-any.whl
@@ -2953,6 +2954,7 @@ environments:
       - pypi: https://files.pythonhosted.org/packages/ee/49/1377b49de7d0c1ce41292161ea0f721913fa8722c19fb9c1e3aa0367eecb/pytest_cov-7.0.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/ca/31/d4e37e9e550c2b92a9cbc2e4d0b7420a27224968580b5a447f420847c975/pytest_xdist-3.8.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/ec/57/56b9bcc3c9c6a792fcbaf139543cee77261f3651ca9da0c93f5c1221264b/python_dateutil-2.9.0.post0-py2.py3-none-any.whl
+      - pypi: https://files.pythonhosted.org/packages/d8/db/795879cc3ddfe338599bddea6388cc5100b088db0a4caf6e6c1af1c27e04/python_discovery-1.2.2-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/de/3d/8161f7711c017e01ac9f008dfddd9410dff3674334c233bde66e7ba65bbf/pywin32_ctypes-0.2.3-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/da/e3/ea007450a105ae919a72393cb06f122f288ef60bba2dc64b26e2646fa315/pyyaml-6.0.3-cp311-cp311-win_amd64.whl
       - pypi: https://files.pythonhosted.org/packages/e1/67/921ec3024056483db83953ae8e48079ad62b92db7880013ca77632921dd0/readme_renderer-44.0-py3-none-any.whl
@@ -2991,7 +2993,7 @@ environments:
       - pypi: https://files.pythonhosted.org/packages/39/08/aaaad47bc4e9dc8c725e68f9d04865dbcb2052843ff09c97b08904852d84/urllib3-2.6.3-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/d5/95/89b9caeb4007d950dfa5fa5923c2ea67de541de0e3582598b7b9c7eeac53/uv-0.10.3-py3-none-win_amd64.whl
       - pypi: https://files.pythonhosted.org/packages/de/c8/a8260b780e4578d7ef19b70343f9717f74ff48f9950138c96c78f209ec01/vhacdx-0.0.10-cp311-cp311-win_amd64.whl
-      - pypi: https://files.pythonhosted.org/packages/42/4b/6cf85b485be7ec29db837ec2a1d8cd68bc1147b1abf23d8636c5bd65b3cc/virtualenv-20.37.0-py3-none-any.whl
+      - pypi: https://files.pythonhosted.org/packages/20/0e/f083a76cb590e60dff3868779558eefefb8dfb7c9ed020babc7aa014ccbf/virtualenv-21.2.1-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/b9/c4/819f000b3d64f654a054a8b24bf8d0c5e799636393135784d0886d8c808e/viser-1.0.21-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/98/93/e36c73f78400a65f5e236cd376713c34182e6663f6889cd45a4a04d8f203/websockets-15.0.1-cp311-cp311-win_amd64.whl
       - pypi: https://files.pythonhosted.org/packages/87/22/b76d483683216dde3d67cba61fb2444be8d5be289bf628c13fc0fd90e5f9/wheel-0.46.3-py3-none-any.whl
@@ -3005,8 +3007,6 @@ environments:
     - url: https://conda.anaconda.org/conda-forge/
     indexes:
     - https://pypi.org/simple
-    options:
-      pypi-prerelease-mode: if-necessary-or-explicit
     packages:
       linux-64:
       - conda: https://conda.anaconda.org/conda-forge/linux-64/_openmp_mutex-4.5-20_gnu.conda
@@ -3123,7 +3123,7 @@ environments:
       - pypi: https://files.pythonhosted.org/packages/09/ec/4b43dae793655b7d8a25f76119624350b4d65eb663459eb9603d7f1f0345/mistune-0.8.4-py2.py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/a4/8e/469e5a4a2f5855992e425f3cb33804cc07bf18d48f2db061aec61ce50270/more_itertools-10.8.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/5c/a2/488517a43ccf5a4b6b6eca6dd4ede0bd82b043d1539dd6bb908a19f8efd3/msgspec-0.20.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl
-      - pypi: https://files.pythonhosted.org/packages/c8/8a/229e4db3692be55532e155e2ca6a1363752243ee79df0e7e22ba00f716cf/mujoco-3.5.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl
+      - pypi: https://files.pythonhosted.org/packages/52/6c/5ec4e93676a65064a6591176772e00cfa02716156a1d0a7d646a8203348f/mujoco-3.6.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl
       - pypi: https://files.pythonhosted.org/packages/9e/c9/b2622292ea83fbb4ec318f5b9ab867d0a28ab43c5717bb85b0a5f6b3b0a4/networkx-3.6.1-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/6f/84/c0dc75c7fb596135f999e59a410d9f45bdabb989f1cb911f0016d22b747b/nh3-0.3.3-cp38-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
       - pypi: https://files.pythonhosted.org/packages/88/b2/d0896bdcdc8d28a7fc5717c305f1a861c26e18c05047949fb371034d98bd/nodeenv-1.10.0-py2.py3-none-any.whl
@@ -3147,6 +3147,7 @@ environments:
       - pypi: https://files.pythonhosted.org/packages/ee/49/1377b49de7d0c1ce41292161ea0f721913fa8722c19fb9c1e3aa0367eecb/pytest_cov-7.0.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/ca/31/d4e37e9e550c2b92a9cbc2e4d0b7420a27224968580b5a447f420847c975/pytest_xdist-3.8.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/ec/57/56b9bcc3c9c6a792fcbaf139543cee77261f3651ca9da0c93f5c1221264b/python_dateutil-2.9.0.post0-py2.py3-none-any.whl
+      - pypi: https://files.pythonhosted.org/packages/d8/db/795879cc3ddfe338599bddea6388cc5100b088db0a4caf6e6c1af1c27e04/python_discovery-1.2.2-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/8b/9d/b3589d3877982d4f2329302ef98a8026e7f4443c765c46cfecc8858c6b4b/pyyaml-6.0.3-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl
       - pypi: https://files.pythonhosted.org/packages/e1/67/921ec3024056483db83953ae8e48079ad62b92db7880013ca77632921dd0/readme_renderer-44.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/2c/58/ca301544e1fa93ed4f80d724bf5b194f6e4b945841c5bfd555878eea9fcb/referencing-0.37.0-py3-none-any.whl
@@ -3185,7 +3186,7 @@ environments:
       - pypi: https://files.pythonhosted.org/packages/39/08/aaaad47bc4e9dc8c725e68f9d04865dbcb2052843ff09c97b08904852d84/urllib3-2.6.3-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/ce/de/0638d8aa83c2038c5f820151006b3bf372d109afbfdb2a8fc59187cc3690/uv-0.10.3-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
       - pypi: https://files.pythonhosted.org/packages/94/98/d2a6aeb1c6570a1fc1be29ee03db795f643ab03c6df7635522f23796b39d/vhacdx-0.0.10-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl
-      - pypi: https://files.pythonhosted.org/packages/42/4b/6cf85b485be7ec29db837ec2a1d8cd68bc1147b1abf23d8636c5bd65b3cc/virtualenv-20.37.0-py3-none-any.whl
+      - pypi: https://files.pythonhosted.org/packages/20/0e/f083a76cb590e60dff3868779558eefefb8dfb7c9ed020babc7aa014ccbf/virtualenv-21.2.1-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/b9/c4/819f000b3d64f654a054a8b24bf8d0c5e799636393135784d0886d8c808e/viser-1.0.21-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/14/8f/aa61f528fba38578ec553c145857a181384c72b98156f858ca5c8e82d9d3/websockets-15.0.1-cp312-cp312-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl
       - pypi: https://files.pythonhosted.org/packages/87/22/b76d483683216dde3d67cba61fb2444be8d5be289bf628c13fc0fd90e5f9/wheel-0.46.3-py3-none-any.whl
@@ -3283,7 +3284,7 @@ environments:
       - pypi: https://files.pythonhosted.org/packages/09/ec/4b43dae793655b7d8a25f76119624350b4d65eb663459eb9603d7f1f0345/mistune-0.8.4-py2.py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/a4/8e/469e5a4a2f5855992e425f3cb33804cc07bf18d48f2db061aec61ce50270/more_itertools-10.8.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/d9/6f/1e25eee957e58e3afb2a44b94fa95e06cebc4c236193ed0de3012fff1e19/msgspec-0.20.0-cp312-cp312-macosx_10_13_x86_64.whl
-      - pypi: https://files.pythonhosted.org/packages/f9/f0/4772421643f1c5aaf46d9e500a8716f59b02c8bf30bfa92cb8a763159efb/mujoco-3.5.0-cp312-cp312-macosx_10_16_x86_64.whl
+      - pypi: https://files.pythonhosted.org/packages/38/c4/f8959e3d5d98b282e081ce08d07cd71ae949cc0ad9f2c39c0a69fcb88c8c/mujoco-3.6.0-cp312-cp312-macosx_10_16_x86_64.whl
       - pypi: https://files.pythonhosted.org/packages/9e/c9/b2622292ea83fbb4ec318f5b9ab867d0a28ab43c5717bb85b0a5f6b3b0a4/networkx-3.6.1-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/13/3e/aef8cf8e0419b530c95e96ae93a5078e9b36c1e6613eeb1df03a80d5194e/nh3-0.3.3-cp38-abi3-macosx_10_12_x86_64.macosx_11_0_arm64.macosx_10_12_universal2.whl
       - pypi: https://files.pythonhosted.org/packages/88/b2/d0896bdcdc8d28a7fc5717c305f1a861c26e18c05047949fb371034d98bd/nodeenv-1.10.0-py2.py3-none-any.whl
@@ -3307,6 +3308,7 @@ environments:
       - pypi: https://files.pythonhosted.org/packages/ee/49/1377b49de7d0c1ce41292161ea0f721913fa8722c19fb9c1e3aa0367eecb/pytest_cov-7.0.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/ca/31/d4e37e9e550c2b92a9cbc2e4d0b7420a27224968580b5a447f420847c975/pytest_xdist-3.8.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/ec/57/56b9bcc3c9c6a792fcbaf139543cee77261f3651ca9da0c93f5c1221264b/python_dateutil-2.9.0.post0-py2.py3-none-any.whl
+      - pypi: https://files.pythonhosted.org/packages/d8/db/795879cc3ddfe338599bddea6388cc5100b088db0a4caf6e6c1af1c27e04/python_discovery-1.2.2-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/d1/33/422b98d2195232ca1826284a76852ad5a86fe23e31b009c9886b2d0fb8b2/pyyaml-6.0.3-cp312-cp312-macosx_10_13_x86_64.whl
       - pypi: https://files.pythonhosted.org/packages/e1/67/921ec3024056483db83953ae8e48079ad62b92db7880013ca77632921dd0/readme_renderer-44.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/2c/58/ca301544e1fa93ed4f80d724bf5b194f6e4b945841c5bfd555878eea9fcb/referencing-0.37.0-py3-none-any.whl
@@ -3344,7 +3346,7 @@ environments:
       - pypi: https://files.pythonhosted.org/packages/39/08/aaaad47bc4e9dc8c725e68f9d04865dbcb2052843ff09c97b08904852d84/urllib3-2.6.3-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/40/db/c5729ec39ed557ac673ca993abd71e5b18d71a92f95f1845152e639000ff/uv-0.10.3-py3-none-macosx_10_12_x86_64.whl
       - pypi: https://files.pythonhosted.org/packages/cf/9c/66375e65634c80f6efb46e81915126bf3e55dc9d6615217590cbc8316d2e/vhacdx-0.0.10-cp312-cp312-macosx_10_13_x86_64.whl
-      - pypi: https://files.pythonhosted.org/packages/42/4b/6cf85b485be7ec29db837ec2a1d8cd68bc1147b1abf23d8636c5bd65b3cc/virtualenv-20.37.0-py3-none-any.whl
+      - pypi: https://files.pythonhosted.org/packages/20/0e/f083a76cb590e60dff3868779558eefefb8dfb7c9ed020babc7aa014ccbf/virtualenv-21.2.1-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/b9/c4/819f000b3d64f654a054a8b24bf8d0c5e799636393135784d0886d8c808e/viser-1.0.21-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/f4/71/809a0f5f6a06522af902e0f2ea2757f71ead94610010cf570ab5c98e99ed/websockets-15.0.1-cp312-cp312-macosx_10_13_x86_64.whl
       - pypi: https://files.pythonhosted.org/packages/87/22/b76d483683216dde3d67cba61fb2444be8d5be289bf628c13fc0fd90e5f9/wheel-0.46.3-py3-none-any.whl
@@ -3442,7 +3444,7 @@ environments:
       - pypi: https://files.pythonhosted.org/packages/09/ec/4b43dae793655b7d8a25f76119624350b4d65eb663459eb9603d7f1f0345/mistune-0.8.4-py2.py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/a4/8e/469e5a4a2f5855992e425f3cb33804cc07bf18d48f2db061aec61ce50270/more_itertools-10.8.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/7f/ee/af51d090ada641d4b264992a486435ba3ef5b5634bc27e6eb002f71cef7d/msgspec-0.20.0-cp312-cp312-macosx_11_0_arm64.whl
-      - pypi: https://files.pythonhosted.org/packages/e1/d4/d0032323f58a9b8080b8464c6aade8d5ac2e101dbed1de64a38b3913b446/mujoco-3.5.0-cp312-cp312-macosx_11_0_arm64.whl
+      - pypi: https://files.pythonhosted.org/packages/26/55/7407eced2c44fbea233302d2c11e778852ea0f2eb0e14610f13a7e0d6ac7/mujoco-3.6.0-cp312-cp312-macosx_11_0_arm64.whl
       - pypi: https://files.pythonhosted.org/packages/9e/c9/b2622292ea83fbb4ec318f5b9ab867d0a28ab43c5717bb85b0a5f6b3b0a4/networkx-3.6.1-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/13/3e/aef8cf8e0419b530c95e96ae93a5078e9b36c1e6613eeb1df03a80d5194e/nh3-0.3.3-cp38-abi3-macosx_10_12_x86_64.macosx_11_0_arm64.macosx_10_12_universal2.whl
       - pypi: https://files.pythonhosted.org/packages/88/b2/d0896bdcdc8d28a7fc5717c305f1a861c26e18c05047949fb371034d98bd/nodeenv-1.10.0-py2.py3-none-any.whl
@@ -3466,6 +3468,7 @@ environments:
       - pypi: https://files.pythonhosted.org/packages/ee/49/1377b49de7d0c1ce41292161ea0f721913fa8722c19fb9c1e3aa0367eecb/pytest_cov-7.0.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/ca/31/d4e37e9e550c2b92a9cbc2e4d0b7420a27224968580b5a447f420847c975/pytest_xdist-3.8.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/ec/57/56b9bcc3c9c6a792fcbaf139543cee77261f3651ca9da0c93f5c1221264b/python_dateutil-2.9.0.post0-py2.py3-none-any.whl
+      - pypi: https://files.pythonhosted.org/packages/d8/db/795879cc3ddfe338599bddea6388cc5100b088db0a4caf6e6c1af1c27e04/python_discovery-1.2.2-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/89/a0/6cf41a19a1f2f3feab0e9c0b74134aa2ce6849093d5517a0c550fe37a648/pyyaml-6.0.3-cp312-cp312-macosx_11_0_arm64.whl
       - pypi: https://files.pythonhosted.org/packages/e1/67/921ec3024056483db83953ae8e48079ad62b92db7880013ca77632921dd0/readme_renderer-44.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/2c/58/ca301544e1fa93ed4f80d724bf5b194f6e4b945841c5bfd555878eea9fcb/referencing-0.37.0-py3-none-any.whl
@@ -3503,7 +3506,7 @@ environments:
       - pypi: https://files.pythonhosted.org/packages/39/08/aaaad47bc4e9dc8c725e68f9d04865dbcb2052843ff09c97b08904852d84/urllib3-2.6.3-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/0b/8c/cce26f6da9f69a9f0df032d642a4cb5603249f145c368e41b0154c463ca9/uv-0.10.3-py3-none-macosx_11_0_arm64.whl
       - pypi: https://files.pythonhosted.org/packages/4e/e3/fc2644d3e7d0b2b52e2f681eb2878c0e1b9cafc53946f66736d0f01e237c/vhacdx-0.0.10-cp312-cp312-macosx_11_0_arm64.whl
-      - pypi: https://files.pythonhosted.org/packages/42/4b/6cf85b485be7ec29db837ec2a1d8cd68bc1147b1abf23d8636c5bd65b3cc/virtualenv-20.37.0-py3-none-any.whl
+      - pypi: https://files.pythonhosted.org/packages/20/0e/f083a76cb590e60dff3868779558eefefb8dfb7c9ed020babc7aa014ccbf/virtualenv-21.2.1-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/b9/c4/819f000b3d64f654a054a8b24bf8d0c5e799636393135784d0886d8c808e/viser-1.0.21-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/3d/69/1a681dd6f02180916f116894181eab8b2e25b31e484c5d0eae637ec01f7c/websockets-15.0.1-cp312-cp312-macosx_11_0_arm64.whl
       - pypi: https://files.pythonhosted.org/packages/87/22/b76d483683216dde3d67cba61fb2444be8d5be289bf628c13fc0fd90e5f9/wheel-0.46.3-py3-none-any.whl
@@ -3597,7 +3600,7 @@ environments:
       - pypi: https://files.pythonhosted.org/packages/09/ec/4b43dae793655b7d8a25f76119624350b4d65eb663459eb9603d7f1f0345/mistune-0.8.4-py2.py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/a4/8e/469e5a4a2f5855992e425f3cb33804cc07bf18d48f2db061aec61ce50270/more_itertools-10.8.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/5a/51/aba940212c23b32eedce752896205912c2668472ed5b205fc33da28a6509/msgspec-0.20.0-cp312-cp312-win_amd64.whl
-      - pypi: https://files.pythonhosted.org/packages/02/37/527d83610b878f27c01dd762e0e41aaa62f095c607f0500ac7f724a2c7a5/mujoco-3.5.0-cp312-cp312-win_amd64.whl
+      - pypi: https://files.pythonhosted.org/packages/92/22/38d82f0c34213af53afbbb248b3442943ef48ffbac1e4c909b321e02ac56/mujoco-3.6.0-cp312-cp312-win_amd64.whl
       - pypi: https://files.pythonhosted.org/packages/9e/c9/b2622292ea83fbb4ec318f5b9ab867d0a28ab43c5717bb85b0a5f6b3b0a4/networkx-3.6.1-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/73/88/1ce287ef8649dc51365b5094bd3713b76454838140a32ab4f8349973883c/nh3-0.3.3-cp38-abi3-win_amd64.whl
       - pypi: https://files.pythonhosted.org/packages/88/b2/d0896bdcdc8d28a7fc5717c305f1a861c26e18c05047949fb371034d98bd/nodeenv-1.10.0-py2.py3-none-any.whl
@@ -3621,6 +3624,7 @@ environments:
       - pypi: https://files.pythonhosted.org/packages/ee/49/1377b49de7d0c1ce41292161ea0f721913fa8722c19fb9c1e3aa0367eecb/pytest_cov-7.0.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/ca/31/d4e37e9e550c2b92a9cbc2e4d0b7420a27224968580b5a447f420847c975/pytest_xdist-3.8.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/ec/57/56b9bcc3c9c6a792fcbaf139543cee77261f3651ca9da0c93f5c1221264b/python_dateutil-2.9.0.post0-py2.py3-none-any.whl
+      - pypi: https://files.pythonhosted.org/packages/d8/db/795879cc3ddfe338599bddea6388cc5100b088db0a4caf6e6c1af1c27e04/python_discovery-1.2.2-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/de/3d/8161f7711c017e01ac9f008dfddd9410dff3674334c233bde66e7ba65bbf/pywin32_ctypes-0.2.3-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/86/bf/899e81e4cce32febab4fb42bb97dcdf66bc135272882d1987881a4b519e9/pyyaml-6.0.3-cp312-cp312-win_amd64.whl
       - pypi: https://files.pythonhosted.org/packages/e1/67/921ec3024056483db83953ae8e48079ad62b92db7880013ca77632921dd0/readme_renderer-44.0-py3-none-any.whl
@@ -3659,7 +3663,7 @@ environments:
       - pypi: https://files.pythonhosted.org/packages/39/08/aaaad47bc4e9dc8c725e68f9d04865dbcb2052843ff09c97b08904852d84/urllib3-2.6.3-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/d5/95/89b9caeb4007d950dfa5fa5923c2ea67de541de0e3582598b7b9c7eeac53/uv-0.10.3-py3-none-win_amd64.whl
       - pypi: https://files.pythonhosted.org/packages/73/e9/f9729603ac75047a257f1b4ddac60cbde72b0abfd49ffed305751ba630a2/vhacdx-0.0.10-cp312-cp312-win_amd64.whl
-      - pypi: https://files.pythonhosted.org/packages/42/4b/6cf85b485be7ec29db837ec2a1d8cd68bc1147b1abf23d8636c5bd65b3cc/virtualenv-20.37.0-py3-none-any.whl
+      - pypi: https://files.pythonhosted.org/packages/20/0e/f083a76cb590e60dff3868779558eefefb8dfb7c9ed020babc7aa014ccbf/virtualenv-21.2.1-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/b9/c4/819f000b3d64f654a054a8b24bf8d0c5e799636393135784d0886d8c808e/viser-1.0.21-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/7d/71/abf2ebc3bbfa40f391ce1428c7168fb20582d0ff57019b69ea20fa698043/websockets-15.0.1-cp312-cp312-win_amd64.whl
       - pypi: https://files.pythonhosted.org/packages/87/22/b76d483683216dde3d67cba61fb2444be8d5be289bf628c13fc0fd90e5f9/wheel-0.46.3-py3-none-any.whl
@@ -3673,8 +3677,6 @@ environments:
     - url: https://conda.anaconda.org/conda-forge/
     indexes:
     - https://pypi.org/simple
-    options:
-      pypi-prerelease-mode: if-necessary-or-explicit
     packages:
       linux-64:
       - conda: https://conda.anaconda.org/conda-forge/linux-64/_openmp_mutex-4.5-20_gnu.conda
@@ -3791,7 +3793,7 @@ environments:
       - pypi: https://files.pythonhosted.org/packages/09/ec/4b43dae793655b7d8a25f76119624350b4d65eb663459eb9603d7f1f0345/mistune-0.8.4-py2.py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/a4/8e/469e5a4a2f5855992e425f3cb33804cc07bf18d48f2db061aec61ce50270/more_itertools-10.8.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/28/83/36557b04cfdc317ed8a525c4993b23e43a8fbcddaddd78619112ca07138c/msgspec-0.20.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl
-      - pypi: https://files.pythonhosted.org/packages/16/47/93c7ac3a9630b49c55d76b0d02aa565543e2f62cecd885f8f574f5c745e7/mujoco-3.5.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl
+      - pypi: https://files.pythonhosted.org/packages/6e/48/c8cd52847d8a973fc606910a5467b8b7b68fa763afbe91f41d87123f957c/mujoco-3.6.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl
       - pypi: https://files.pythonhosted.org/packages/9e/c9/b2622292ea83fbb4ec318f5b9ab867d0a28ab43c5717bb85b0a5f6b3b0a4/networkx-3.6.1-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/6f/84/c0dc75c7fb596135f999e59a410d9f45bdabb989f1cb911f0016d22b747b/nh3-0.3.3-cp38-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
       - pypi: https://files.pythonhosted.org/packages/88/b2/d0896bdcdc8d28a7fc5717c305f1a861c26e18c05047949fb371034d98bd/nodeenv-1.10.0-py2.py3-none-any.whl
@@ -3815,6 +3817,7 @@ environments:
       - pypi: https://files.pythonhosted.org/packages/ee/49/1377b49de7d0c1ce41292161ea0f721913fa8722c19fb9c1e3aa0367eecb/pytest_cov-7.0.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/ca/31/d4e37e9e550c2b92a9cbc2e4d0b7420a27224968580b5a447f420847c975/pytest_xdist-3.8.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/ec/57/56b9bcc3c9c6a792fcbaf139543cee77261f3651ca9da0c93f5c1221264b/python_dateutil-2.9.0.post0-py2.py3-none-any.whl
+      - pypi: https://files.pythonhosted.org/packages/d8/db/795879cc3ddfe338599bddea6388cc5100b088db0a4caf6e6c1af1c27e04/python_discovery-1.2.2-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/74/27/e5b8f34d02d9995b80abcef563ea1f8b56d20134d8f4e5e81733b1feceb2/pyyaml-6.0.3-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl
       - pypi: https://files.pythonhosted.org/packages/e1/67/921ec3024056483db83953ae8e48079ad62b92db7880013ca77632921dd0/readme_renderer-44.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/2c/58/ca301544e1fa93ed4f80d724bf5b194f6e4b945841c5bfd555878eea9fcb/referencing-0.37.0-py3-none-any.whl
@@ -3853,7 +3856,7 @@ environments:
       - pypi: https://files.pythonhosted.org/packages/39/08/aaaad47bc4e9dc8c725e68f9d04865dbcb2052843ff09c97b08904852d84/urllib3-2.6.3-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/ce/de/0638d8aa83c2038c5f820151006b3bf372d109afbfdb2a8fc59187cc3690/uv-0.10.3-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
       - pypi: https://files.pythonhosted.org/packages/b8/b4/07ab1c828bae0eb5c72cd9a4cbe8b0376d374509be3c7055e1a399bf85c3/vhacdx-0.0.10-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl
-      - pypi: https://files.pythonhosted.org/packages/42/4b/6cf85b485be7ec29db837ec2a1d8cd68bc1147b1abf23d8636c5bd65b3cc/virtualenv-20.37.0-py3-none-any.whl
+      - pypi: https://files.pythonhosted.org/packages/20/0e/f083a76cb590e60dff3868779558eefefb8dfb7c9ed020babc7aa014ccbf/virtualenv-21.2.1-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/b9/c4/819f000b3d64f654a054a8b24bf8d0c5e799636393135784d0886d8c808e/viser-1.0.21-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/ff/b2/83a6ddf56cdcbad4e3d841fcc55d6ba7d19aeb89c50f24dd7e859ec0805f/websockets-15.0.1-cp313-cp313-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl
       - pypi: https://files.pythonhosted.org/packages/87/22/b76d483683216dde3d67cba61fb2444be8d5be289bf628c13fc0fd90e5f9/wheel-0.46.3-py3-none-any.whl
@@ -3953,7 +3956,7 @@ environments:
       - pypi: https://files.pythonhosted.org/packages/09/ec/4b43dae793655b7d8a25f76119624350b4d65eb663459eb9603d7f1f0345/mistune-0.8.4-py2.py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/a4/8e/469e5a4a2f5855992e425f3cb33804cc07bf18d48f2db061aec61ce50270/more_itertools-10.8.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/8a/d1/b902d38b6e5ba3bdddbec469bba388d647f960aeed7b5b3623a8debe8a76/msgspec-0.20.0-cp313-cp313-macosx_10_13_x86_64.whl
-      - pypi: https://files.pythonhosted.org/packages/87/2a/371033684e4ddcda47c97661fb6e9617c0e5e3749af082a9b4d5d1bf9f27/mujoco-3.5.0-cp313-cp313-macosx_10_16_x86_64.whl
+      - pypi: https://files.pythonhosted.org/packages/29/0e/f3ea1fc9d1a25f19b173b13c23797805480cdb0a1026d43cf6b37dc2de6e/mujoco-3.6.0-cp313-cp313-macosx_10_16_x86_64.whl
       - pypi: https://files.pythonhosted.org/packages/9e/c9/b2622292ea83fbb4ec318f5b9ab867d0a28ab43c5717bb85b0a5f6b3b0a4/networkx-3.6.1-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/13/3e/aef8cf8e0419b530c95e96ae93a5078e9b36c1e6613eeb1df03a80d5194e/nh3-0.3.3-cp38-abi3-macosx_10_12_x86_64.macosx_11_0_arm64.macosx_10_12_universal2.whl
       - pypi: https://files.pythonhosted.org/packages/88/b2/d0896bdcdc8d28a7fc5717c305f1a861c26e18c05047949fb371034d98bd/nodeenv-1.10.0-py2.py3-none-any.whl
@@ -3977,6 +3980,7 @@ environments:
       - pypi: https://files.pythonhosted.org/packages/ee/49/1377b49de7d0c1ce41292161ea0f721913fa8722c19fb9c1e3aa0367eecb/pytest_cov-7.0.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/ca/31/d4e37e9e550c2b92a9cbc2e4d0b7420a27224968580b5a447f420847c975/pytest_xdist-3.8.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/ec/57/56b9bcc3c9c6a792fcbaf139543cee77261f3651ca9da0c93f5c1221264b/python_dateutil-2.9.0.post0-py2.py3-none-any.whl
+      - pypi: https://files.pythonhosted.org/packages/d8/db/795879cc3ddfe338599bddea6388cc5100b088db0a4caf6e6c1af1c27e04/python_discovery-1.2.2-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/d1/11/0fd08f8192109f7169db964b5707a2f1e8b745d4e239b784a5a1dd80d1db/pyyaml-6.0.3-cp313-cp313-macosx_10_13_x86_64.whl
       - pypi: https://files.pythonhosted.org/packages/e1/67/921ec3024056483db83953ae8e48079ad62b92db7880013ca77632921dd0/readme_renderer-44.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/2c/58/ca301544e1fa93ed4f80d724bf5b194f6e4b945841c5bfd555878eea9fcb/referencing-0.37.0-py3-none-any.whl
@@ -4014,7 +4018,7 @@ environments:
       - pypi: https://files.pythonhosted.org/packages/39/08/aaaad47bc4e9dc8c725e68f9d04865dbcb2052843ff09c97b08904852d84/urllib3-2.6.3-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/40/db/c5729ec39ed557ac673ca993abd71e5b18d71a92f95f1845152e639000ff/uv-0.10.3-py3-none-macosx_10_12_x86_64.whl
       - pypi: https://files.pythonhosted.org/packages/0e/54/c2fc08d9324bbd92735caf9207cbabada3a8dd9d270d6e46ca69eb7f883d/vhacdx-0.0.10-cp313-cp313-macosx_10_13_x86_64.whl
-      - pypi: https://files.pythonhosted.org/packages/42/4b/6cf85b485be7ec29db837ec2a1d8cd68bc1147b1abf23d8636c5bd65b3cc/virtualenv-20.37.0-py3-none-any.whl
+      - pypi: https://files.pythonhosted.org/packages/20/0e/f083a76cb590e60dff3868779558eefefb8dfb7c9ed020babc7aa014ccbf/virtualenv-21.2.1-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/b9/c4/819f000b3d64f654a054a8b24bf8d0c5e799636393135784d0886d8c808e/viser-1.0.21-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/8a/05/aa116ec9943c718905997412c5989f7ed671bc0188ee2ba89520e8765d7b/websockets-15.0.1-cp313-cp313-macosx_10_13_x86_64.whl
       - pypi: https://files.pythonhosted.org/packages/87/22/b76d483683216dde3d67cba61fb2444be8d5be289bf628c13fc0fd90e5f9/wheel-0.46.3-py3-none-any.whl
@@ -4114,7 +4118,7 @@ environments:
       - pypi: https://files.pythonhosted.org/packages/09/ec/4b43dae793655b7d8a25f76119624350b4d65eb663459eb9603d7f1f0345/mistune-0.8.4-py2.py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/a4/8e/469e5a4a2f5855992e425f3cb33804cc07bf18d48f2db061aec61ce50270/more_itertools-10.8.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/57/b6/eff0305961a1d9447ec2b02f8c73c8946f22564d302a504185b730c9a761/msgspec-0.20.0-cp313-cp313-macosx_11_0_arm64.whl
-      - pypi: https://files.pythonhosted.org/packages/6f/c9/26bd4979d503d03f7a6ded851c3094a5708cb534cf0dc80b4db6672da2b0/mujoco-3.5.0-cp313-cp313-macosx_11_0_arm64.whl
+      - pypi: https://files.pythonhosted.org/packages/25/f2/2cbeabc6b69110e743100f550ec00ae8c60352b6975cf95470add299ed7a/mujoco-3.6.0-cp313-cp313-macosx_11_0_arm64.whl
       - pypi: https://files.pythonhosted.org/packages/9e/c9/b2622292ea83fbb4ec318f5b9ab867d0a28ab43c5717bb85b0a5f6b3b0a4/networkx-3.6.1-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/13/3e/aef8cf8e0419b530c95e96ae93a5078e9b36c1e6613eeb1df03a80d5194e/nh3-0.3.3-cp38-abi3-macosx_10_12_x86_64.macosx_11_0_arm64.macosx_10_12_universal2.whl
       - pypi: https://files.pythonhosted.org/packages/88/b2/d0896bdcdc8d28a7fc5717c305f1a861c26e18c05047949fb371034d98bd/nodeenv-1.10.0-py2.py3-none-any.whl
@@ -4138,6 +4142,7 @@ environments:
       - pypi: https://files.pythonhosted.org/packages/ee/49/1377b49de7d0c1ce41292161ea0f721913fa8722c19fb9c1e3aa0367eecb/pytest_cov-7.0.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/ca/31/d4e37e9e550c2b92a9cbc2e4d0b7420a27224968580b5a447f420847c975/pytest_xdist-3.8.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/ec/57/56b9bcc3c9c6a792fcbaf139543cee77261f3651ca9da0c93f5c1221264b/python_dateutil-2.9.0.post0-py2.py3-none-any.whl
+      - pypi: https://files.pythonhosted.org/packages/d8/db/795879cc3ddfe338599bddea6388cc5100b088db0a4caf6e6c1af1c27e04/python_discovery-1.2.2-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/b1/16/95309993f1d3748cd644e02e38b75d50cbc0d9561d21f390a76242ce073f/pyyaml-6.0.3-cp313-cp313-macosx_11_0_arm64.whl
       - pypi: https://files.pythonhosted.org/packages/e1/67/921ec3024056483db83953ae8e48079ad62b92db7880013ca77632921dd0/readme_renderer-44.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/2c/58/ca301544e1fa93ed4f80d724bf5b194f6e4b945841c5bfd555878eea9fcb/referencing-0.37.0-py3-none-any.whl
@@ -4175,7 +4180,7 @@ environments:
       - pypi: https://files.pythonhosted.org/packages/39/08/aaaad47bc4e9dc8c725e68f9d04865dbcb2052843ff09c97b08904852d84/urllib3-2.6.3-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/0b/8c/cce26f6da9f69a9f0df032d642a4cb5603249f145c368e41b0154c463ca9/uv-0.10.3-py3-none-macosx_11_0_arm64.whl
       - pypi: https://files.pythonhosted.org/packages/3b/9e/42adb642a12915acc9cb2bfab21710a6aabf045c26967ba0ff0e08a872d0/vhacdx-0.0.10-cp313-cp313-macosx_11_0_arm64.whl
-      - pypi: https://files.pythonhosted.org/packages/42/4b/6cf85b485be7ec29db837ec2a1d8cd68bc1147b1abf23d8636c5bd65b3cc/virtualenv-20.37.0-py3-none-any.whl
+      - pypi: https://files.pythonhosted.org/packages/20/0e/f083a76cb590e60dff3868779558eefefb8dfb7c9ed020babc7aa014ccbf/virtualenv-21.2.1-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/b9/c4/819f000b3d64f654a054a8b24bf8d0c5e799636393135784d0886d8c808e/viser-1.0.21-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/ff/0b/33cef55ff24f2d92924923c99926dcce78e7bd922d649467f0eda8368923/websockets-15.0.1-cp313-cp313-macosx_11_0_arm64.whl
       - pypi: https://files.pythonhosted.org/packages/87/22/b76d483683216dde3d67cba61fb2444be8d5be289bf628c13fc0fd90e5f9/wheel-0.46.3-py3-none-any.whl
@@ -4271,7 +4276,7 @@ environments:
       - pypi: https://files.pythonhosted.org/packages/09/ec/4b43dae793655b7d8a25f76119624350b4d65eb663459eb9603d7f1f0345/mistune-0.8.4-py2.py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/a4/8e/469e5a4a2f5855992e425f3cb33804cc07bf18d48f2db061aec61ce50270/more_itertools-10.8.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/f1/25/5e8080fe0117f799b1b68008dc29a65862077296b92550632de015128579/msgspec-0.20.0-cp313-cp313-win_amd64.whl
-      - pypi: https://files.pythonhosted.org/packages/ab/53/54a0815d43c83e1074cfc7da98a3dea88d7dda48c03edfd225a387a3767b/mujoco-3.5.0-cp313-cp313-win_amd64.whl
+      - pypi: https://files.pythonhosted.org/packages/c9/f4/17e3962681c141182616db9ec556ad902311ec154fffda7e9b35ed9677c2/mujoco-3.6.0-cp313-cp313-win_amd64.whl
       - pypi: https://files.pythonhosted.org/packages/9e/c9/b2622292ea83fbb4ec318f5b9ab867d0a28ab43c5717bb85b0a5f6b3b0a4/networkx-3.6.1-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/73/88/1ce287ef8649dc51365b5094bd3713b76454838140a32ab4f8349973883c/nh3-0.3.3-cp38-abi3-win_amd64.whl
       - pypi: https://files.pythonhosted.org/packages/88/b2/d0896bdcdc8d28a7fc5717c305f1a861c26e18c05047949fb371034d98bd/nodeenv-1.10.0-py2.py3-none-any.whl
@@ -4295,6 +4300,7 @@ environments:
       - pypi: https://files.pythonhosted.org/packages/ee/49/1377b49de7d0c1ce41292161ea0f721913fa8722c19fb9c1e3aa0367eecb/pytest_cov-7.0.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/ca/31/d4e37e9e550c2b92a9cbc2e4d0b7420a27224968580b5a447f420847c975/pytest_xdist-3.8.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/ec/57/56b9bcc3c9c6a792fcbaf139543cee77261f3651ca9da0c93f5c1221264b/python_dateutil-2.9.0.post0-py2.py3-none-any.whl
+      - pypi: https://files.pythonhosted.org/packages/d8/db/795879cc3ddfe338599bddea6388cc5100b088db0a4caf6e6c1af1c27e04/python_discovery-1.2.2-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/de/3d/8161f7711c017e01ac9f008dfddd9410dff3674334c233bde66e7ba65bbf/pywin32_ctypes-0.2.3-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/97/c9/39d5b874e8b28845e4ec2202b5da735d0199dbe5b8fb85f91398814a9a46/pyyaml-6.0.3-cp313-cp313-win_amd64.whl
       - pypi: https://files.pythonhosted.org/packages/e1/67/921ec3024056483db83953ae8e48079ad62b92db7880013ca77632921dd0/readme_renderer-44.0-py3-none-any.whl
@@ -4333,7 +4339,7 @@ environments:
       - pypi: https://files.pythonhosted.org/packages/39/08/aaaad47bc4e9dc8c725e68f9d04865dbcb2052843ff09c97b08904852d84/urllib3-2.6.3-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/d5/95/89b9caeb4007d950dfa5fa5923c2ea67de541de0e3582598b7b9c7eeac53/uv-0.10.3-py3-none-win_amd64.whl
       - pypi: https://files.pythonhosted.org/packages/32/25/f0e6806824f88d47ab8bc1c9bf6f11634fd7b382d635d0696825f3b5672f/vhacdx-0.0.10-cp313-cp313-win_amd64.whl
-      - pypi: https://files.pythonhosted.org/packages/42/4b/6cf85b485be7ec29db837ec2a1d8cd68bc1147b1abf23d8636c5bd65b3cc/virtualenv-20.37.0-py3-none-any.whl
+      - pypi: https://files.pythonhosted.org/packages/20/0e/f083a76cb590e60dff3868779558eefefb8dfb7c9ed020babc7aa014ccbf/virtualenv-21.2.1-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/b9/c4/819f000b3d64f654a054a8b24bf8d0c5e799636393135784d0886d8c808e/viser-1.0.21-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/1b/6c/c65773d6cab416a64d191d6ee8a8b1c68a09970ea6909d16965d26bfed1e/websockets-15.0.1-cp313-cp313-win_amd64.whl
       - pypi: https://files.pythonhosted.org/packages/87/22/b76d483683216dde3d67cba61fb2444be8d5be289bf628c13fc0fd90e5f9/wheel-0.46.3-py3-none-any.whl
@@ -5839,10 +5845,10 @@ packages:
 - pypi: ./
   name: judo-rai
   version: 0.0.7
-  sha256: 03bc4bd68a496e5d81ee23ac8989b8d27193ca4fcb1e452796e700d4aca6f9dd
+  sha256: 320dfb5029eaaf28f8e40ea7d8bb9008e9ee598714bbe294b467a643c9873320
   requires_dist:
   - dora-utils
-  - mujoco>=3.5.0,<3.6
+  - mujoco>=3.6.0,<3.7.0
   - numpy
   - pillow
   - pycparser
@@ -5870,6 +5876,7 @@ packages:
   - twine ; extra == 'dev'
   - wheel ; extra == 'dev'
   requires_python: '>=3.10'
+  editable: true
 - conda: https://conda.anaconda.org/conda-forge/noarch/kernel-headers_linux-64-4.18.0-he073ed8_9.conda
   sha256: 41557eeadf641de6aeae49486cef30d02a6912d8da98585d687894afd65b356a
   md5: 86d9cba083cd041bfbf242a01a7a1999
@@ -7389,10 +7396,10 @@ packages:
   - tomli-w ; extra == 'toml'
   - pyyaml ; extra == 'yaml'
   requires_python: '>=3.9'
-- pypi: https://files.pythonhosted.org/packages/02/37/527d83610b878f27c01dd762e0e41aaa62f095c607f0500ac7f724a2c7a5/mujoco-3.5.0-cp312-cp312-win_amd64.whl
+- pypi: https://files.pythonhosted.org/packages/25/f2/2cbeabc6b69110e743100f550ec00ae8c60352b6975cf95470add299ed7a/mujoco-3.6.0-cp313-cp313-macosx_11_0_arm64.whl
   name: mujoco
-  version: 3.5.0
-  sha256: 4b3a62af174ab59b9b6d816dca0786b7fd85ac081d6c2a931a2b22dd6e821f50
+  version: 3.6.0
+  sha256: d593e9373a61db82a506485f7f34533fb6d7e2bff7602f2310aa03e3a93b292f
   requires_dist:
   - absl-py
   - etils[epath]
@@ -7411,11 +7418,11 @@ packages:
   - typing-extensions ; extra == 'sysid'
   - usd-core ; extra == 'usd'
   - pillow ; extra == 'usd'
-  requires_python: '>=3.9'
-- pypi: https://files.pythonhosted.org/packages/16/47/93c7ac3a9630b49c55d76b0d02aa565543e2f62cecd885f8f574f5c745e7/mujoco-3.5.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl
+  requires_python: '>=3.10'
+- pypi: https://files.pythonhosted.org/packages/26/55/7407eced2c44fbea233302d2c11e778852ea0f2eb0e14610f13a7e0d6ac7/mujoco-3.6.0-cp312-cp312-macosx_11_0_arm64.whl
   name: mujoco
-  version: 3.5.0
-  sha256: a956520adb275ce8e878da29e2586eac3affc7b7ac772065ef01f2380a9e8784
+  version: 3.6.0
+  sha256: ea71750f8cbe24b02a091093592f08fb71c95692b43c25e87dabe496ace0bb55
   requires_dist:
   - absl-py
   - etils[epath]
@@ -7434,11 +7441,11 @@ packages:
   - typing-extensions ; extra == 'sysid'
   - usd-core ; extra == 'usd'
   - pillow ; extra == 'usd'
-  requires_python: '>=3.9'
-- pypi: https://files.pythonhosted.org/packages/47/14/d5bf98385354318ec2e6c466a8c7cf7fd76f8b711ed6d11d155e2baa81fb/mujoco-3.5.0-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl
+  requires_python: '>=3.10'
+- pypi: https://files.pythonhosted.org/packages/29/0e/f3ea1fc9d1a25f19b173b13c23797805480cdb0a1026d43cf6b37dc2de6e/mujoco-3.6.0-cp313-cp313-macosx_10_16_x86_64.whl
   name: mujoco
-  version: 3.5.0
-  sha256: ba54826121c6857fc4ca82df642d9a89174ce5537677c6ead34844bb692437e3
+  version: 3.6.0
+  sha256: 29c8c05061798fc3b80269ab3661fa915b890e9623bda4bc6bc9e237db81e885
   requires_dist:
   - absl-py
   - etils[epath]
@@ -7457,11 +7464,11 @@ packages:
   - typing-extensions ; extra == 'sysid'
   - usd-core ; extra == 'usd'
   - pillow ; extra == 'usd'
-  requires_python: '>=3.9'
-- pypi: https://files.pythonhosted.org/packages/6f/c9/26bd4979d503d03f7a6ded851c3094a5708cb534cf0dc80b4db6672da2b0/mujoco-3.5.0-cp313-cp313-macosx_11_0_arm64.whl
+  requires_python: '>=3.10'
+- pypi: https://files.pythonhosted.org/packages/38/c4/f8959e3d5d98b282e081ce08d07cd71ae949cc0ad9f2c39c0a69fcb88c8c/mujoco-3.6.0-cp312-cp312-macosx_10_16_x86_64.whl
   name: mujoco
-  version: 3.5.0
-  sha256: 82416804ae96c69ed779330bd4f4af0a43632e2bbbcc60e5b193642db48e84ca
+  version: 3.6.0
+  sha256: e7e60ee4c07f6fecd63c23e6f47b8d7cdacad75d311739d50d50b5107a630af2
   requires_dist:
   - absl-py
   - etils[epath]
@@ -7480,11 +7487,11 @@ packages:
   - typing-extensions ; extra == 'sysid'
   - usd-core ; extra == 'usd'
   - pillow ; extra == 'usd'
-  requires_python: '>=3.9'
-- pypi: https://files.pythonhosted.org/packages/7e/91/d82dd3c16892e1b0e27a2f537eec8aad54d91d939cb3cd37db2e8c09ecc2/mujoco-3.5.0-cp310-cp310-win_amd64.whl
+  requires_python: '>=3.10'
+- pypi: https://files.pythonhosted.org/packages/3f/a5/a84a9edc6234a2ee29a6b008d432e1c3855795f10a51786ee2260128ed12/mujoco-3.6.0-cp311-cp311-win_amd64.whl
   name: mujoco
-  version: 3.5.0
-  sha256: 2328358d2f0031175897092560dd6d04b14bab1cc22caa145ce99b843c17daa2
+  version: 3.6.0
+  sha256: 99c4b2ec48e988d7ab2dce38e65f237c326954c06323cdd405718034da0b2077
   requires_dist:
   - absl-py
   - etils[epath]
@@ -7503,11 +7510,11 @@ packages:
   - typing-extensions ; extra == 'sysid'
   - usd-core ; extra == 'usd'
   - pillow ; extra == 'usd'
-  requires_python: '>=3.9'
-- pypi: https://files.pythonhosted.org/packages/82/02/aa6057ac4c50fb36558208005d6da19518f9a7857ef9b5fd2ed8f9262fe2/mujoco-3.5.0-cp311-cp311-macosx_11_0_arm64.whl
+  requires_python: '>=3.10'
+- pypi: https://files.pythonhosted.org/packages/52/6c/5ec4e93676a65064a6591176772e00cfa02716156a1d0a7d646a8203348f/mujoco-3.6.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl
   name: mujoco
-  version: 3.5.0
-  sha256: e4fbb00809de98e8a65f2002745c5bca39076f8118b0fe08e973e7a99603c92b
+  version: 3.6.0
+  sha256: 8714fab312c7ee58f45bda7ef8762da2184e3a6a1d780a5093e93a160d66bd3d
   requires_dist:
   - absl-py
   - etils[epath]
@@ -7526,11 +7533,11 @@ packages:
   - typing-extensions ; extra == 'sysid'
   - usd-core ; extra == 'usd'
   - pillow ; extra == 'usd'
-  requires_python: '>=3.9'
-- pypi: https://files.pythonhosted.org/packages/87/2a/371033684e4ddcda47c97661fb6e9617c0e5e3749af082a9b4d5d1bf9f27/mujoco-3.5.0-cp313-cp313-macosx_10_16_x86_64.whl
+  requires_python: '>=3.10'
+- pypi: https://files.pythonhosted.org/packages/6e/48/c8cd52847d8a973fc606910a5467b8b7b68fa763afbe91f41d87123f957c/mujoco-3.6.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl
   name: mujoco
-  version: 3.5.0
-  sha256: 74b05ec4a6a3d728b2da6944d2ae17cac4af9b7a9293f2c2e9e7332fa7535714
+  version: 3.6.0
+  sha256: 5dc7adceab3a7dbf8b4d52176d4aa629aca5f83dfce5ae06abc1a8c93980d67b
   requires_dist:
   - absl-py
   - etils[epath]
@@ -7549,11 +7556,11 @@ packages:
   - typing-extensions ; extra == 'sysid'
   - usd-core ; extra == 'usd'
   - pillow ; extra == 'usd'
-  requires_python: '>=3.9'
-- pypi: https://files.pythonhosted.org/packages/8b/47/e923589301c197c3ea0776b60cc0d57383b3cc51639ca75e4e4b6c5334d6/mujoco-3.5.0-cp311-cp311-macosx_10_16_x86_64.whl
+  requires_python: '>=3.10'
+- pypi: https://files.pythonhosted.org/packages/7f/ed/afb86dfbe53ff4aec0eb265713255899db9947539cd3f235450d82b669f7/mujoco-3.6.0-cp310-cp310-macosx_10_16_x86_64.whl
   name: mujoco
-  version: 3.5.0
-  sha256: 6b3ae97c3f84d093e84dc445a093c893d9f4b6f6bbb1a441e56d77074c450553
+  version: 3.6.0
+  sha256: 8f4d3d2e08472647550f2167a7c72b75e7a749cd9ce62820582d47518b3cabf9
   requires_dist:
   - absl-py
   - etils[epath]
@@ -7572,11 +7579,11 @@ packages:
   - typing-extensions ; extra == 'sysid'
   - usd-core ; extra == 'usd'
   - pillow ; extra == 'usd'
-  requires_python: '>=3.9'
-- pypi: https://files.pythonhosted.org/packages/8d/6b/fdac8ed97086e12ac930fb44e419eda1626e339010df73678cb1f22527d7/mujoco-3.5.0-cp310-cp310-macosx_11_0_arm64.whl
+  requires_python: '>=3.10'
+- pypi: https://files.pythonhosted.org/packages/92/22/38d82f0c34213af53afbbb248b3442943ef48ffbac1e4c909b321e02ac56/mujoco-3.6.0-cp312-cp312-win_amd64.whl
   name: mujoco
-  version: 3.5.0
-  sha256: 5f3803ff0dd7bc04d6c47d53a794343843bde06f0aeefeac28bb62b4cf2baab3
+  version: 3.6.0
+  sha256: 3d4ec53e4e20fcc85843d607fa1648e0b12d2d2de81ee6f85926e95a7e84e8d8
   requires_dist:
   - absl-py
   - etils[epath]
@@ -7595,11 +7602,11 @@ packages:
   - typing-extensions ; extra == 'sysid'
   - usd-core ; extra == 'usd'
   - pillow ; extra == 'usd'
-  requires_python: '>=3.9'
-- pypi: https://files.pythonhosted.org/packages/ab/53/54a0815d43c83e1074cfc7da98a3dea88d7dda48c03edfd225a387a3767b/mujoco-3.5.0-cp313-cp313-win_amd64.whl
+  requires_python: '>=3.10'
+- pypi: https://files.pythonhosted.org/packages/94/27/043852c6d46b9f5b092430e334b2102753f2b63b262f4f466c75c0b9daac/mujoco-3.6.0-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl
   name: mujoco
-  version: 3.5.0
-  sha256: 646b26f545cfdd60ae65ee90d44f63f50fc7ea5b8242777964ef0148830e72df
+  version: 3.6.0
+  sha256: 44b923861cade57044d21cb30f8da78460208cc3f3870f58d83651af35c5f063
   requires_dist:
   - absl-py
   - etils[epath]
@@ -7618,11 +7625,11 @@ packages:
   - typing-extensions ; extra == 'sysid'
   - usd-core ; extra == 'usd'
   - pillow ; extra == 'usd'
-  requires_python: '>=3.9'
-- pypi: https://files.pythonhosted.org/packages/b8/98/c1fac334cc764068e6c5d7eb01d6ed2a3392bab51952c816888b2dfe78c2/mujoco-3.5.0-cp311-cp311-win_amd64.whl
+  requires_python: '>=3.10'
+- pypi: https://files.pythonhosted.org/packages/9b/a2/93ce08c5fcbe04dd997fe207bdc008e856b664ecf69db6442035aacf7fba/mujoco-3.6.0-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl
   name: mujoco
-  version: 3.5.0
-  sha256: ec0e35678773b34ee8b15741c34a745e027db062efcae790315aa83a5581c505
+  version: 3.6.0
+  sha256: 6708a62f4c85bef51c47d0835d29e116da96b4f6a4cf5beef3467dca8af8c407
   requires_dist:
   - absl-py
   - etils[epath]
@@ -7641,11 +7648,11 @@ packages:
   - typing-extensions ; extra == 'sysid'
   - usd-core ; extra == 'usd'
   - pillow ; extra == 'usd'
-  requires_python: '>=3.9'
-- pypi: https://files.pythonhosted.org/packages/c8/8a/229e4db3692be55532e155e2ca6a1363752243ee79df0e7e22ba00f716cf/mujoco-3.5.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl
+  requires_python: '>=3.10'
+- pypi: https://files.pythonhosted.org/packages/bc/6a/b1b3446f3b99bdbc56b4b63dadf4c5cd47ae08e3bf110f36dc5d3fde7836/mujoco-3.6.0-cp310-cp310-macosx_11_0_arm64.whl
   name: mujoco
-  version: 3.5.0
-  sha256: 66fe37276644c28fab497929c55580725de81afc6d511a40cc27525a8dd99efa
+  version: 3.6.0
+  sha256: 69dec85b1a4eeba1be9b6c986dbe6651e993c413a93f072213b175b7d7e19afd
   requires_dist:
   - absl-py
   - etils[epath]
@@ -7664,11 +7671,11 @@ packages:
   - typing-extensions ; extra == 'sysid'
   - usd-core ; extra == 'usd'
   - pillow ; extra == 'usd'
-  requires_python: '>=3.9'
-- pypi: https://files.pythonhosted.org/packages/ca/d6/a5a7b615b257867b7c97db6b3ce07dec9351d5d9d5a5aca881cbb583d7a3/mujoco-3.5.0-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl
+  requires_python: '>=3.10'
+- pypi: https://files.pythonhosted.org/packages/c9/f4/17e3962681c141182616db9ec556ad902311ec154fffda7e9b35ed9677c2/mujoco-3.6.0-cp313-cp313-win_amd64.whl
   name: mujoco
-  version: 3.5.0
-  sha256: 01b12896ae906f157e18d8b1b7c24a8b72d2576fffa09869047150f186e92b33
+  version: 3.6.0
+  sha256: 8068182e134ad8a7786a8d24e3198f485e2c531be1149d7793cfda1cc7fb7122
   requires_dist:
   - absl-py
   - etils[epath]
@@ -7687,11 +7694,11 @@ packages:
   - typing-extensions ; extra == 'sysid'
   - usd-core ; extra == 'usd'
   - pillow ; extra == 'usd'
-  requires_python: '>=3.9'
-- pypi: https://files.pythonhosted.org/packages/e1/d4/d0032323f58a9b8080b8464c6aade8d5ac2e101dbed1de64a38b3913b446/mujoco-3.5.0-cp312-cp312-macosx_11_0_arm64.whl
+  requires_python: '>=3.10'
+- pypi: https://files.pythonhosted.org/packages/d7/fb/5335c0ba2e88f4b8f8300c15966823dbb96ecc906a61c86bcf9cdac77311/mujoco-3.6.0-cp311-cp311-macosx_11_0_arm64.whl
   name: mujoco
-  version: 3.5.0
-  sha256: 94cf4285b46bc2d74fbe86e39a93ecfb3b0e584477fff7e38d293d47b88576e7
+  version: 3.6.0
+  sha256: 5f1a57423e49e6a35ba9cc8335fe023973c11fc29569fee60e14725b384d557d
   requires_dist:
   - absl-py
   - etils[epath]
@@ -7710,11 +7717,11 @@ packages:
   - typing-extensions ; extra == 'sysid'
   - usd-core ; extra == 'usd'
   - pillow ; extra == 'usd'
-  requires_python: '>=3.9'
-- pypi: https://files.pythonhosted.org/packages/e8/20/9e0595e653543df3e4233bc3ad7e50b371b81dbe48d45ffbc867ed7c379d/mujoco-3.5.0-cp310-cp310-macosx_10_16_x86_64.whl
+  requires_python: '>=3.10'
+- pypi: https://files.pythonhosted.org/packages/e0/75/d8afb4e98b58a119be3cba8da88b75cf53ff16f83baa9a14d37aa15a426e/mujoco-3.6.0-cp311-cp311-macosx_10_16_x86_64.whl
   name: mujoco
-  version: 3.5.0
-  sha256: c4324161cb4f334dd984fbb4a4f7d7db9f914f40d06174b02dcf05463d8275e4
+  version: 3.6.0
+  sha256: 5ae08e9249dc04b9da2bb22fe1657277996ad96632f3835cdf3ad60e47beda68
   requires_dist:
   - absl-py
   - etils[epath]
@@ -7733,11 +7740,11 @@ packages:
   - typing-extensions ; extra == 'sysid'
   - usd-core ; extra == 'usd'
   - pillow ; extra == 'usd'
-  requires_python: '>=3.9'
-- pypi: https://files.pythonhosted.org/packages/f9/f0/4772421643f1c5aaf46d9e500a8716f59b02c8bf30bfa92cb8a763159efb/mujoco-3.5.0-cp312-cp312-macosx_10_16_x86_64.whl
+  requires_python: '>=3.10'
+- pypi: https://files.pythonhosted.org/packages/eb/15/ae5df3d9b8fde39f7d75d5143bd9d0a74c1b189b385248667c072136cdea/mujoco-3.6.0-cp310-cp310-win_amd64.whl
   name: mujoco
-  version: 3.5.0
-  sha256: ec0587cc423385a8d45343a981df58511cb69758ba99164a71567af2d41be3c9
+  version: 3.6.0
+  sha256: 3ebaa0f71de1aa3aca4c47990f9b29d65f9829d7dbbaeb05c1a23a9fce176c47
   requires_dist:
   - absl-py
   - etils[epath]
@@ -7756,7 +7763,7 @@ packages:
   - typing-extensions ; extra == 'sysid'
   - usd-core ; extra == 'usd'
   - pillow ; extra == 'usd'
-  requires_python: '>=3.9'
+  requires_python: '>=3.10'
 - conda: https://conda.anaconda.org/conda-forge/linux-64/ncurses-6.5-h2d0b736_3.conda
   sha256: 3fde293232fa3fca98635e1167de6b7c7fda83caf24b9d6c91ec9eefb4f4d586
   md5: 47e340acb35de30501a76c7c799c41d7
@@ -9232,6 +9239,23 @@ packages:
   requires_dist:
   - six>=1.5
   requires_python: '>=2.7,!=3.0.*,!=3.1.*,!=3.2.*'
+- pypi: https://files.pythonhosted.org/packages/d8/db/795879cc3ddfe338599bddea6388cc5100b088db0a4caf6e6c1af1c27e04/python_discovery-1.2.2-py3-none-any.whl
+  name: python-discovery
+  version: 1.2.2
+  sha256: e1ae95d9af875e78f15e19aed0c6137ab1bb49c200f21f5061786490c9585c7a
+  requires_dist:
+  - filelock>=3.15.4
+  - platformdirs>=4.3.6,<5
+  - furo>=2025.12.19 ; extra == 'docs'
+  - sphinx-autodoc-typehints>=3.6.3 ; extra == 'docs'
+  - sphinx>=9.1 ; extra == 'docs'
+  - sphinxcontrib-mermaid>=2 ; extra == 'docs'
+  - covdefaults>=2.3 ; extra == 'testing'
+  - coverage>=7.5.4 ; extra == 'testing'
+  - pytest-mock>=3.14 ; extra == 'testing'
+  - pytest>=8.3.5 ; extra == 'testing'
+  - setuptools>=75.1 ; extra == 'testing'
+  requires_python: '>=3.8'
 - conda: https://conda.anaconda.org/conda-forge/noarch/python_abi-3.13-8_cp313.conda
   build_number: 8
   sha256: 210bffe7b121e651419cb196a2a63687b087497595c9be9d20ebe97dd06060a7
@@ -11125,42 +11149,18 @@ packages:
   - numpy
   - pytest ; extra == 'test'
   requires_python: '>=3.7'
-- pypi: https://files.pythonhosted.org/packages/42/4b/6cf85b485be7ec29db837ec2a1d8cd68bc1147b1abf23d8636c5bd65b3cc/virtualenv-20.37.0-py3-none-any.whl
+- pypi: https://files.pythonhosted.org/packages/20/0e/f083a76cb590e60dff3868779558eefefb8dfb7c9ed020babc7aa014ccbf/virtualenv-21.2.1-py3-none-any.whl
   name: virtualenv
-  version: 20.37.0
-  sha256: 5d3951c32d57232ae3569d4de4cc256c439e045135ebf43518131175d9be435d
+  version: 21.2.1
+  sha256: bd16b49c53562b28cf1a3ad2f36edb805ad71301dee70ddc449e5c88a9f919a2
   requires_dist:
   - distlib>=0.3.7,<1
   - filelock>=3.24.2,<4 ; python_full_version >= '3.10'
   - filelock>=3.16.1,<=3.19.1 ; python_full_version < '3.10'
   - importlib-metadata>=6.6 ; python_full_version < '3.8'
   - platformdirs>=3.9.1,<5
+  - python-discovery>=1
   - typing-extensions>=4.13.2 ; python_full_version < '3.11'
-  - furo>=2023.7.26 ; extra == 'docs'
-  - pre-commit-uv>=4.1.4 ; extra == 'docs'
-  - proselint>=0.13 ; extra == 'docs'
-  - sphinx>=7.1.2,!=7.3 ; extra == 'docs'
-  - sphinx-argparse>=0.4 ; extra == 'docs'
-  - sphinx-autodoc-typehints>=3.6.2 ; extra == 'docs'
-  - sphinx-copybutton>=0.5.2 ; extra == 'docs'
-  - sphinx-inline-tabs>=2025.12.21.14 ; extra == 'docs'
-  - sphinxcontrib-mermaid>=2 ; extra == 'docs'
-  - sphinxcontrib-towncrier>=0.2.1a0 ; extra == 'docs'
-  - towncrier>=23.6 ; extra == 'docs'
-  - covdefaults>=2.3 ; extra == 'test'
-  - coverage-enable-subprocess>=1 ; extra == 'test'
-  - coverage>=7.2.7 ; extra == 'test'
-  - flaky>=3.7 ; extra == 'test'
-  - packaging>=23.1 ; extra == 'test'
-  - pytest-env>=0.8.2 ; extra == 'test'
-  - pytest-freezer>=0.4.8 ; (python_full_version >= '3.13' and platform_python_implementation == 'CPython' and sys_platform == 'win32' and extra == 'test') or (platform_python_implementation == 'GraalVM' and extra == 'test') or (platform_python_implementation == 'PyPy' and extra == 'test')
-  - pytest-mock>=3.11.1 ; extra == 'test'
-  - pytest-randomly>=3.12 ; extra == 'test'
-  - pytest-timeout>=2.1 ; extra == 'test'
-  - pytest-xdist>=3.5 ; extra == 'test'
-  - pytest>=7.4 ; extra == 'test'
-  - setuptools>=68 ; extra == 'test'
-  - time-machine>=2.10 ; platform_python_implementation == 'CPython' and extra == 'test'
   requires_python: '>=3.8'
 - pypi: https://files.pythonhosted.org/packages/b9/c4/819f000b3d64f654a054a8b24bf8d0c5e799636393135784d0886d8c808e/viser-1.0.21-py3-none-any.whl
   name: viser
diff --git a/pyproject.toml b/pyproject.toml
index 72bdb06a..b6f5e3a7 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -30,7 +30,7 @@ classifiers = [
 
 dependencies = [
     "dora-utils",  # dora + qol utils for writing nodes
-    "mujoco>=3.5.0,<3.6",
+    "mujoco>=3.6.0, <3.7.0",
     "numpy",
     "pillow",  # for displaying app logo
     "pycparser",

From 9acaa753131e982a87e5912bcd638a1281e48ca4 Mon Sep 17 00:00:00 2001
From: Brandon Hung <bhung@theaiinstitute.com>
Date: Thu, 16 Apr 2026 16:15:32 -0400
Subject: [PATCH 11/19] Reverting version to 3.5

---
 pixi.lock      | 4 ++--
 pyproject.toml | 2 +-
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/pixi.lock b/pixi.lock
index fbc2092d..a8acfc0b 100644
--- a/pixi.lock
+++ b/pixi.lock
@@ -5845,10 +5845,10 @@ packages:
 - pypi: ./
   name: judo-rai
   version: 0.0.7
-  sha256: 320dfb5029eaaf28f8e40ea7d8bb9008e9ee598714bbe294b467a643c9873320
+  sha256: 1c8c5075097ad38d1acc5f454593067e39d708122350bed5413b95397684d6bf
   requires_dist:
   - dora-utils
-  - mujoco>=3.6.0,<3.7.0
+  - mujoco>=3.5.0
   - numpy
   - pillow
   - pycparser
diff --git a/pyproject.toml b/pyproject.toml
index b6f5e3a7..d91fba07 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -30,7 +30,7 @@ classifiers = [
 
 dependencies = [
     "dora-utils",  # dora + qol utils for writing nodes
-    "mujoco>=3.6.0, <3.7.0",
+    "mujoco>=3.5.0",
     "numpy",
     "pillow",  # for displaying app logo
     "pycparser",

From b326d25ad56497928cfa1fb86a1bbb90bce13653 Mon Sep 17 00:00:00 2001
From: Duy Ta <dta@rai-inst.com>
Date: Tue, 5 May 2026 15:09:19 -0400
Subject: [PATCH 12/19] clean up backend selection mess, rename classes for
 clarity

---
 .../{controller.py => controller_node.py}     |  41 ++++--
 .../{simulation.py => simulation_node.py}     |  43 ++++---
 ...visualization.py => visualization_node.py} |  21 +++-
 judo/app/utils.py                             |  23 +++-
 judo/cli.py                                   |  13 ++
 judo/configs/judo_dora_default.yaml           |   8 +-
 judo/controller/controller.py                 | 117 +++++++++++++-----
 judo/simulation/__init__.py                   |  21 ++--
 judo/simulation/base.py                       |   3 +-
 ...ation.py => hierarchical_mj_simulation.py} |  59 +++++----
 judo/tasks/__init__.py                        | 100 ++++++++++++---
 judo/tasks/base.py                            |  11 +-
 judo/tasks/spot/spot_base.py                  |   7 +-
 ....py => hierarchical_mj_rollout_backend.py} |  16 +--
 judo/visualizers/visualizer.py                |   7 +-
 tests/test_simulation/test_simulation.py      |  14 ++-
 tests/test_tasks/test_spot.py                 |   7 +-
 17 files changed, 342 insertions(+), 169 deletions(-)
 rename judo/app/dora/{controller.py => controller_node.py} (79%)
 rename judo/app/dora/{simulation.py => simulation_node.py} (75%)
 rename judo/app/dora/{visualization.py => visualization_node.py} (84%)
 rename judo/simulation/{policy_mj_simulation.py => hierarchical_mj_simulation.py} (65%)
 rename judo/utils/{policy_mj_rollout_backend.py => hierarchical_mj_rollout_backend.py} (88%)

diff --git a/judo/app/dora/controller.py b/judo/app/dora/controller_node.py
similarity index 79%
rename from judo/app/dora/controller.py
rename to judo/app/dora/controller_node.py
index cda4d40b..412a9691 100644
--- a/judo/app/dora/controller.py
+++ b/judo/app/dora/controller_node.py
@@ -11,6 +11,7 @@
 
 from judo.app.structs import MujocoState
 from judo.controller import Controller, make_controller
+from judo.tasks import get_task_registration
 
 
 class ControllerNode(DoraNode):
@@ -24,23 +25,40 @@ def __init__(
         max_workers: int | None = None,
         task_registration_cfg: DictConfig | None = None,
         optimizer_registration_cfg: DictConfig | None = None,
-        controller_cls: type[Controller] | None = None,
         make_controller_fn: Callable | None = None,
     ) -> None:
-        """Initialize the controller node."""
+        """Initialize the controller node.
+
+        Args:
+            init_task: Name of the task to initialize.
+            init_optimizer: Name of the optimizer to initialize (e.g., "cem", "ps", "mppi").
+            node_id: Identifier for this dora node.
+            max_workers: Maximum number of worker threads for dora (None = auto).
+            task_registration_cfg: Optional config for task registration overrides.
+            optimizer_registration_cfg: Optional config for optimizer registration overrides.
+            make_controller_fn: Optional factory function to create Controller instances.
+                Defaults to judo.controller.make_controller. Allows custom controller creation.
+        """
         super().__init__(node_id=node_id, max_workers=max_workers)
-        self._controller_cls = controller_cls or Controller
         self._make_controller_fn = make_controller_fn or make_controller
-        self.controller = self._make_controller_fn(
-            init_task=init_task,
-            init_optimizer=init_optimizer,
-            task_registration_cfg=task_registration_cfg,
-            optimizer_registration_cfg=optimizer_registration_cfg,
-        )
+        self._task_registration_cfg = task_registration_cfg
+        self._optimizer_registration_cfg = optimizer_registration_cfg
+        self.controller = self._build_controller(init_task, init_optimizer)
         self._paused = False
         self.write_controls()
         self.lock = Lock()
 
+    def _build_controller(self, task_name: str, optimizer_name: str) -> Controller:
+        """Build controller using the task's registered rollout backend."""
+        rollout_backend = get_task_registration(task_name).rollout_backend
+        return self._make_controller_fn(
+            init_task=task_name,
+            init_optimizer=optimizer_name,
+            task_registration_cfg=self._task_registration_cfg,
+            optimizer_registration_cfg=self._optimizer_registration_cfg,
+            rollout_backend=rollout_backend,
+        )
+
     def _current_optimizer_name(self) -> str:
         """Look up the name of the current optimizer from the registry."""
         for name, (cls, _) in self.controller.available_optimizers.items():
@@ -57,10 +75,7 @@ def update_task(self, event: dict) -> None:
             raise ValueError(f"Task {new_task} not found in task registry.")
 
         with self.lock:
-            self.controller = self._make_controller_fn(
-                init_task=new_task,
-                init_optimizer=self._current_optimizer_name(),
-            )
+            self.controller = self._build_controller(new_task, self._current_optimizer_name())
             self.write_controls()
 
     @on_event("INPUT", "task_reset")
diff --git a/judo/app/dora/simulation.py b/judo/app/dora/simulation_node.py
similarity index 75%
rename from judo/app/dora/simulation.py
rename to judo/app/dora/simulation_node.py
index 7d4849e0..79059c7a 100644
--- a/judo/app/dora/simulation.py
+++ b/judo/app/dora/simulation_node.py
@@ -9,6 +9,7 @@
 from omegaconf import DictConfig
 
 from judo.app.structs import SplineData
+from judo.simulation import DEFAULT_SIMULATION_BACKEND_REGISTRY
 from judo.simulation.base import Simulation
 from judo.tasks import get_registered_tasks
 
@@ -22,43 +23,41 @@ def __init__(
         init_task: str = "cylinder_push",
         max_workers: int | None = None,
         task_registration_cfg: DictConfig | None = None,
-        simulation_backend: str = "mujoco",
-        custom_backends: dict[str, type[Simulation]] | None = None,
+        backend_registry: dict[str, type[Simulation]] | None = None,
     ) -> None:
-        """Initialize the simulation node."""
+        """Initialize the simulation node.
+
+        Args:
+            node_id: Identifier for this dora node.
+            init_task: Name of the task to initialize.
+            max_workers: Maximum number of worker threads for dora (None = auto).
+            task_registration_cfg: Optional config for task registration overrides.
+            backend_registry: Optional mapping of backend names → Simulation classes. Checked first before built-in registry.
+        """
         super().__init__(node_id=node_id, max_workers=max_workers)
-        self._simulation_backend = simulation_backend
         self._task_registration_cfg = task_registration_cfg
-        self._custom_backends = custom_backends or {}
+        self._backend_registry = dict(DEFAULT_SIMULATION_BACKEND_REGISTRY)
+        self._backend_registry.update(backend_registry or {})
         self._init_sim(init_task)
         self.control_spline: Callable | None = None
         self.write_states()
 
-    def _resolve_backend(self, backend: str) -> type:
-        """Resolve a simulation backend class by name, checking custom backends first."""
-        if backend in self._custom_backends:
-            return self._custom_backends[backend]
-        from judo.simulation import get_simulation_backend  # noqa: PLC0415, I001
-
-        return get_simulation_backend(backend)
+    def _resolve_backend(self, backend_name: str) -> type[Simulation]:
+        """Resolve a simulation backend class by name from merged registry."""
+        backend_cls = self._backend_registry.get(backend_name)
+        if backend_cls is None:
+            raise KeyError(f"Unknown simulation backend: {backend_name!r}")
+        return backend_cls
 
     def _init_sim(self, task_name: str) -> None:
-        """Initialize simulation, auto-upgrading to policy backend if needed."""
+        """Initialize simulation using the task's registered simulation backend."""
         task_entry = get_registered_tasks().get(task_name)
         if task_entry is None:
             raise ValueError(f"Task {task_name} not found in task registry.")
 
-        task_cls, _ = task_entry
-        backend = getattr(task_cls, "default_backend", None) or self._simulation_backend
-
-        sim_backend_cls = self._resolve_backend(backend)
+        sim_backend_cls = self._resolve_backend(task_entry.simulation_backend)
         self.sim = sim_backend_cls(init_task=task_name, task_registration_cfg=self._task_registration_cfg)
 
-        # Auto-upgrade to policy backend if task requires locomotion policy
-        if backend == "mujoco" and self.sim.task.uses_locomotion_policy:
-            sim_backend_cls = self._resolve_backend("mujoco_policy")
-            self.sim = sim_backend_cls(init_task=task_name, task_registration_cfg=self._task_registration_cfg)
-
     @on_event("INPUT", "task")
     def update_task(self, event: dict) -> None:
         """Event handler for processing task updates."""
diff --git a/judo/app/dora/visualization.py b/judo/app/dora/visualization_node.py
similarity index 84%
rename from judo/app/dora/visualization.py
rename to judo/app/dora/visualization_node.py
index 08d8a2c1..a68fa56a 100644
--- a/judo/app/dora/visualization.py
+++ b/judo/app/dora/visualization_node.py
@@ -9,7 +9,7 @@
 from viser import GuiFolderHandle, GuiImageHandle, GuiInputHandle, IcosphereHandle, MeshHandle
 
 from judo.app.structs import WorldState
-from judo.tasks import Task, TaskConfig
+from judo.tasks import TaskRegistration
 from judo.visualizers.visualizer import Visualizer
 
 ElementType = GuiImageHandle | GuiInputHandle | GuiFolderHandle | MeshHandle | IcosphereHandle
@@ -30,9 +30,24 @@ def __init__(
         optimizer_override_cfg: DictConfig | None = None,
         sim_pause_button: bool = True,
         geom_exclude_substring: str = "collision",
-        available_tasks: dict[str, tuple[type[Task], type[TaskConfig]]] | None = None,
+        available_tasks: dict[str, TaskRegistration] | None = None,
     ) -> None:
-        """Initialize the visualization node."""
+        """Initialize the visualization node (Viser web GUI for task/optimizer control).
+
+        Args:
+            node_id: Identifier for this dora node.
+            max_workers: Maximum number of worker threads for dora (None = auto).
+            init_task: Name of the task to initialize.
+            init_optimizer: Name of the optimizer to initialize (e.g., "cem", "ps").
+            task_registration_cfg: Optional config for task registration overrides.
+            optimizer_registration_cfg: Optional config for optimizer registration overrides.
+            controller_override_cfg: Optional config overrides for the controller.
+            optimizer_override_cfg: Optional config overrides for the optimizer.
+            sim_pause_button: Whether to display a simulation pause button in the GUI.
+            geom_exclude_substring: Geometry name substring to exclude from visualization (default "collision" hides collision shapes).
+            available_tasks: Optional pre-computed mapping of task names to TaskRegistration entries
+                for the task selector. If None, tasks are inferred from the task registry.
+        """
         super().__init__(node_id=node_id, max_workers=max_workers)
         self.visualizer = Visualizer(
             init_task=init_task,
diff --git a/judo/app/utils.py b/judo/app/utils.py
index 09aa632f..ed2c9a53 100644
--- a/judo/app/utils.py
+++ b/judo/app/utils.py
@@ -20,14 +20,31 @@ def register_tasks_from_cfg(task_registration_cfg: DictConfig) -> None:
     """Register custom tasks."""
     for task_name in task_registration_cfg.keys():
         task_dict = task_registration_cfg.get(task_name, {})
-        assert set(task_dict.keys()) == {"task", "config"}, (
-            "Task registration must be a dict with keys 'task' and 'config'."
+        allowed_keys = {"task", "config", "rollout_backend", "simulation_backend", "locomotion_policy_path"}
+        assert set(task_dict.keys()).issubset(allowed_keys) and {"task", "config"}.issubset(task_dict.keys()), (
+            "Task registration must include 'task' and 'config', and may optionally include "
+            "'rollout_backend', 'simulation_backend', and 'locomotion_policy_path'."
         )
         assert isinstance(task_dict["task"], str), "Task must be a string path to the task class."
         assert isinstance(task_dict["config"], str), "Task config must be a string path to the config class."
         task_cls = get_class_from_string(task_dict["task"])
         task_config_cls = get_class_from_string(task_dict["config"])
-        register_task(str(task_name), task_cls, task_config_cls)
+        rollout_backend = task_dict.get("rollout_backend", "mujoco")
+        simulation_backend = task_dict.get("simulation_backend", "mujoco")
+        locomotion_policy_path = task_dict.get("locomotion_policy_path", None)
+        assert isinstance(rollout_backend, str), "rollout_backend must be a string."
+        assert isinstance(simulation_backend, str), "simulation_backend must be a string."
+        assert locomotion_policy_path is None or isinstance(
+            locomotion_policy_path, str
+        ), "locomotion_policy_path must be a string if provided."
+        register_task(
+            str(task_name),
+            task_cls,
+            task_config_cls,
+            rollout_backend=rollout_backend,
+            simulation_backend=simulation_backend,
+            locomotion_policy_path=locomotion_policy_path,
+        )
 
 
 def register_optimizers_from_cfg(optimizer_registration_cfg: DictConfig) -> None:
diff --git a/judo/cli.py b/judo/cli.py
index f5cabf69..42e49541 100644
--- a/judo/cli.py
+++ b/judo/cli.py
@@ -141,8 +141,20 @@ def _warm_caches() -> None:
         pass  # non-Spot tasks don't need this
 
 
+def _require_mujoco_extensions() -> None:
+    """Fail fast if mujoco_extensions is unavailable in the current environment."""
+    try:
+        import mujoco_extensions  # noqa: F401, PLC0415
+    except Exception as e:  # pragma: no cover - environment dependent
+        raise RuntimeError(
+            "mujoco_extensions is required but could not be imported. "
+            "Build it with: pixi run build"
+        ) from e
+
+
 def app() -> None:
     """Entry point for the judo CLI."""
+    _require_mujoco_extensions()
     _warm_caches()
     # we store judo_dora_default in the config store so that custom dora configs outside of judo can inherit from it
     cs = ConfigStore.instance()
@@ -168,6 +180,7 @@ def main_benchmark(cfg: DictConfig) -> None:
 
 def benchmark() -> None:
     """Entry point for benchmarking."""
+    _require_mujoco_extensions()
     # we store benchmark_default in the config store so that custom configs located outside of judo can inherit from it
     cs = ConfigStore.instance()
     with initialize_config_dir(config_dir=str(CONFIG_PATH), version_base="1.3"):
diff --git a/judo/configs/judo_dora_default.yaml b/judo/configs/judo_dora_default.yaml
index 6487039d..79c59b47 100644
--- a/judo/configs/judo_dora_default.yaml
+++ b/judo/configs/judo_dora_default.yaml
@@ -75,14 +75,13 @@ dataflow:
 
 node_definitions:
   simulation:
-    _target_: judo.app.dora.simulation.SimulationNode
+    _target_: judo.app.dora.simulation_node.SimulationNode
     node_id: simulation
     max_workers: null
     init_task: ${task}
     task_registration_cfg: ${custom_tasks}
-    simulation_backend: ${simulation_backend}
   visualization:
-    _target_: judo.app.dora.visualization.VisualizationNode
+    _target_: judo.app.dora.visualization_node.VisualizationNode
     node_id: visualization
     max_workers: null
     init_task: ${task}
@@ -93,7 +92,7 @@ node_definitions:
     optimizer_override_cfg: ${optimizer_config_overrides}
     sim_pause_button: true
   controller:
-    _target_: judo.app.dora.controller.ControllerNode
+    _target_: judo.app.dora.controller_node.ControllerNode
     node_id: controller
     max_workers: null
     init_task: ${task}
@@ -105,4 +104,3 @@ custom_tasks: null
 custom_optimizers: null
 controller_config_overrides: null
 optimizer_config_overrides: null
-simulation_backend: mujoco
diff --git a/judo/controller/controller.py b/judo/controller/controller.py
index c57ff32f..4bcde6c7 100644
--- a/judo/controller/controller.py
+++ b/judo/controller/controller.py
@@ -13,8 +13,9 @@
 from judo.config import OverridableConfig
 from judo.gui import slider
 from judo.optimizers import Optimizer, OptimizerConfig, get_registered_optimizers
-from judo.tasks import Task, TaskConfig, get_registered_tasks
+from judo.tasks import Task, TaskConfig, get_registered_tasks, get_task_registration
 from judo.tasks.spot.spot_constants import POLICY_OUTPUT_DIM
+from judo.utils.hierarchical_mj_rollout_backend import HierarchicalMJRolloutBackend
 from judo.utils.mj_rollout_backend import MJRolloutBackend
 from judo.utils.normalization import (
     IdentityNormalizer,
@@ -23,11 +24,19 @@
     make_normalizer,
     normalizer_registry,
 )
-from judo.utils.policy_mj_rollout_backend import PolicyMJRolloutBackend
 from judo.utils.rollout_backend import RolloutBackend
 from judo.visualizers.utils import get_trace_sensors
 
 
+RolloutBackendEntry = type[RolloutBackend]
+
+
+DEFAULT_ROLLOUT_BACKEND_REGISTRY: dict[str, RolloutBackendEntry] = {
+    "mujoco": MJRolloutBackend,
+    "mujoco_hierarchical": HierarchicalMJRolloutBackend,
+}
+
+
 @slider("horizon", 0.1, 10.0, bounded=True)
 @slider("control_freq", 0.25, 50.0)
 @dataclass
@@ -50,8 +59,9 @@ def __init__(
         controller_config: ControllerConfig,
         task: Task,
         optimizer: Optimizer,
-        rollout_backend: Literal["mujoco"] = "mujoco",
-        custom_rollout_backends: dict[str, type] | None = None,
+        rollout_backend: str = "mujoco",
+        rollout_backend_registry: dict[str, RolloutBackendEntry] | None = None,
+        rollout_backend_kwargs: dict[str, Any] | None = None,
     ) -> None:
         """Initialize the controller.
 
@@ -59,15 +69,21 @@ def __init__(
             controller_config: The controller configuration.
             task: The task to use.
             optimizer: The optimizer to use.
-            rollout_backend: The backend to use for rollouts. Currently only "mujoco" is supported.
-            custom_rollout_backends: Optional mapping of backend names to backend classes.
-                If the task's ``default_backend`` matches a key, that class is instantiated
-                with ``model`` and ``num_threads`` keyword arguments.
+            rollout_backend: Name of the backend to use for rollouts (e.g., "mujoco", "mujoco_hierarchical").
+            rollout_backend_registry: Optional mapping of backend names to backend classes.
+                Overrides entries in DEFAULT_ROLLOUT_BACKEND_REGISTRY.
+            rollout_backend_kwargs: Optional extra kwargs for rollout backend constructor.
+                For "mujoco_hierarchical" backend, 'physics_substeps' and 'policy_path' cannot be
+                specified here—they are sourced from the task and task registry respectively.
+                Raises ValueError if either is provided. To use different values, create or
+                update a task registry entry.
         """
         self._controller_cfg = controller_config
         self.task = task
         self.optimizer = optimizer
-        self._custom_rollout_backends = custom_rollout_backends or {}
+        self._rollout_backend_registry = dict(DEFAULT_ROLLOUT_BACKEND_REGISTRY)
+        self._rollout_backend_registry.update(rollout_backend_registry or {})
+        self._rollout_backend_kwargs = rollout_backend_kwargs or {}
 
         self.available_optimizers = get_registered_optimizers()
         self.available_tasks = get_registered_tasks()
@@ -75,27 +91,14 @@ def __init__(
         self.model = self.task.model
 
         # Initialize rollout backend
-        default_backend = getattr(self.task, "default_backend", None)
-        if default_backend and default_backend in self._custom_rollout_backends:
-            self.rollout_backend: RolloutBackend = self._custom_rollout_backends[default_backend](
-                model=self.model,
-                num_threads=self.optimizer_cfg.num_rollouts,
-            )
-        elif self.task.uses_locomotion_policy:
-            assert self.task.locomotion_policy_path is not None
-            self.rollout_backend = PolicyMJRolloutBackend(
-                model=self.model,
-                num_threads=self.optimizer_cfg.num_rollouts,
-                policy_path=self.task.locomotion_policy_path,
-                physics_substeps=self.task.physics_substeps,
-            )
-        else:
-            self.rollout_backend = MJRolloutBackend(
-                model=self.model,
-                num_threads=self.optimizer_cfg.num_rollouts,
-            )
+        self.rollout_backend: RolloutBackend = self._make_rollout_backend(
+            rollout_backend,
+            backend_kwargs=self._rollout_backend_kwargs,
+        )
         self._last_policy_output = (
-            np.zeros((self.optimizer_cfg.num_rollouts, POLICY_OUTPUT_DIM)) if self.task.uses_locomotion_policy else None
+            np.zeros((self.optimizer_cfg.num_rollouts, POLICY_OUTPUT_DIM))
+            if isinstance(self.rollout_backend, HierarchicalMJRolloutBackend)
+            else None
         )
         self.action_normalizer = self._init_action_normalizer()
 
@@ -327,8 +330,8 @@ def reset(self) -> None:
         self.candidate_knots = np.tile(self.nominal_knots, (self.optimizer_cfg.num_rollouts, 1, 1))
         self.times = self.task.data.time + self.spline_timesteps
         self.update_spline(self.times, self.nominal_knots)
-        # Reset policy output state for locomotion policy tasks
-        if self.task.uses_locomotion_policy:
+        # Reset policy output state for policy rollout backends
+        if isinstance(self.rollout_backend, HierarchicalMJRolloutBackend):
             self._last_policy_output = np.zeros((self.optimizer_cfg.num_rollouts, POLICY_OUTPUT_DIM))
 
     def update_traces(self) -> None:
@@ -389,6 +392,53 @@ def _init_action_normalizer(self) -> Normalizer:
             action_normalizer_kwargs["init_std"] = 1.0  # TODO(yunhai): make this configurable
         return make_normalizer(self.action_normalizer_type, self.task.nu, **action_normalizer_kwargs)
 
+    def _make_rollout_backend(
+        self,
+        backend_name: str,
+        backend_kwargs: dict[str, Any] | None = None,
+    ) -> RolloutBackend:
+        """Instantiate a rollout backend from the merged backend registry."""
+        backend_cls = self._rollout_backend_registry.get(backend_name)
+        if backend_cls is None:
+            raise ValueError(
+                f"Unknown rollout backend '{backend_name}'. "
+                "Provide it via rollout_backend_registry or choose a built-in backend."
+            )
+
+        final_kwargs = {
+            "model": self.model,
+            "num_threads": self.optimizer_cfg.num_rollouts,
+        }
+        final_kwargs.update(backend_kwargs or {})
+
+        if backend_name == "mujoco_hierarchical":
+            # physics_substeps must come from task, cannot be overridden in kwargs
+            if "physics_substeps" in final_kwargs:
+                raise ValueError(
+                    f"Cannot specify 'physics_substeps' in rollout_backend_kwargs. "
+                    f"It is determined by the task configuration (task.physics_substeps). "
+                    f"Current task '{self.task.name}' has physics_substeps={self.task.physics_substeps}."
+                )
+            final_kwargs["physics_substeps"] = self.task.physics_substeps
+
+            # policy_path must come from task registry, cannot be overridden in kwargs
+            if "policy_path" in final_kwargs:
+                raise ValueError(
+                    f"Cannot specify 'policy_path' in rollout_backend_kwargs. "
+                    f"It must be defined in the task registry entry for '{self.task.name}'. "
+                    f"To use a different policy path, create or update a task registry entry with the desired path."
+                )
+
+            task_policy_path = get_task_registration(self.task.name).locomotion_policy_path
+            if task_policy_path is None:
+                raise ValueError(
+                    f"Backend '{backend_name}' requires 'policy_path'. "
+                    f"Task '{self.task.name}' must have a locomotion_policy_path registered in the task registry."
+                )
+            final_kwargs["policy_path"] = task_policy_path
+
+        return backend_cls(**final_kwargs)
+
 
 def make_spline(times: np.ndarray, controls: np.ndarray, spline_order: str) -> interp1d:
     """Helper function for creating spline objects.
@@ -417,7 +467,7 @@ def make_controller(
     init_optimizer: str,
     task_registration_cfg: DictConfig | None = None,
     optimizer_registration_cfg: DictConfig | None = None,
-    rollout_backend: Literal["mujoco"] = "mujoco",
+    rollout_backend: str = "mujoco",
     controller_cls: type[Controller] | None = None,
     **controller_kwargs: Any,
 ) -> Controller:
@@ -436,8 +486,7 @@ def make_controller(
     assert optimizer_entry is not None, f"Optimizer {init_optimizer} not found in optimizer registry."
 
     # instantiate the task/optimizer/controller
-    task_cls, _ = task_entry
-    task = task_cls()
+    task = task_entry.task_type()
 
     optimizer_cls, optimizer_config_cls = optimizer_entry
     optimizer_cfg = optimizer_config_cls()
diff --git a/judo/simulation/__init__.py b/judo/simulation/__init__.py
index fc522e94..f6d2f043 100644
--- a/judo/simulation/__init__.py
+++ b/judo/simulation/__init__.py
@@ -1,19 +1,13 @@
 # Copyright (c) 2025 Robotics and AI Institute LLC. All rights reserved.
 
 from judo.simulation.base import Simulation
+from judo.simulation.hierarchical_mj_simulation import HierarchicalMJSimulation
 from judo.simulation.mj_simulation import MJSimulation
-from judo.simulation.policy_mj_simulation import PolicyMJSimulation
 
 
-def _get_policy_mj_simulation() -> type[MJSimulation]:
-    from judo.simulation.policy_mj_simulation import PolicyMJSimulation  # noqa: PLC0415, I001
-
-    return PolicyMJSimulation
-
-
-_simulation_registry = {
-    "mujoco": lambda: MJSimulation,
-    "mujoco_policy": _get_policy_mj_simulation,
+DEFAULT_SIMULATION_BACKEND_REGISTRY: dict[str, type[Simulation]] = {
+    "mujoco": MJSimulation,
+    "mujoco_hierarchical": HierarchicalMJSimulation,
 }
 
 
@@ -26,14 +20,15 @@ def get_simulation_backend(simulation_backend: str) -> type:
     Returns:
         The simulation class for the given backend.
     """
-    if simulation_backend not in _simulation_registry:
+    if simulation_backend not in DEFAULT_SIMULATION_BACKEND_REGISTRY:
         raise KeyError(f"Unknown simulation backend: {simulation_backend!r}")
-    return _simulation_registry[simulation_backend]()
+    return DEFAULT_SIMULATION_BACKEND_REGISTRY[simulation_backend]
 
 
 __all__ = [
     "Simulation",
     "MJSimulation",
-    "PolicyMJSimulation",
+    "HierarchicalMJSimulation",
+    "DEFAULT_SIMULATION_BACKEND_REGISTRY",
     "get_simulation_backend",
 ]
diff --git a/judo/simulation/base.py b/judo/simulation/base.py
index 9575fad9..e76ac743 100644
--- a/judo/simulation/base.py
+++ b/judo/simulation/base.py
@@ -38,8 +38,7 @@ def set_task(self, task_name: str) -> None:
         if task_entry is None:
             raise ValueError(f"Task {task_name} not found in task registry")
 
-        task_cls, _ = task_entry
-        self.task: Task = task_cls()
+        self.task: Task = task_entry.task_type()
         self.task.reset()
 
     @abstractmethod
diff --git a/judo/simulation/policy_mj_simulation.py b/judo/simulation/hierarchical_mj_simulation.py
similarity index 65%
rename from judo/simulation/policy_mj_simulation.py
rename to judo/simulation/hierarchical_mj_simulation.py
index ad60b39f..e72c160a 100644
--- a/judo/simulation/policy_mj_simulation.py
+++ b/judo/simulation/hierarchical_mj_simulation.py
@@ -1,6 +1,6 @@
 # Copyright (c) 2025 Robotics and AI Institute LLC. All rights reserved.
 
-"""MuJoCo Simulation with locomotion policy support."""
+"""MuJoCo simulation with hierarchical low-level policy support."""
 
 from pathlib import Path
 
@@ -9,6 +9,7 @@
 from omegaconf import DictConfig
 
 from judo.simulation.mj_simulation import MJSimulation
+from judo.tasks import get_task_registration
 from judo.tasks.spot.spot_constants import DEFAULT_SPOT_ROLLOUT_CUTOFF_TIME, POLICY_OUTPUT_DIM
 
 try:
@@ -21,14 +22,15 @@
     ) from e
 
 
-class PolicyMJSimulation(MJSimulation):
-    """MuJoCo simulation with locomotion policy support.
+class HierarchicalMJSimulation(MJSimulation):
+    """MuJoCo simulation with a hierarchical low-level policy layer.
 
-    For tasks with locomotion_policy_path set, uses C++ mujoco_extensions
-    threaded_rollout to run the neural network policy at 50Hz.
+    For tasks with uses_locomotion_policy=True, this routes control through
+    mujoco_extensions threaded_rollout so a lower-level policy can refine the
+    high-level command before physics integration.
 
-    The simulation maintains internal state for the locomotion policy
-    (last_policy_output) to ensure smooth transitions between timesteps.
+    The simulation maintains internal policy state (last_policy_output) to
+    ensure smooth transitions between timesteps.
     """
 
     def __init__(
@@ -36,7 +38,7 @@ def __init__(
         init_task: str = "spot_base",
         task_registration_cfg: DictConfig | None = None,
     ) -> None:
-        """Initialize the policy simulation.
+        """Initialize the hierarchical simulation.
 
         Args:
             init_task: Name of the task to initialize.
@@ -47,15 +49,20 @@ def __init__(
         self._systems = None
         self._last_policy_output = np.zeros(POLICY_OUTPUT_DIM)
 
-        # Initialize C++ systems if task uses locomotion policy
-        if self.task.locomotion_policy_path is not None:
-            self._init_cpp_systems(self.task.locomotion_policy_path)
+        # Initialize C++ systems if the task uses a hierarchical policy layer.
+        if self.task.uses_locomotion_policy:
+            policy_path = get_task_registration(self.task.name).locomotion_policy_path
+            if policy_path is None:
+                raise ValueError(
+                    f"Task '{self.task.name}' uses locomotion policy but no locomotion_policy_path is registered."
+                )
+            self._init_cpp_systems(policy_path)
 
     def _init_cpp_systems(self, policy_path: str | Path) -> None:
         """Initialize the C++ systems vector for threaded rollout.
 
         Args:
-            policy_path: Path to the ONNX locomotion policy file.
+            policy_path: Path to the ONNX low-level policy file.
         """
         self._systems = create_systems_vector(
             self.task.model,  # Pass the MjModel directly
@@ -66,26 +73,27 @@ def _init_cpp_systems(self, policy_path: str | Path) -> None:
     def step(self, command: np.ndarray) -> None:
         """Step the simulation forward.
 
-        Routes to the C++ policy rollout if systems are initialized,
+        Routes to the C++ hierarchical rollout if systems are initialized,
         otherwise falls back to direct actuator control.
 
         Args:
             command: Control array in task format (task.nu dimensions).
-                For locomotion tasks, will be converted to policy command internally.
+                For hierarchical tasks, this is converted to the low-level
+                policy command internally.
         """
         if self._systems is not None:
             if self.paused:
                 return
             command = self.task.task_to_sim_ctrl(command)
-            self._step_with_locomotion_policy(command)
+            self._step_with_hierarchical_policy(command)
         else:
             super().step(command)
 
-    def _step_with_locomotion_policy(self, command: np.ndarray) -> None:
-        """Execute a single step using the C++ rollout backend.
+    def _step_with_hierarchical_policy(self, command: np.ndarray) -> None:
+        """Execute a single step using the hierarchical rollout backend.
 
         Args:
-            command: Command array for the locomotion policy.
+            command: Command array for the low-level policy.
         """
         # Get current state
         state = np.concatenate([self.task.data.qpos, self.task.data.qvel])
@@ -124,7 +132,7 @@ def _step_with_locomotion_policy(self, command: np.ndarray) -> None:
         # Compute derived quantities (xpos, xquat, etc.) for visualization
         mj_forward(self.task.model, self.task.data)
 
-        # Update last policy output for continuity
+        # Update last policy output for continuity.
         self._last_policy_output = np.array(policy_outputs[0])
 
     def reset_policy_state(self) -> None:
@@ -139,14 +147,19 @@ def set_task(self, task_name: str) -> None:
         """
         super().set_task(task_name)
 
-        # Reinitialize systems based on new task's policy
-        if self.task.locomotion_policy_path is not None:
-            self._init_cpp_systems(self.task.locomotion_policy_path)
+        # Reinitialize systems based on the new task's policy layer.
+        if self.task.uses_locomotion_policy:
+            policy_path = get_task_registration(self.task.name).locomotion_policy_path
+            if policy_path is None:
+                raise ValueError(
+                    f"Task '{self.task.name}' uses locomotion policy but no locomotion_policy_path is registered."
+                )
+            self._init_cpp_systems(policy_path)
             self._last_policy_output = np.zeros(POLICY_OUTPUT_DIM)
         else:
             self._systems = None
 
     @property
     def last_policy_output(self) -> np.ndarray:
-        """Returns the last policy output (12-dim leg actions)."""
+        """Return the last low-level policy output."""
         return self._last_policy_output.copy()
diff --git a/judo/tasks/__init__.py b/judo/tasks/__init__.py
index 38a077e2..027b3cf2 100644
--- a/judo/tasks/__init__.py
+++ b/judo/tasks/__init__.py
@@ -1,6 +1,7 @@
 # Copyright (c) 2025 Robotics and AI Institute LLC. All rights reserved.
 
-from typing import Dict, Tuple, Type
+from dataclasses import dataclass
+from typing import Dict, Type
 
 from judo.tasks.base import Task, TaskConfig
 from judo.tasks.caltech_leap_cube import CaltechLeapCube, CaltechLeapCubeConfig
@@ -21,35 +22,100 @@
     SpotTireUpright,
     SpotTireUprightConfig,
 )
+from judo.tasks.spot.spot_constants import SPOT_LOCOMOTION_POLICY_PATH
 
-_registered_tasks: Dict[str, Tuple[Type[Task], Type[TaskConfig]]] = {
-    CylinderPush.name: (CylinderPush, CylinderPushConfig),
-    Cartpole.name: (Cartpole, CartpoleConfig),
-    FR3Pick.name: (FR3Pick, FR3PickConfig),
-    LeapCube.name: (LeapCube, LeapCubeConfig),
-    LeapCubeDown.name: (LeapCubeDown, LeapCubeDownConfig),
-    CaltechLeapCube.name: (CaltechLeapCube, CaltechLeapCubeConfig),
-    SpotBase.name: (SpotBase, SpotBaseConfig),
-    SpotBoxPush.name: (SpotBoxPush, SpotBoxPushConfig),
-    SpotNavigate.name: (SpotNavigate, SpotNavigateConfig),
-    SpotTireRoll.name: (SpotTireRoll, SpotTireRollConfig),
-    SpotTireUpright.name: (SpotTireUpright, SpotTireUprightConfig),
+@dataclass(frozen=True)
+class TaskRegistration:
+    """Complete registration metadata for a task."""
+
+    task_type: Type[Task]
+    task_config_type: Type[TaskConfig]
+    rollout_backend: str = "mujoco"
+    simulation_backend: str = "mujoco"
+    locomotion_policy_path: str | None = None
+
+
+_registered_tasks: Dict[str, TaskRegistration] = {
+    CylinderPush.name: TaskRegistration(CylinderPush, CylinderPushConfig),
+    Cartpole.name: TaskRegistration(Cartpole, CartpoleConfig),
+    FR3Pick.name: TaskRegistration(FR3Pick, FR3PickConfig),
+    LeapCube.name: TaskRegistration(LeapCube, LeapCubeConfig),
+    LeapCubeDown.name: TaskRegistration(LeapCubeDown, LeapCubeDownConfig),
+    CaltechLeapCube.name: TaskRegistration(CaltechLeapCube, CaltechLeapCubeConfig),
+    SpotBase.name: TaskRegistration(
+        SpotBase,
+        SpotBaseConfig,
+        rollout_backend="mujoco_hierarchical",
+        simulation_backend="mujoco_hierarchical",
+        locomotion_policy_path=str(SPOT_LOCOMOTION_POLICY_PATH),
+    ),
+    SpotBoxPush.name: TaskRegistration(
+        SpotBoxPush,
+        SpotBoxPushConfig,
+        rollout_backend="mujoco_hierarchical",
+        simulation_backend="mujoco_hierarchical",
+        locomotion_policy_path=str(SPOT_LOCOMOTION_POLICY_PATH),
+    ),
+    SpotNavigate.name: TaskRegistration(
+        SpotNavigate,
+        SpotNavigateConfig,
+        rollout_backend="mujoco_hierarchical",
+        simulation_backend="mujoco_hierarchical",
+        locomotion_policy_path=str(SPOT_LOCOMOTION_POLICY_PATH),
+    ),
+    SpotTireRoll.name: TaskRegistration(
+        SpotTireRoll,
+        SpotTireRollConfig,
+        rollout_backend="mujoco_hierarchical",
+        simulation_backend="mujoco_hierarchical",
+        locomotion_policy_path=str(SPOT_LOCOMOTION_POLICY_PATH),
+    ),
+    SpotTireUpright.name: TaskRegistration(
+        SpotTireUpright,
+        SpotTireUprightConfig,
+        rollout_backend="mujoco_hierarchical",
+        simulation_backend="mujoco_hierarchical",
+        locomotion_policy_path=str(SPOT_LOCOMOTION_POLICY_PATH),
+    ),
 }
 
 
-def get_registered_tasks() -> Dict[str, Tuple[Type[Task], Type[TaskConfig]]]:
+def get_registered_tasks() -> Dict[str, TaskRegistration]:
     """Returns a dictionary of registered tasks."""
     return _registered_tasks
 
 
-def register_task(name: str, task_type: Type[Task], task_config_type: Type[TaskConfig]) -> None:
-    """Registers a new task."""
-    _registered_tasks[name] = (task_type, task_config_type)
+def get_task_registration(task_name: str) -> TaskRegistration:
+    """Return full registration metadata for a task."""
+    task_entry = _registered_tasks.get(task_name)
+    if task_entry is None:
+        raise ValueError(f"Task {task_name} not found in task registry.")
+    return task_entry
+
+
+def register_task(
+    name: str,
+    task_type: Type[Task],
+    task_config_type: Type[TaskConfig],
+    rollout_backend: str = "mujoco",
+    simulation_backend: str = "mujoco",
+    locomotion_policy_path: str | None = None,
+) -> None:
+    """Registers a new task and its default controller/simulation backends."""
+    _registered_tasks[name] = TaskRegistration(
+        task_type=task_type,
+        task_config_type=task_config_type,
+        rollout_backend=rollout_backend,
+        simulation_backend=simulation_backend,
+        locomotion_policy_path=locomotion_policy_path,
+    )
 
 
 __all__ = [
     "get_registered_tasks",
+    "get_task_registration",
     "register_task",
+    "TaskRegistration",
     "Task",
     "TaskConfig",
     "CaltechLeapCube",
diff --git a/judo/tasks/base.py b/judo/tasks/base.py
index a2676ecc..670cd074 100644
--- a/judo/tasks/base.py
+++ b/judo/tasks/base.py
@@ -81,19 +81,10 @@ def nu(self) -> int:
         """Number of control inputs. The same as the MjModel for this task."""
         return self.model.nu
 
-    @property
-    def locomotion_policy_path(self) -> str | None:
-        """Path to locomotion policy for this task, or None if not used.
-
-        Override in tasks that use a learned locomotion policy
-        (e.g., Spot tasks that run an ONNX policy at 50Hz).
-        """
-        return None
-
     @property
     def uses_locomotion_policy(self) -> bool:
         """Whether this task uses a locomotion policy for simulation."""
-        return self.locomotion_policy_path is not None
+        return False
 
     @property
     def actuator_ctrlrange(self) -> np.ndarray:
diff --git a/judo/tasks/spot/spot_base.py b/judo/tasks/spot/spot_base.py
index 32ab7925..683001a9 100644
--- a/judo/tasks/spot/spot_base.py
+++ b/judo/tasks/spot/spot_base.py
@@ -32,7 +32,6 @@
     LEG_SOFT_UPPER_JOINT_LIMITS,
     LEGS_STANDING_POS,
     LEGS_STANDING_POS_RL,
-    SPOT_LOCOMOTION_POLICY_PATH,
     STANDING_HEIGHT,
     STANDING_HEIGHT_CMD,
     TORSO_CMD_INDS,
@@ -108,9 +107,9 @@ def physics_substeps(self) -> int:  # type: ignore[override]
         return 2
 
     @property
-    def locomotion_policy_path(self) -> str:
-        """Path to Spot locomotion policy."""
-        return str(SPOT_LOCOMOTION_POLICY_PATH)
+    def uses_locomotion_policy(self) -> bool:  # type: ignore[override]
+        """Spot tasks always use a locomotion policy backend."""
+        return True
 
     def __init__(
         self,
diff --git a/judo/utils/policy_mj_rollout_backend.py b/judo/utils/hierarchical_mj_rollout_backend.py
similarity index 88%
rename from judo/utils/policy_mj_rollout_backend.py
rename to judo/utils/hierarchical_mj_rollout_backend.py
index ca19f1a7..fc216d72 100644
--- a/judo/utils/policy_mj_rollout_backend.py
+++ b/judo/utils/hierarchical_mj_rollout_backend.py
@@ -1,6 +1,6 @@
 # Copyright (c) 2025 Robotics and AI Institute LLC. All rights reserved.
 
-"""MuJoCo rollout backend with locomotion policy support."""
+"""MuJoCo rollout backend with hierarchical low-level policy support."""
 
 from pathlib import Path
 
@@ -11,8 +11,8 @@
 from judo.utils.rollout_backend import RolloutBackend
 
 
-class PolicyMJRolloutBackend(RolloutBackend):
-    """Rollout backend with C++ mujoco_extensions and ONNX locomotion policy inference.
+class HierarchicalMJRolloutBackend(RolloutBackend):
+    """Rollout backend with C++ mujoco_extensions and ONNX low-level policy inference.
 
     For Spot tasks, the command format is a 25-dim vector:
     [base_vel(3), arm(7), legs(12), torso(3)]
@@ -25,12 +25,12 @@ def __init__(
         policy_path: str | Path,
         physics_substeps: int = 2,
     ) -> None:
-        """Initialize the policy rollout backend.
+        """Initialize the hierarchical rollout backend.
 
         Args:
             model: MuJoCo model for the scene.
             num_threads: Number of parallel rollout threads.
-            policy_path: Path to ONNX locomotion policy.
+            policy_path: Path to the ONNX low-level policy.
             physics_substeps: Physics steps per control step.
         """
         self.num_threads = num_threads
@@ -41,7 +41,7 @@ def __init__(
         self._setup_mujoco_extensions(model, policy_path, num_threads)
 
     def _setup_mujoco_extensions(self, model: MjModel, policy_path: str | Path, num_threads: int) -> None:
-        """Setup the mujoco_extensions C++ rollout backend with ONNX policy."""
+        """Setup the mujoco_extensions C++ rollout backend with ONNX low-level policy."""
         try:
             from mujoco_extensions.policy_rollout import create_systems_vector, threaded_rollout  # type: ignore  # noqa: PLC0415, I001
         except ImportError as e:
@@ -60,7 +60,7 @@ def rollout(
         controls: np.ndarray,
         last_policy_output: np.ndarray | None = None,
     ) -> tuple[np.ndarray, np.ndarray, np.ndarray | None]:
-        """Conduct parallel rollouts with policy inference.
+        """Conduct parallel rollouts with hierarchical policy inference.
 
         Args:
             x0: Initial state, shape (nq+nv,). Will be tiled to num_threads internally.
@@ -78,7 +78,7 @@ def rollout(
             x0 = np.tile(x0, (self.num_threads, 1))
 
         if last_policy_output is None:
-            raise ValueError("last_policy_output is required for PolicyMJRolloutBackend")
+            raise ValueError("last_policy_output is required for HierarchicalMJRolloutBackend")
 
         x0 = np.asarray(x0, dtype=np.float64)
         controls = np.asarray(controls, dtype=np.float64)
diff --git a/judo/visualizers/visualizer.py b/judo/visualizers/visualizer.py
index b058d88a..d69e1c3a 100644
--- a/judo/visualizers/visualizer.py
+++ b/judo/visualizers/visualizer.py
@@ -15,7 +15,7 @@
 from judo.controller import ControllerConfig
 from judo.gui import create_gui_elements
 from judo.optimizers import get_registered_optimizers
-from judo.tasks import Task, TaskConfig, get_registered_tasks
+from judo.tasks import TaskRegistration, get_registered_tasks
 from judo.visualizers.model import ViserMjModel
 
 ElementType = GuiImageHandle | GuiInputHandle | GuiFolderHandle | MeshHandle | IcosphereHandle
@@ -39,7 +39,7 @@ def __init__(
         optimizer_override_cfg: DictConfig | None = None,
         sim_pause_button: bool = True,
         geom_exclude_substring: str = "collision",
-        available_tasks: dict[str, tuple[type[Task], type[TaskConfig]]] | None = None,
+        available_tasks: dict[str, TaskRegistration] | None = None,
     ) -> None:
         """Initialize the visualization node."""
         # handling custom task and optimizer registration
@@ -106,8 +106,7 @@ def set_task(self, task_name: str, optimizer_name: str) -> None:
         if task_entry is None:
             raise ValueError(f"Task {task_name} not found in the task registry.")
 
-        task_cls, _ = task_entry
-        self.task = task_cls()
+        self.task = task_entry.task_type()
         self.task_config = self.task.config
         self.data = mujoco.MjData(self.task.model)
         self.viser_model = ViserMjModel(
diff --git a/tests/test_simulation/test_simulation.py b/tests/test_simulation/test_simulation.py
index eb9d4a25..5254de27 100644
--- a/tests/test_simulation/test_simulation.py
+++ b/tests/test_simulation/test_simulation.py
@@ -4,7 +4,8 @@
 
 import numpy as np
 
-from judo.simulation import MJSimulation, PolicyMJSimulation
+from judo.simulation import HierarchicalMJSimulation, MJSimulation
+from judo.tasks import get_task_registration
 from judo.tasks.cartpole import Cartpole, CartpoleConfig
 from judo.tasks.cylinder_push import CylinderPush, CylinderPushConfig
 
@@ -41,15 +42,16 @@ def test_simulation_data_step(temp_np_seed: Callable) -> None:
 
 
 def test_spot_simulation_init() -> None:
-    """Test PolicyMJSimulation initializes with a Spot task and C++ systems."""
-    sim = PolicyMJSimulation(init_task="spot_base")
+    """Test HierarchicalMJSimulation initializes with a Spot task and C++ systems."""
+    sim = HierarchicalMJSimulation(init_task="spot_base")
     assert sim._systems is not None
-    assert sim.task.locomotion_policy_path is not None
+    assert sim.task.uses_locomotion_policy
+    assert get_task_registration(sim.task.name).locomotion_policy_path is not None
 
 
 def test_spot_simulation_step() -> None:
-    """Test PolicyMJSimulation steps correctly with Spot locomotion policy."""
-    sim = PolicyMJSimulation(init_task="spot_base")
+    """Test HierarchicalMJSimulation steps correctly with Spot locomotion policy."""
+    sim = HierarchicalMJSimulation(init_task="spot_base")
     qpos_before = sim.task.data.qpos.copy()
     command = np.zeros(sim.task.nu)
     sim.step(command)
diff --git a/tests/test_tasks/test_spot.py b/tests/test_tasks/test_spot.py
index 33d96cac..530357b9 100644
--- a/tests/test_tasks/test_spot.py
+++ b/tests/test_tasks/test_spot.py
@@ -4,6 +4,7 @@
 
 import numpy as np
 
+from judo.tasks import get_task_registration
 from judo.tasks.spot import SpotTireUpright
 from judo.tasks.spot.spot_base import SpotBase
 
@@ -13,7 +14,8 @@ def test_spot_base_init() -> None:
     task = SpotBase()
     assert task.name == "spot_base"
     assert task.physics_substeps == 2
-    assert task.locomotion_policy_path is not None
+    assert task.uses_locomotion_policy
+    assert get_task_registration(task.name).locomotion_policy_path is not None
     # Base + arm = 3 + 7 = 10
     assert task.nu == 10
 
@@ -23,7 +25,8 @@ def test_spot_tire_upright_init() -> None:
     task = SpotTireUpright()
     assert task.name == "spot_tire_upright"
     assert task.physics_substeps == 2
-    assert task.locomotion_policy_path is not None
+    assert task.uses_locomotion_policy
+    assert get_task_registration(task.name).locomotion_policy_path is not None
     # Base + arm + legs + leg_selection = 3 + 7 + 6 + 1 = 17
     assert task.nu == 17
 

From 970114abe3a9a69debf3d8691a0c277b8c20fb56 Mon Sep 17 00:00:00 2001
From: Duy Ta <dta@rai-inst.com>
Date: Tue, 5 May 2026 16:08:47 -0400
Subject: [PATCH 13/19] precommit

---
 judo/app/utils.py             | 6 +++---
 judo/cli.py                   | 3 +--
 judo/controller/controller.py | 1 -
 judo/simulation/__init__.py   | 1 -
 judo/tasks/__init__.py        | 1 +
 5 files changed, 5 insertions(+), 7 deletions(-)

diff --git a/judo/app/utils.py b/judo/app/utils.py
index ed2c9a53..4c0486d1 100644
--- a/judo/app/utils.py
+++ b/judo/app/utils.py
@@ -34,9 +34,9 @@ def register_tasks_from_cfg(task_registration_cfg: DictConfig) -> None:
         locomotion_policy_path = task_dict.get("locomotion_policy_path", None)
         assert isinstance(rollout_backend, str), "rollout_backend must be a string."
         assert isinstance(simulation_backend, str), "simulation_backend must be a string."
-        assert locomotion_policy_path is None or isinstance(
-            locomotion_policy_path, str
-        ), "locomotion_policy_path must be a string if provided."
+        assert locomotion_policy_path is None or isinstance(locomotion_policy_path, str), (
+            "locomotion_policy_path must be a string if provided."
+        )
         register_task(
             str(task_name),
             task_cls,
diff --git a/judo/cli.py b/judo/cli.py
index 42e49541..7a9a29d4 100644
--- a/judo/cli.py
+++ b/judo/cli.py
@@ -147,8 +147,7 @@ def _require_mujoco_extensions() -> None:
         import mujoco_extensions  # noqa: F401, PLC0415
     except Exception as e:  # pragma: no cover - environment dependent
         raise RuntimeError(
-            "mujoco_extensions is required but could not be imported. "
-            "Build it with: pixi run build"
+            "mujoco_extensions is required but could not be imported. Build it with: pixi run build"
         ) from e
 
 
diff --git a/judo/controller/controller.py b/judo/controller/controller.py
index 422f6588..87815273 100644
--- a/judo/controller/controller.py
+++ b/judo/controller/controller.py
@@ -30,7 +30,6 @@
 from judo.utils.timer import Timer
 from judo.visualizers.utils import get_trace_sensors
 
-
 RolloutBackendEntry = type[RolloutBackend]
 
 
diff --git a/judo/simulation/__init__.py b/judo/simulation/__init__.py
index f6d2f043..cb16104d 100644
--- a/judo/simulation/__init__.py
+++ b/judo/simulation/__init__.py
@@ -4,7 +4,6 @@
 from judo.simulation.hierarchical_mj_simulation import HierarchicalMJSimulation
 from judo.simulation.mj_simulation import MJSimulation
 
-
 DEFAULT_SIMULATION_BACKEND_REGISTRY: dict[str, type[Simulation]] = {
     "mujoco": MJSimulation,
     "mujoco_hierarchical": HierarchicalMJSimulation,
diff --git a/judo/tasks/__init__.py b/judo/tasks/__init__.py
index 027b3cf2..30be02ad 100644
--- a/judo/tasks/__init__.py
+++ b/judo/tasks/__init__.py
@@ -24,6 +24,7 @@
 )
 from judo.tasks.spot.spot_constants import SPOT_LOCOMOTION_POLICY_PATH
 
+
 @dataclass(frozen=True)
 class TaskRegistration:
     """Complete registration metadata for a task."""

From 51e4f29b78e155fe0ad0be14a890f3b489bfc3aa Mon Sep 17 00:00:00 2001
From: Duy Ta <dta@rai-inst.com>
Date: Tue, 5 May 2026 16:40:59 -0400
Subject: [PATCH 14/19] cosmetic changes

---
 judo/app/dora/controller_node.py              |   3 -
 judo/app/dora/simulation.py                   | 122 ++++++++++++++
 judo/app/dora/simulation_node.py              |   4 +-
 judo/app/dora/visualization.py                | 157 ++++++++++++++++++
 judo/app/dora/visualization_node.py           |  10 +-
 judo/app/structs.py                           |   4 +-
 judo/app/utils.py                             |  39 ++++-
 judo/configs/judo_dora_default.yaml           |   6 +-
 judo/controller/controller.py                 |  18 +-
 judo/simulation/base.py                       |   8 +-
 .../test_action_normalization.py              |   4 -
 tests/test_controller/test_controller.py      |   2 -
 12 files changed, 344 insertions(+), 33 deletions(-)
 create mode 100644 judo/app/dora/simulation.py
 create mode 100644 judo/app/dora/visualization.py

diff --git a/judo/app/dora/controller_node.py b/judo/app/dora/controller_node.py
index 412a9691..c97cb010 100644
--- a/judo/app/dora/controller_node.py
+++ b/judo/app/dora/controller_node.py
@@ -11,7 +11,6 @@
 
 from judo.app.structs import MujocoState
 from judo.controller import Controller, make_controller
-from judo.tasks import get_task_registration
 
 
 class ControllerNode(DoraNode):
@@ -50,13 +49,11 @@ def __init__(
 
     def _build_controller(self, task_name: str, optimizer_name: str) -> Controller:
         """Build controller using the task's registered rollout backend."""
-        rollout_backend = get_task_registration(task_name).rollout_backend
         return self._make_controller_fn(
             init_task=task_name,
             init_optimizer=optimizer_name,
             task_registration_cfg=self._task_registration_cfg,
             optimizer_registration_cfg=self._optimizer_registration_cfg,
-            rollout_backend=rollout_backend,
         )
 
     def _current_optimizer_name(self) -> str:
diff --git a/judo/app/dora/simulation.py b/judo/app/dora/simulation.py
new file mode 100644
index 00000000..930892fe
--- /dev/null
+++ b/judo/app/dora/simulation.py
@@ -0,0 +1,122 @@
+# Copyright (c) 2025 Robotics and AI Institute LLC. All rights reserved.
+
+import time
+import warnings
+from typing import Callable
+
+from dora_utils.dataclasses import from_arrow, to_arrow
+from dora_utils.node import DoraNode, on_event
+from omegaconf import DictConfig
+
+from judo.app.structs import SplineData
+from judo.simulation.base import Simulation
+from judo.tasks import get_registered_tasks
+
+
+class SimulationNode(DoraNode):
+    """The simulation node."""
+
+    def __init__(
+        self,
+        node_id: str = "simulation",
+        init_task: str = "cylinder_push",
+        max_workers: int | None = None,
+        task_registration_cfg: DictConfig | None = None,
+        simulation_backend: str = "mujoco",
+        custom_backends: dict[str, type[Simulation]] | None = None,
+    ) -> None:
+        """Initialize the simulation node."""
+        super().__init__(node_id=node_id, max_workers=max_workers)
+        self._simulation_backend = simulation_backend
+        self._task_registration_cfg = task_registration_cfg
+        self._custom_backends = custom_backends or {}
+        self._init_sim(init_task)
+        self.control_spline: Callable | None = None
+        self.write_states()
+
+    def _resolve_backend(self, backend: str) -> type:
+        """Resolve a simulation backend class by name, checking custom backends first."""
+        if backend in self._custom_backends:
+            return self._custom_backends[backend]
+        from judo.simulation import get_simulation_backend  # noqa: PLC0415, I001
+
+        return get_simulation_backend(backend)
+
+    def _init_sim(self, task_name: str) -> None:
+        """Initialize simulation, auto-upgrading to policy backend if needed."""
+        task_entry = get_registered_tasks().get(task_name)
+        if task_entry is None:
+            raise ValueError(f"Task {task_name} not found in task registry.")
+
+        task_cls, _ = task_entry
+        backend = getattr(task_cls, "default_backend", None) or self._simulation_backend
+
+        sim_backend_cls = self._resolve_backend(backend)
+        self.sim = sim_backend_cls(init_task=task_name, task_registration_cfg=self._task_registration_cfg)
+
+        # Auto-upgrade to policy backend if task requires locomotion policy
+        if backend == "mujoco" and self.sim.task.uses_locomotion_policy:
+            sim_backend_cls = self._resolve_backend("mujoco_policy")
+            self.sim = sim_backend_cls(init_task=task_name, task_registration_cfg=self._task_registration_cfg)
+
+    @on_event("INPUT", "task")
+    def update_task(self, event: dict) -> None:
+        """Event handler for processing task updates."""
+        new_task = event["value"].to_numpy(zero_copy_only=False)[0]
+        self._init_sim(new_task)
+        self.control_spline = None  # Clear stale spline
+
+    def spin(self) -> None:
+        """Spin logic for the simulation node."""
+        try:
+            while True:
+                start_time = time.time()
+                self.parse_messages()
+
+                if self.control_spline is not None:
+                    command = self.control_spline(self.sim.task.data.time)
+                    if command.shape[-1] == self.sim.task.nu:
+                        self.sim.step(command)
+                    else:
+                        warnings.warn(
+                            f"Control command has wrong number of dimensions! Expected {self.sim.task.nu}, got {command.shape[-1]}",
+                            stacklevel=2,
+                        )
+
+                self.write_states()
+
+                # Force simulation node to run at fixed rate specified by simulation timestep (specified in the model).
+                dt_des = self.sim.timestep
+                dt_elapsed = time.time() - start_time
+                if dt_elapsed < dt_des:
+                    time.sleep(dt_des - dt_elapsed)
+                else:
+                    warnings.warn(
+                        f"Sim step {dt_elapsed:.3f} longer than desired step {dt_des:.3f}!",
+                        stacklevel=2,
+                    )
+        except KeyboardInterrupt:
+            pass
+
+    def write_states(self) -> None:
+        """Reads data from simulation and writes to output topic."""
+        arr, metadata = to_arrow(self.sim.sim_state)
+        self.node.send_output("states", arr, metadata)
+        arr, metadata = to_arrow(self.sim.render_pose)
+        self.node.send_output("render_pose", arr, metadata)
+
+    @on_event("INPUT", "sim_pause")
+    def set_paused_status(self, event: dict) -> None:
+        """Event handler for processing pause status updates."""
+        self.sim.pause()
+
+    @on_event("INPUT", "task_reset")
+    def reset_task(self, event: dict) -> None:
+        """Resets the task."""
+        self.sim.task.reset()
+
+    @on_event("INPUT", "controls")
+    def update_control(self, event: dict) -> None:
+        """Event handler for processing controls received from controller node."""
+        spline_data = from_arrow(event["value"], event["metadata"], SplineData)
+        self.control_spline = spline_data.spline()
diff --git a/judo/app/dora/simulation_node.py b/judo/app/dora/simulation_node.py
index 79059c7a..d3925847 100644
--- a/judo/app/dora/simulation_node.py
+++ b/judo/app/dora/simulation_node.py
@@ -101,8 +101,8 @@ def write_states(self) -> None:
         """Reads data from simulation and writes to output topic."""
         arr, metadata = to_arrow(self.sim.sim_state)
         self.node.send_output("states", arr, metadata)
-        arr, metadata = to_arrow(self.sim.world_state)
-        self.node.send_output("world_states", arr, metadata)
+        arr, metadata = to_arrow(self.sim.render_pose)
+        self.node.send_output("render_pose", arr, metadata)
 
     @on_event("INPUT", "sim_pause")
     def set_paused_status(self, event: dict) -> None:
diff --git a/judo/app/dora/visualization.py b/judo/app/dora/visualization.py
new file mode 100644
index 00000000..bed42747
--- /dev/null
+++ b/judo/app/dora/visualization.py
@@ -0,0 +1,157 @@
+# Copyright (c) 2025 Robotics and AI Institute LLC. All rights reserved.
+
+import warnings
+
+import pyarrow as pa
+from dora_utils.dataclasses import from_arrow, to_arrow
+from dora_utils.node import DoraNode, on_event
+from omegaconf import DictConfig
+from viser import GuiFolderHandle, GuiImageHandle, GuiInputHandle, IcosphereHandle, MeshHandle
+
+from judo.app.structs import RenderPose
+from judo.tasks import Task, TaskConfig
+from judo.visualizers.visualizer import Visualizer
+
+ElementType = GuiImageHandle | GuiInputHandle | GuiFolderHandle | MeshHandle | IcosphereHandle
+
+
+class VisualizationNode(DoraNode):
+    """The visualization node."""
+
+    def __init__(
+        self,
+        node_id: str = "visualization",
+        max_workers: int | None = None,
+        init_task: str = "cylinder_push",
+        init_optimizer: str = "cem",
+        task_registration_cfg: DictConfig | None = None,
+        optimizer_registration_cfg: DictConfig | None = None,
+        controller_override_cfg: DictConfig | None = None,
+        optimizer_override_cfg: DictConfig | None = None,
+        sim_pause_button: bool = True,
+        geom_exclude_substring: str = "collision",
+        available_tasks: dict[str, tuple[type[Task], type[TaskConfig]]] | None = None,
+    ) -> None:
+        """Initialize the visualization node."""
+        super().__init__(node_id=node_id, max_workers=max_workers)
+        self.visualizer = Visualizer(
+            init_task=init_task,
+            init_optimizer=init_optimizer,
+            task_registration_cfg=task_registration_cfg,
+            optimizer_registration_cfg=optimizer_registration_cfg,
+            controller_override_cfg=controller_override_cfg,
+            optimizer_override_cfg=optimizer_override_cfg,
+            sim_pause_button=sim_pause_button,
+            geom_exclude_substring=geom_exclude_substring,
+            available_tasks=available_tasks,
+        )
+
+    def write_sim_pause(self) -> None:
+        """Write the sim pause signal to the GUI."""
+        with self.visualizer.sim_pause_lock:
+            self.node.send_output("sim_pause", pa.array([1]))  # dummy value
+        self.visualizer.sim_pause_updated.clear()
+
+    def write_task(self) -> None:
+        """Write the task name to the GUI."""
+        with self.visualizer.task_lock:
+            self.node.send_output("task", pa.array([self.visualizer.task_name]))
+        self.visualizer.task_updated.clear()
+
+    def write_task_reset(self) -> None:
+        """Write the task reset signal to the GUI."""
+        with self.visualizer.task_lock:
+            self.node.send_output("task_reset", pa.array([1]))  # dummy value
+        self.visualizer.task_reset_updated.clear()
+
+    def write_optimizer(self) -> None:
+        """Write the optimizer name to the GUI."""
+        with self.visualizer.optimizer_lock:
+            self.node.send_output("optimizer", pa.array([self.visualizer.optimizer_name]))
+        self.visualizer.optimizer_updated.clear()
+
+    def write_controller_config(self) -> None:
+        """Write the controller config to the GUI."""
+        with self.visualizer.controller_config_lock:
+            self.node.send_output("controller_config", *to_arrow(self.visualizer.controller_config))
+        self.visualizer.controller_config_updated.clear()
+
+    def write_optimizer_config(self) -> None:
+        """Write the optimizer config to the GUI."""
+        with self.visualizer.optimizer_config_lock:
+            self.node.send_output("optimizer_config", *to_arrow(self.visualizer.optimizer_config))
+        self.visualizer.optimizer_config_updated.clear()
+
+    def write_task_config(self) -> None:
+        """Write the task config to the GUI."""
+        with self.visualizer.task_config_lock:
+            self.node.send_output("task_config", *to_arrow(self.visualizer.task_config))
+        self.visualizer.task_config_updated.clear()
+
+    @on_event("INPUT", "render_pose")
+    def update_states(self, event: dict) -> None:
+        """Callback to update states on receiving a new state measurement."""
+        if self.visualizer.controller_config.spline_order == "cubic" and self.visualizer.optimizer_config.num_nodes < 4:
+            warnings.warn("Cubic splines require at least 4 nodes. Setting num_nodes=4.", stacklevel=2)
+            for e in self.visualizer.gui_elements["optimizer_params"]:
+                if e.label == "num_nodes":
+                    e.value = 4
+                    break
+            self.visualizer.optimizer_config_updated.set()
+
+        # TODO: change the mujoco state
+        render_pose_msg = from_arrow(event["value"], event["metadata"], RenderPose)
+        try:
+            with self.visualizer.task_lock:
+                self.visualizer.data.xpos[:] = render_pose_msg.xpos
+                self.visualizer.data.xquat[:] = render_pose_msg.xquat
+                self.visualizer.viser_model.set_data(self.visualizer.data)
+        except ValueError:
+            # we're switching tasks and the new task has a different number of xpos/xquat
+            return
+
+    @on_event("INPUT", "traces")
+    def update_traces(self, event: dict) -> None:
+        """Callback to update traces on receiving a new trace measurement."""
+        traces_flat = event["value"].to_numpy()
+        all_traces_rollout_size = int(event["metadata"]["all_traces_rollout_size"])
+        shape = event["metadata"]["shape"]
+        traces = traces_flat.reshape(*shape)
+        with self.visualizer.task_lock:
+            self.visualizer.viser_model.set_traces(traces, all_traces_rollout_size)
+
+    @on_event("INPUT", "plan_time")
+    def update_plan_time(self, event: dict) -> None:
+        """Callback to update plan time on receiving a new plan time measurement."""
+        plan_time_s = event["value"].to_numpy(zero_copy_only=False)[0]
+        # GUI element might not exist during task switch
+        if "plan_time_display" in self.visualizer.gui_elements:
+            self.visualizer.gui_elements["plan_time_display"].value = plan_time_s * 1000  # ms
+
+    def spin(self) -> None:
+        """Spin logic for the visualization node."""
+        try:
+            for event in self.node:
+                if self.visualizer.sim_pause_updated.is_set():
+                    self.write_sim_pause()
+                if self.visualizer.task_updated.is_set():
+                    self.write_task()
+                if self.visualizer.task_reset_updated.is_set():
+                    self.write_task_reset()
+                if self.visualizer.optimizer_updated.is_set():
+                    self.write_optimizer()
+                if self.visualizer.controller_config_updated.is_set():
+                    self.write_controller_config()
+                if self.visualizer.optimizer_config_updated.is_set():
+                    self.write_optimizer_config()
+                if self.visualizer.task_config_updated.is_set():
+                    self.write_task_config()
+
+                self.handle(event)
+        except KeyboardInterrupt:
+            pass
+
+    def cleanup(self) -> None:
+        """Cleanup the visualization node."""
+        self.visualizer.cleanup()
+        super().cleanup()
diff --git a/judo/app/dora/visualization_node.py b/judo/app/dora/visualization_node.py
index a68fa56a..4b145951 100644
--- a/judo/app/dora/visualization_node.py
+++ b/judo/app/dora/visualization_node.py
@@ -8,7 +8,7 @@
 from omegaconf import DictConfig
 from viser import GuiFolderHandle, GuiImageHandle, GuiInputHandle, IcosphereHandle, MeshHandle
 
-from judo.app.structs import WorldState
+from judo.app.structs import RenderPose
 from judo.tasks import TaskRegistration
 from judo.visualizers.visualizer import Visualizer
 
@@ -103,7 +103,7 @@ def write_task_config(self) -> None:
             self.node.send_output("task_config", *to_arrow(self.visualizer.task_config))
         self.visualizer.task_config_updated.clear()
 
-    @on_event("INPUT", "world_states")
+    @on_event("INPUT", "render_pose")
     def update_states(self, event: dict) -> None:
         """Callback to update states on receiving a new state measurement."""
         if self.visualizer.controller_config.spline_order == "cubic" and self.visualizer.optimizer_config.num_nodes < 4:
@@ -115,11 +115,11 @@ def update_states(self, event: dict) -> None:
             self.visualizer.optimizer_config_updated.set()
 
         # TODO: change the mujoco state
-        world_state_msg = from_arrow(event["value"], event["metadata"], WorldState)
+        render_pose_msg = from_arrow(event["value"], event["metadata"], RenderPose)
         try:
             with self.visualizer.task_lock:
-                self.visualizer.data.xpos[:] = world_state_msg.xpos
-                self.visualizer.data.xquat[:] = world_state_msg.xquat
+                self.visualizer.data.xpos[:] = render_pose_msg.xpos
+                self.visualizer.data.xquat[:] = render_pose_msg.xquat
                 self.visualizer.viser_model.set_data(self.visualizer.data)
         except ValueError:
             # we're switching tasks and the new task has a different number of xpos/xquat
diff --git a/judo/app/structs.py b/judo/app/structs.py
index ac143caf..3769efe4 100644
--- a/judo/app/structs.py
+++ b/judo/app/structs.py
@@ -83,8 +83,8 @@ def spline(self) -> interp1d:
 
 
 @dataclass
-class WorldState:
-    """Struct for world states used in visualization."""
+class RenderPose:
+    """Struct for visualization poses used by the renderer."""
 
     xpos: np.ndarray
     xquat: np.ndarray
diff --git a/judo/app/utils.py b/judo/app/utils.py
index 4c0486d1..340609da 100644
--- a/judo/app/utils.py
+++ b/judo/app/utils.py
@@ -17,7 +17,28 @@ def get_class_from_string(class_path: str) -> type:
 
 
 def register_tasks_from_cfg(task_registration_cfg: DictConfig) -> None:
-    """Register custom tasks."""
+    """Register custom tasks.
+
+    Args:
+        task_registration_cfg: Mapping keyed by task name. Each value must contain:
+            `task`: import path to the Task class
+            `config`: import path to the TaskConfig class
+
+            Optional keys:
+            `rollout_backend`: rollout backend registry key for controllers
+            `simulation_backend`: simulation backend registry key for the simulation node
+            `locomotion_policy_path`: path to a low-level policy used by hierarchical tasks
+
+            Example schema:
+            {
+                "cylinder_push": {
+                    "task": "judo.tasks.cylinder_push.CylinderPush",
+                    "config": "judo.tasks.cylinder_push.CylinderPushConfig",
+                    "rollout_backend": "mujoco",
+                    "simulation_backend": "mujoco",
+                }
+            }
+    """
     for task_name in task_registration_cfg.keys():
         task_dict = task_registration_cfg.get(task_name, {})
         allowed_keys = {"task", "config", "rollout_backend", "simulation_backend", "locomotion_policy_path"}
@@ -48,7 +69,21 @@ def register_tasks_from_cfg(task_registration_cfg: DictConfig) -> None:
 
 
 def register_optimizers_from_cfg(optimizer_registration_cfg: DictConfig) -> None:
-    """Register custom optimizers."""
+    """Register custom optimizers.
+
+    Args:
+        optimizer_registration_cfg: Mapping keyed by optimizer name. Each value must contain:
+            `optimizer`: import path to the Optimizer class
+            `config`: import path to the OptimizerConfig class
+
+            Example schema:
+            {
+                "cem": {
+                    "optimizer": "judo.optimizers.cem.CrossEntropyMethod",
+                    "config": "judo.optimizers.cem.CrossEntropyMethodConfig",
+                }
+            }
+    """
     for optimizer_name in optimizer_registration_cfg.keys():
         optimizer_dict = optimizer_registration_cfg.get(optimizer_name, {})
         assert set(optimizer_dict.keys()) == {"optimizer", "config"}, (
diff --git a/judo/configs/judo_dora_default.yaml b/judo/configs/judo_dora_default.yaml
index 79c59b47..0ee7f442 100644
--- a/judo/configs/judo_dora_default.yaml
+++ b/judo/configs/judo_dora_default.yaml
@@ -20,12 +20,12 @@ dataflow:
           queue_size: 1
       outputs:
         - states
-        - world_states
+        - render_pose
     - id: visualization
       path: dynamic
       inputs:
-        world_states:
-          source: simulation/world_states
+        render_pose:
+          source: simulation/render_pose
           queue_size: 1
         traces:
           source: controller/traces
diff --git a/judo/controller/controller.py b/judo/controller/controller.py
index 87815273..a5fcefe8 100644
--- a/judo/controller/controller.py
+++ b/judo/controller/controller.py
@@ -717,7 +717,6 @@ def make_controller(
     init_optimizer: str,
     task_registration_cfg: DictConfig | None = None,
     optimizer_registration_cfg: DictConfig | None = None,
-    rollout_backend: str = "mujoco",
     controller_cls: type[Controller] | None = None,
     **controller_kwargs: Any,
 ) -> Controller:
@@ -726,10 +725,16 @@ def make_controller(
     Args:
         init_task: The task name to use.
         init_optimizer: The optimizer name to use.
-        task_registration_cfg: Optional task registration config.
-        optimizer_registration_cfg: Optional optimizer registration config.
-        rollout_backend: Either a backend type string ("mujoco") to create a new backend,
-            or an existing WarpRolloutBackend instance to share with other controllers.
+        task_registration_cfg: Optional task registration overrides keyed by task name.
+            Each entry must contain `task` and `config` import paths, and may also define
+            `rollout_backend`, `simulation_backend`, and `locomotion_policy_path`.
+            See register_tasks_from_cfg for the exact supported schema.
+        optimizer_registration_cfg: Optional optimizer registration overrides keyed by
+            optimizer name. Each entry must contain `optimizer` and `config` import paths.
+            See register_optimizers_from_cfg for the exact supported schema.
+        controller_cls: Optional controller class to instantiate instead of Controller.
+        **controller_kwargs: Additional keyword arguments forwarded to the controller
+            constructor.
 
     Returns:
         The created Controller instance.
@@ -746,6 +751,7 @@ def make_controller(
 
     assert task_entry is not None, f"Task {init_task} not found in task registry."
     assert optimizer_entry is not None, f"Optimizer {init_optimizer} not found in optimizer registry."
+    task_registration = get_task_registration(init_task)
 
     # instantiate the task/optimizer/controller
     task = task_entry.task_type()
@@ -763,6 +769,6 @@ def make_controller(
         controller_config=controller_cfg,
         task=task,
         optimizer=optimizer,
-        rollout_backend=rollout_backend,
+        rollout_backend=task_registration.rollout_backend,
         **controller_kwargs,
     )
diff --git a/judo/simulation/base.py b/judo/simulation/base.py
index e76ac743..bf9cbc77 100644
--- a/judo/simulation/base.py
+++ b/judo/simulation/base.py
@@ -5,7 +5,7 @@
 import numpy as np
 from omegaconf import DictConfig
 
-from judo.app.structs import MujocoState, WorldState
+from judo.app.structs import MujocoState, RenderPose
 from judo.app.utils import register_tasks_from_cfg
 from judo.tasks import get_registered_tasks
 from judo.tasks.base import Task
@@ -66,9 +66,9 @@ def sim_state(self) -> MujocoState:
         )
 
     @property
-    def world_state(self) -> WorldState:
-        """Returns the current global state."""
-        return WorldState(
+    def render_pose(self) -> RenderPose:
+        """Returns the current pose data used for visualization."""
+        return RenderPose(
             xpos=self.task.data.xpos,  # type: ignore
             xquat=self.task.data.xquat,  # type: ignore
         )
diff --git a/tests/test_controller/test_action_normalization.py b/tests/test_controller/test_action_normalization.py
index e9ba2d03..4c662d7f 100644
--- a/tests/test_controller/test_action_normalization.py
+++ b/tests/test_controller/test_action_normalization.py
@@ -122,7 +122,6 @@ def test_normalizer_type_change() -> None:
     controller = make_controller(
         init_task="cylinder_push",
         init_optimizer="cem",
-        rollout_backend="mujoco",
     )
 
     # Initially should be IdentityNormalizer
@@ -147,7 +146,6 @@ def test_normalizer_in_update_action_loop() -> None:
         controller = make_controller(
             init_task="cylinder_push",
             init_optimizer="cem",
-            rollout_backend="mujoco",
         )
         controller.controller_cfg = ControllerConfig(action_normalizer=normalizer_type)
 
@@ -164,7 +162,6 @@ def test_min_max_normalizer_with_task_control_ranges() -> None:
     controller = make_controller(
         init_task="cylinder_push",
         init_optimizer="cem",
-        rollout_backend="mujoco",
     )
     controller.controller_cfg = ControllerConfig(action_normalizer="min_max", max_opt_iters=1)
 
@@ -196,7 +193,6 @@ def test_running_normalizer_updates_with_optimizer_data() -> None:
     controller = make_controller(
         init_task="cylinder_push",
         init_optimizer="cem",
-        rollout_backend="mujoco",
     )
     controller.controller_cfg = ControllerConfig(action_normalizer="running", max_opt_iters=1)
 
diff --git a/tests/test_controller/test_controller.py b/tests/test_controller/test_controller.py
index 5bf0d7c6..b43765e6 100644
--- a/tests/test_controller/test_controller.py
+++ b/tests/test_controller/test_controller.py
@@ -49,7 +49,6 @@ def _setup_controller(max_opt_iters: int) -> tuple[MockOptimizerTrackNominalKnot
         controller = make_controller(
             init_task="cylinder_push",
             init_optimizer="cem",
-            rollout_backend="mujoco",
         )
         controller.controller_cfg = ControllerConfig(max_opt_iters=max_opt_iters)
         controller.optimizer = opt
@@ -87,7 +86,6 @@ def _setup_controller(opt_cls: type[Optimizer], opt_cfg: OptimizerConfig) -> Con
         controller = make_controller(
             init_task="cylinder_push",
             init_optimizer="cem",
-            rollout_backend="mujoco",
         )
         controller.optimizer = opt
         return controller

From e215bb1caf991f3d81d01e5f566803f59076cf2c Mon Sep 17 00:00:00 2001
From: Duy Ta <dta@rai-inst.com>
Date: Tue, 5 May 2026 16:57:22 -0400
Subject: [PATCH 15/19] more fixes

---
 judo/app/dora/simulation.py       | 122 -----------------------
 judo/app/dora/visualization.py    | 157 ------------------------------
 judo/controller/controller.py     |  22 +++--
 run_mpc/mpc_batch.py              |  18 ++--
 run_mpc/mpc_config.py             |   3 +-
 run_mpc/mpc_setup.py              |   4 +-
 run_mpc/visualize_trajectories.py |   3 +-
 7 files changed, 29 insertions(+), 300 deletions(-)
 delete mode 100644 judo/app/dora/simulation.py
 delete mode 100644 judo/app/dora/visualization.py

diff --git a/judo/app/dora/simulation.py b/judo/app/dora/simulation.py
deleted file mode 100644
index 930892fe..00000000
--- a/judo/app/dora/simulation.py
+++ /dev/null
@@ -1,122 +0,0 @@
-# Copyright (c) 2025 Robotics and AI Institute LLC. All rights reserved.
-
-import time
-import warnings
-from typing import Callable
-
-from dora_utils.dataclasses import from_arrow, to_arrow
-from dora_utils.node import DoraNode, on_event
-from omegaconf import DictConfig
-
-from judo.app.structs import SplineData
-from judo.simulation.base import Simulation
-from judo.tasks import get_registered_tasks
-
-
-class SimulationNode(DoraNode):
-    """The simulation node."""
-
-    def __init__(
-        self,
-        node_id: str = "simulation",
-        init_task: str = "cylinder_push",
-        max_workers: int | None = None,
-        task_registration_cfg: DictConfig | None = None,
-        simulation_backend: str = "mujoco",
-        custom_backends: dict[str, type[Simulation]] | None = None,
-    ) -> None:
-        """Initialize the simulation node."""
-        super().__init__(node_id=node_id, max_workers=max_workers)
-        self._simulation_backend = simulation_backend
-        self._task_registration_cfg = task_registration_cfg
-        self._custom_backends = custom_backends or {}
-        self._init_sim(init_task)
-        self.control_spline: Callable | None = None
-        self.write_states()
-
-    def _resolve_backend(self, backend: str) -> type:
-        """Resolve a simulation backend class by name, checking custom backends first."""
-        if backend in self._custom_backends:
-            return self._custom_backends[backend]
-        from judo.simulation import get_simulation_backend  # noqa: PLC0415, I001
-
-        return get_simulation_backend(backend)
-
-    def _init_sim(self, task_name: str) -> None:
-        """Initialize simulation, auto-upgrading to policy backend if needed."""
-        task_entry = get_registered_tasks().get(task_name)
-        if task_entry is None:
-            raise ValueError(f"Task {task_name} not found in task registry.")
-
-        task_cls, _ = task_entry
-        backend = getattr(task_cls, "default_backend", None) or self._simulation_backend
-
-        sim_backend_cls = self._resolve_backend(backend)
-        self.sim = sim_backend_cls(init_task=task_name, task_registration_cfg=self._task_registration_cfg)
-
-        # Auto-upgrade to policy backend if task requires locomotion policy
-        if backend == "mujoco" and self.sim.task.uses_locomotion_policy:
-            sim_backend_cls = self._resolve_backend("mujoco_policy")
-            self.sim = sim_backend_cls(init_task=task_name, task_registration_cfg=self._task_registration_cfg)
-
-    @on_event("INPUT", "task")
-    def update_task(self, event: dict) -> None:
-        """Event handler for processing task updates."""
-        new_task = event["value"].to_numpy(zero_copy_only=False)[0]
-        self._init_sim(new_task)
-        self.control_spline = None  # Clear stale spline
-
-    def spin(self) -> None:
-        """Spin logic for the simulation node."""
-        try:
-            while True:
-                start_time = time.time()
-                self.parse_messages()
-
-                if self.control_spline is not None:
-                    command = self.control_spline(self.sim.task.data.time)
-                    if command.shape[-1] == self.sim.task.nu:
-                        self.sim.step(command)
-                    else:
-                        warnings.warn(
-                            f"Control command has wrong number of dimensions! Expected {self.sim.task.nu}, got {command.shape[-1]}",
-                            stacklevel=2,
-                        )
-
-                self.write_states()
-
-                # Force simulation node to run at fixed rate specified by simulation timestep (specified in the model).
-                dt_des = self.sim.timestep
-                dt_elapsed = time.time() - start_time
-                if dt_elapsed < dt_des:
-                    time.sleep(dt_des - dt_elapsed)
-                else:
-                    warnings.warn(
-                        f"Sim step {dt_elapsed:.3f} longer than desired step {dt_des:.3f}!",
-                        stacklevel=2,
-                    )
-        except KeyboardInterrupt:
-            pass
-
-    def write_states(self) -> None:
-        """Reads data from simulation and writes to output topic."""
-        arr, metadata = to_arrow(self.sim.sim_state)
-        self.node.send_output("states", arr, metadata)
-        arr, metadata = to_arrow(self.sim.render_pose)
-        self.node.send_output("render_pose", arr, metadata)
-
-    @on_event("INPUT", "sim_pause")
-    def set_paused_status(self, event: dict) -> None:
-        """Event handler for processing pause status updates."""
-        self.sim.pause()
-
-    @on_event("INPUT", "task_reset")
-    def reset_task(self, event: dict) -> None:
-        """Resets the task."""
-        self.sim.task.reset()
-
-    @on_event("INPUT", "controls")
-    def update_control(self, event: dict) -> None:
-        """Event handler for processing controls received from controller node."""
-        spline_data = from_arrow(event["value"], event["metadata"], SplineData)
-        self.control_spline = spline_data.spline()
diff --git a/judo/app/dora/visualization.py b/judo/app/dora/visualization.py
deleted file mode 100644
index bed42747..00000000
--- a/judo/app/dora/visualization.py
+++ /dev/null
@@ -1,157 +0,0 @@
-# Copyright (c) 2025 Robotics and AI Institute LLC. All rights reserved.
-
-import warnings
-
-import pyarrow as pa
-from dora_utils.dataclasses import from_arrow, to_arrow
-from dora_utils.node import DoraNode, on_event
-from omegaconf import DictConfig
-from viser import GuiFolderHandle, GuiImageHandle, GuiInputHandle, IcosphereHandle, MeshHandle
-
-from judo.app.structs import RenderPose
-from judo.tasks import Task, TaskConfig
-from judo.visualizers.visualizer import Visualizer
-
-ElementType = GuiImageHandle | GuiInputHandle | GuiFolderHandle | MeshHandle | IcosphereHandle
-
-
-class VisualizationNode(DoraNode):
-    """The visualization node."""
-
-    def __init__(
-        self,
-        node_id: str = "visualization",
-        max_workers: int | None = None,
-        init_task: str = "cylinder_push",
-        init_optimizer: str = "cem",
-        task_registration_cfg: DictConfig | None = None,
-        optimizer_registration_cfg: DictConfig | None = None,
-        controller_override_cfg: DictConfig | None = None,
-        optimizer_override_cfg: DictConfig | None = None,
-        sim_pause_button: bool = True,
-        geom_exclude_substring: str = "collision",
-        available_tasks: dict[str, tuple[type[Task], type[TaskConfig]]] | None = None,
-    ) -> None:
-        """Initialize the visualization node."""
-        super().__init__(node_id=node_id, max_workers=max_workers)
-        self.visualizer = Visualizer(
-            init_task=init_task,
-            init_optimizer=init_optimizer,
-            task_registration_cfg=task_registration_cfg,
-            optimizer_registration_cfg=optimizer_registration_cfg,
-            controller_override_cfg=controller_override_cfg,
-            optimizer_override_cfg=optimizer_override_cfg,
-            sim_pause_button=sim_pause_button,
-            geom_exclude_substring=geom_exclude_substring,
-            available_tasks=available_tasks,
-        )
-
-    def write_sim_pause(self) -> None:
-        """Write the sim pause signal to the GUI."""
-        with self.visualizer.sim_pause_lock:
-            self.node.send_output("sim_pause", pa.array([1]))  # dummy value
-        self.visualizer.sim_pause_updated.clear()
-
-    def write_task(self) -> None:
-        """Write the task name to the GUI."""
-        with self.visualizer.task_lock:
-            self.node.send_output("task", pa.array([self.visualizer.task_name]))
-        self.visualizer.task_updated.clear()
-
-    def write_task_reset(self) -> None:
-        """Write the task reset signal to the GUI."""
-        with self.visualizer.task_lock:
-            self.node.send_output("task_reset", pa.array([1]))  # dummy value
-        self.visualizer.task_reset_updated.clear()
-
-    def write_optimizer(self) -> None:
-        """Write the optimizer name to the GUI."""
-        with self.visualizer.optimizer_lock:
-            self.node.send_output("optimizer", pa.array([self.visualizer.optimizer_name]))
-        self.visualizer.optimizer_updated.clear()
-
-    def write_controller_config(self) -> None:
-        """Write the controller config to the GUI."""
-        with self.visualizer.controller_config_lock:
-            self.node.send_output("controller_config", *to_arrow(self.visualizer.controller_config))
-        self.visualizer.controller_config_updated.clear()
-
-    def write_optimizer_config(self) -> None:
-        """Write the optimizer config to the GUI."""
-        with self.visualizer.optimizer_config_lock:
-            self.node.send_output("optimizer_config", *to_arrow(self.visualizer.optimizer_config))
-        self.visualizer.optimizer_config_updated.clear()
-
-    def write_task_config(self) -> None:
-        """Write the task config to the GUI."""
-        with self.visualizer.task_config_lock:
-            self.node.send_output("task_config", *to_arrow(self.visualizer.task_config))
-        self.visualizer.task_config_updated.clear()
-
-    @on_event("INPUT", "render_pose")
-    def update_states(self, event: dict) -> None:
-        """Callback to update states on receiving a new state measurement."""
-        if self.visualizer.controller_config.spline_order == "cubic" and self.visualizer.optimizer_config.num_nodes < 4:
-            warnings.warn("Cubic splines require at least 4 nodes. Setting num_nodes=4.", stacklevel=2)
-            for e in self.visualizer.gui_elements["optimizer_params"]:
-                if e.label == "num_nodes":
-                    e.value = 4
-                    break
-            self.visualizer.optimizer_config_updated.set()
-
-        # TODO: change the mujoco state
-        render_pose_msg = from_arrow(event["value"], event["metadata"], RenderPose)
-        try:
-            with self.visualizer.task_lock:
-                self.visualizer.data.xpos[:] = render_pose_msg.xpos
-                self.visualizer.data.xquat[:] = render_pose_msg.xquat
-                self.visualizer.viser_model.set_data(self.visualizer.data)
-        except ValueError:
-            # we're switching tasks and the new task has a different number of xpos/xquat
-            return
-
-    @on_event("INPUT", "traces")
-    def update_traces(self, event: dict) -> None:
-        """Callback to update traces on receiving a new trace measurement."""
-        traces_flat = event["value"].to_numpy()
-        all_traces_rollout_size = int(event["metadata"]["all_traces_rollout_size"])
-        shape = event["metadata"]["shape"]
-        traces = traces_flat.reshape(*shape)
-        with self.visualizer.task_lock:
-            self.visualizer.viser_model.set_traces(traces, all_traces_rollout_size)
-
-    @on_event("INPUT", "plan_time")
-    def update_plan_time(self, event: dict) -> None:
-        """Callback to update plan time on receiving a new plan time measurement."""
-        plan_time_s = event["value"].to_numpy(zero_copy_only=False)[0]
-        # GUI element might not exist during task switch
-        if "plan_time_display" in self.visualizer.gui_elements:
-            self.visualizer.gui_elements["plan_time_display"].value = plan_time_s * 1000  # ms
-
-    def spin(self) -> None:
-        """Spin logic for the visualization node."""
-        try:
-            for event in self.node:
-                if self.visualizer.sim_pause_updated.is_set():
-                    self.write_sim_pause()
-                if self.visualizer.task_updated.is_set():
-                    self.write_task()
-                if self.visualizer.task_reset_updated.is_set():
-                    self.write_task_reset()
-                if self.visualizer.optimizer_updated.is_set():
-                    self.write_optimizer()
-                if self.visualizer.controller_config_updated.is_set():
-                    self.write_controller_config()
-                if self.visualizer.optimizer_config_updated.is_set():
-                    self.write_optimizer_config()
-                if self.visualizer.task_config_updated.is_set():
-                    self.write_task_config()
-
-                self.handle(event)
-        except KeyboardInterrupt:
-            pass
-
-    def cleanup(self) -> None:
-        """Cleanup the visualization node."""
-        self.visualizer.cleanup()
-        super().cleanup()
diff --git a/judo/controller/controller.py b/judo/controller/controller.py
index a5fcefe8..45e8e344 100644
--- a/judo/controller/controller.py
+++ b/judo/controller/controller.py
@@ -3,7 +3,7 @@
 import copy
 import warnings
 from dataclasses import dataclass
-from typing import Any, Literal
+from typing import TYPE_CHECKING, Any, Literal
 
 import numpy as np
 from mujoco import MjData
@@ -30,6 +30,9 @@
 from judo.utils.timer import Timer
 from judo.visualizers.utils import get_trace_sensors
 
+if TYPE_CHECKING:
+    from judo.utils.mjwarp_rollout_backend import MJWarpRolloutBackend
+
 RolloutBackendEntry = type[RolloutBackend]
 
 
@@ -61,7 +64,7 @@ def __init__(
         controller_config: ControllerConfig,
         task: Task,
         optimizer: Optimizer,
-        rollout_backend: str = "mujoco",
+        rollout_backend: "str | MJWarpRolloutBackend" = "mujoco",
         rollout_backend_registry: dict[str, RolloutBackendEntry] | None = None,
         rollout_backend_kwargs: dict[str, Any] | None = None,
     ) -> None:
@@ -71,7 +74,8 @@ def __init__(
             controller_config: The controller configuration.
             task: The task to use.
             optimizer: The optimizer to use.
-            rollout_backend: Name of the backend to use for rollouts (e.g., "mujoco", "mujoco_hierarchical").
+            rollout_backend: Name of the backend to use for rollouts (e.g., "mujoco", "mujoco_hierarchical"),
+                or a pre-built RolloutBackend instance to use directly (e.g., a shared MJWarpRolloutBackend).
             rollout_backend_registry: Optional mapping of backend names to backend classes.
                 Overrides entries in DEFAULT_ROLLOUT_BACKEND_REGISTRY.
             rollout_backend_kwargs: Optional extra kwargs for rollout backend constructor.
@@ -93,10 +97,14 @@ def __init__(
         self.model = self.task.model
 
         # Initialize rollout backend
-        self.rollout_backend: RolloutBackend = self._make_rollout_backend(
-            rollout_backend,
-            backend_kwargs=self._rollout_backend_kwargs,
-        )
+        if isinstance(rollout_backend, str):
+            self.rollout_backend: RolloutBackend = self._make_rollout_backend(
+                rollout_backend,
+                backend_kwargs=self._rollout_backend_kwargs,
+            )
+        else:
+            # Pre-built backend instance (e.g. shared MJWarpRolloutBackend from BatchedControllers)
+            self.rollout_backend = rollout_backend
         self._last_policy_output = (
             np.zeros((self.optimizer_cfg.num_rollouts, POLICY_OUTPUT_DIM))
             if isinstance(self.rollout_backend, HierarchicalMJRolloutBackend)
diff --git a/run_mpc/mpc_batch.py b/run_mpc/mpc_batch.py
index 1c9f367e..38a0bdda 100644
--- a/run_mpc/mpc_batch.py
+++ b/run_mpc/mpc_batch.py
@@ -9,24 +9,24 @@
 import numpy as np
 from tqdm import tqdm
 
-from judo.app.structs import MujocoState
+from judo.app.structs import MujocoState, RenderPose
 from judo.controller import BatchedControllers as JudoBatchedController
 from judo.controller import Controller as JudoController
+from judo.simulation.hierarchical_mj_simulation import HierarchicalMJSimulation
 from judo.simulation.mj_simulation import MJSimulation
-from judo.simulation.policy_mj_simulation import PolicyMJSimulation
 from judo.visualizers.visualizer import Visualizer
 from run_mpc.mpc_config import MPCTimers, PublicMPCConfig, SizeData
 
 
 def _get_previous_actions(sims: list[MJSimulation]) -> list[np.ndarray | None]:
     """Get previous actions from sims for hierarchical control sync."""
-    return [sim.last_policy_output if isinstance(sim, PolicyMJSimulation) else None for sim in sims]
+    return [sim.last_policy_output if isinstance(sim, HierarchicalMJSimulation) else None for sim in sims]
 
 
-def update_visualization(visualizer: Visualizer, sim_state: MujocoState, traces: np.ndarray) -> None:
-    """Update the viser visualization with current sim state and traces."""
-    visualizer.data.xpos[:] = sim_state.xpos
-    visualizer.data.xquat[:] = sim_state.xquat
+def update_visualization(visualizer: Visualizer, render_pose: RenderPose, traces: np.ndarray) -> None:
+    """Update the viser visualization with current render pose and traces."""
+    visualizer.data.xpos[:] = render_pose.xpos
+    visualizer.data.xquat[:] = render_pose.xquat
     visualizer.viser_model.set_data(visualizer.data)
     sensor_rollout_size = traces.shape[1]
     num_trace_sensors = traces.shape[2]
@@ -154,7 +154,7 @@ def run_mpc_batch(
     # Reset all simulations and controllers, sample new initial conditions + goals
     for sim, ctrl in zip(sims, controllers, strict=True):
         sim.task.reset()
-        if isinstance(sim, PolicyMJSimulation):
+        if isinstance(sim, HierarchicalMJSimulation):
             sim.reset_policy_state()
         ctrl.reset()
         ctrl.update_states(sim.sim_state)
@@ -202,7 +202,7 @@ def run_mpc_batch(
         timers.sim_step.toc()
 
         if vis is not None and num_parallel == 1:
-            update_visualization(vis, sims[0].sim_state, controllers[0].traces)
+            update_visualization(vis, sims[0].render_pose, controllers[0].traces)
             time.sleep(sims[0].timestep)
 
     results = storage.package_results(config.max_num_task_steps)
diff --git a/run_mpc/mpc_config.py b/run_mpc/mpc_config.py
index e0f1649f..e06b57a1 100644
--- a/run_mpc/mpc_config.py
+++ b/run_mpc/mpc_config.py
@@ -119,7 +119,8 @@ def load_configs_from_json_data(json_data: Any) -> tuple[JudoTask, Optimizer, Co
     optimizer_entry = available_optimizers.get(json_data["optimizer"])
     assert optimizer_entry is not None, f"Optimizer {json_data['optimizer']} is not registered!"
 
-    task_cls, task_config_cls = task_entry
+    task_cls = task_entry.task_type
+    task_config_cls = task_entry.task_config_type
     task_config: JudoTaskConfig = dacite.from_dict(task_config_cls, json_data["task_config"])
     task: JudoTask = task_cls()
     task.config = task_config
diff --git a/run_mpc/mpc_setup.py b/run_mpc/mpc_setup.py
index b161ad8c..f3b5b856 100644
--- a/run_mpc/mpc_setup.py
+++ b/run_mpc/mpc_setup.py
@@ -14,8 +14,8 @@
 from judo.controller import ControllerConfig
 from judo.controller.batched_spot_locomotion import BatchedSpotLocomotion
 from judo.optimizers import Optimizer
+from judo.simulation.hierarchical_mj_simulation import HierarchicalMJSimulation
 from judo.simulation.mj_simulation import MJSimulation
-from judo.simulation.policy_mj_simulation import PolicyMJSimulation
 from judo.tasks import Task as JudoTask
 from judo.utils.mjwarp_rollout_backend import MJWarpRolloutBackend
 from run_mpc.mpc_config import PublicMPCConfig, SizeData, make_size_data
@@ -58,7 +58,7 @@ def setup_mpc(
 
     for _ in range(num_parallel):
         if use_spot:
-            sim = PolicyMJSimulation(init_task=json_configs["task"])
+            sim = HierarchicalMJSimulation(init_task=json_configs["task"])
         else:
             sim = MJSimulation(init_task=json_configs["task"])
         sim.task.config = copy.deepcopy(task.config)
diff --git a/run_mpc/visualize_trajectories.py b/run_mpc/visualize_trajectories.py
index ac29971e..e1c933a0 100644
--- a/run_mpc/visualize_trajectories.py
+++ b/run_mpc/visualize_trajectories.py
@@ -56,8 +56,7 @@ def visualize_trajectory_batch(
     registered_tasks = get_registered_tasks()
     task_entry = registered_tasks.get(task)
     assert task_entry is not None, f"Task {task} is not registered!"
-    task_cls, _ = task_entry
-    task_instance = task_cls()
+    task_instance = task_entry.task_type()
     # Increase constraint buffers for contact-heavy scenes (e.g. Spot + tire).
     task_instance.spec.nconmax = 512
     task_instance.spec.njmax = 2048

From 3307be1ad0f1dca757fcd2e4f594f5d5197ba647 Mon Sep 17 00:00:00 2001
From: Duy Ta <dta@rai-inst.com>
Date: Tue, 5 May 2026 18:56:27 -0400
Subject: [PATCH 16/19] clean up

---
 judo/app/dora/visualization_node.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/judo/app/dora/visualization_node.py b/judo/app/dora/visualization_node.py
index 4b145951..462155ae 100644
--- a/judo/app/dora/visualization_node.py
+++ b/judo/app/dora/visualization_node.py
@@ -114,7 +114,6 @@ def update_states(self, event: dict) -> None:
                     break
             self.visualizer.optimizer_config_updated.set()
 
-        # TODO: change the mujoco state
         render_pose_msg = from_arrow(event["value"], event["metadata"], RenderPose)
         try:
             with self.visualizer.task_lock:

From e0d14bf04a5051be7f84516b6779afbf9d0a7760 Mon Sep 17 00:00:00 2001
From: Duy Ta <dta@rai-inst.com>
Date: Thu, 7 May 2026 15:08:18 -0400
Subject: [PATCH 17/19] clean up controllers rolloutbackend API

---
 judo/controller/controller.py        | 74 +++++++++++++++-------------
 judo/utils/mjwarp_rollout_backend.py |  4 +-
 judo/utils/rollout_backend.py        |  7 +++
 3 files changed, 49 insertions(+), 36 deletions(-)

diff --git a/judo/controller/controller.py b/judo/controller/controller.py
index 45e8e344..260b35af 100644
--- a/judo/controller/controller.py
+++ b/judo/controller/controller.py
@@ -3,7 +3,7 @@
 import copy
 import warnings
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any, Literal
+from typing import Any, Callable, Literal
 
 import numpy as np
 from mujoco import MjData
@@ -26,14 +26,11 @@
     make_normalizer,
     normalizer_registry,
 )
-from judo.utils.rollout_backend import RolloutBackend
+from judo.utils.rollout_backend import BatchedRolloutBackend, RolloutBackend
 from judo.utils.timer import Timer
 from judo.visualizers.utils import get_trace_sensors
 
-if TYPE_CHECKING:
-    from judo.utils.mjwarp_rollout_backend import MJWarpRolloutBackend
-
-RolloutBackendEntry = type[RolloutBackend]
+RolloutBackendEntry = type[RolloutBackend] | Callable[..., RolloutBackend]
 
 
 DEFAULT_ROLLOUT_BACKEND_REGISTRY: dict[str, RolloutBackendEntry] = {
@@ -64,7 +61,7 @@ def __init__(
         controller_config: ControllerConfig,
         task: Task,
         optimizer: Optimizer,
-        rollout_backend: "str | MJWarpRolloutBackend" = "mujoco",
+        rollout_backend: str = "mujoco",
         rollout_backend_registry: dict[str, RolloutBackendEntry] | None = None,
         rollout_backend_kwargs: dict[str, Any] | None = None,
     ) -> None:
@@ -74,8 +71,7 @@ def __init__(
             controller_config: The controller configuration.
             task: The task to use.
             optimizer: The optimizer to use.
-            rollout_backend: Name of the backend to use for rollouts (e.g., "mujoco", "mujoco_hierarchical"),
-                or a pre-built RolloutBackend instance to use directly (e.g., a shared MJWarpRolloutBackend).
+            rollout_backend: Name of the backend to use for rollouts (e.g., "mujoco", "mujoco_hierarchical").
             rollout_backend_registry: Optional mapping of backend names to backend classes.
                 Overrides entries in DEFAULT_ROLLOUT_BACKEND_REGISTRY.
             rollout_backend_kwargs: Optional extra kwargs for rollout backend constructor.
@@ -96,15 +92,10 @@ def __init__(
 
         self.model = self.task.model
 
-        # Initialize rollout backend
-        if isinstance(rollout_backend, str):
-            self.rollout_backend: RolloutBackend = self._make_rollout_backend(
-                rollout_backend,
-                backend_kwargs=self._rollout_backend_kwargs,
-            )
-        else:
-            # Pre-built backend instance (e.g. shared MJWarpRolloutBackend from BatchedControllers)
-            self.rollout_backend = rollout_backend
+        self.rollout_backend: RolloutBackend = self._make_rollout_backend(
+            rollout_backend,
+            backend_kwargs=self._rollout_backend_kwargs,
+        )
         self._last_policy_output = (
             np.zeros((self.optimizer_cfg.num_rollouts, POLICY_OUTPUT_DIM))
             if isinstance(self.rollout_backend, HierarchicalMJRolloutBackend)
@@ -455,8 +446,8 @@ def _make_rollout_backend(
         backend_kwargs: dict[str, Any] | None = None,
     ) -> RolloutBackend:
         """Instantiate a rollout backend from the merged backend registry."""
-        backend_cls = self._rollout_backend_registry.get(backend_name)
-        if backend_cls is None:
+        backend_factory = self._rollout_backend_registry.get(backend_name)
+        if backend_factory is None:
             raise ValueError(
                 f"Unknown rollout backend '{backend_name}'. "
                 "Provide it via rollout_backend_registry or choose a built-in backend."
@@ -494,11 +485,17 @@ def _make_rollout_backend(
                 )
             final_kwargs["policy_path"] = task_policy_path
 
-        return backend_cls(**final_kwargs)
+        backend = backend_factory(**final_kwargs)
+        if not isinstance(backend, RolloutBackend):
+            raise TypeError(
+                f"Rollout backend factory for '{backend_name}' must return a RolloutBackend, "
+                f"got {type(backend).__name__}."
+            )
+        return backend
 
 
 class BatchedControllers:
-    """Coordinates multiple controllers sharing a single RolloutBackend.
+    """Coordinates multiple controllers sharing a single BatchedRolloutBackend.
 
     This class manages batched rollouts across multiple controllers, executing
     a single GPU rollout for all controllers at each optimization iteration.
@@ -507,7 +504,7 @@ class BatchedControllers:
         # Create shared backend with num_threads per problem and num_problems
         num_rollouts = 64  # rollouts per controller
         num_problems = 3   # number of controllers
-        backend = RolloutBackend(model, num_threads=num_rollouts, num_problems=num_problems)
+        backend = BatchedRolloutBackend(model, num_threads=num_rollouts, num_problems=num_problems)
 
         # Create batched controller coordinator
         batched = BatchedControllers(config, task, optimizer, backend)
@@ -521,7 +518,7 @@ def __init__(
         controller_config: ControllerConfig,
         task: Task,
         optimizer: Optimizer,
-        rollout_backend: "MJWarpRolloutBackend",
+        rollout_backend: BatchedRolloutBackend,
     ) -> None:
         """Initialize the batched controllers.
 
@@ -529,7 +526,7 @@ def __init__(
             controller_config: Configuration for all controllers.
             task: Template task instance (new instances created from its class and model_path).
             optimizer: Template optimizer instance (deep copied for each controller).
-            rollout_backend: The shared WarpRolloutBackend instance. Should be initialized with
+            rollout_backend: Shared batched rollout backend instance. Should be initialized with
                 num_problems equal to len(controllers).
         """
         self.num_problems = rollout_backend.num_problems
@@ -537,11 +534,19 @@ def __init__(
         for _ in range(self.num_problems):
             new_task = task.__class__(model_path=task.model_path)
             new_task.config = copy.deepcopy(task.config)
+
+            # Construct controllers through normal backend-name resolution, but route
+            # the shared backend instance via a per-controller registry override.
+            shared_backend_name = "__shared_batched_backend__"
+            shared_backend_registry: dict[str, RolloutBackendEntry] = {
+                shared_backend_name: (lambda **_: rollout_backend)
+            }
             controller = Controller(
                 controller_config=controller_config,
                 task=new_task,
                 optimizer=copy.deepcopy(optimizer),
-                rollout_backend=rollout_backend,
+                rollout_backend=shared_backend_name,
+                rollout_backend_registry=shared_backend_registry,
             )
             self.controllers.append(controller)
         self.rollout_backend = rollout_backend
@@ -555,9 +560,9 @@ def __init__(
                 )
 
         # Validate num_problems matches number of controllers
-        if rollout_backend.num_problems != len(self.controllers):
+        if self.num_problems != len(self.controllers):
             raise ValueError(
-                f"RolloutBackend num_problems ({rollout_backend.num_problems}) does not match "
+                f"RolloutBackend num_problems ({self.num_problems}) does not match "
                 f"number of controllers ({len(self.controllers)}). "
                 f"Initialize backend with num_problems={len(self.controllers)}."
             )
@@ -656,8 +661,9 @@ def print_timer_stats(self) -> None:
         self.timer_rewards.print_stats()
         self.timer_update_iter.print_stats()
         self.timer_post_opt.print_stats()
-        if hasattr(self.rollout_backend, "print_timer_stats"):
-            self.rollout_backend.print_timer_stats()
+        backend_print_timer_stats = getattr(self.rollout_backend, "print_timer_stats", None)
+        if callable(backend_print_timer_stats):
+            backend_print_timer_stats()
 
     def reset_timers(self) -> None:
         """Reset all timers."""
@@ -666,8 +672,9 @@ def reset_timers(self) -> None:
         self.timer_rewards.reset()
         self.timer_update_iter.reset()
         self.timer_post_opt.reset()
-        if hasattr(self.rollout_backend, "reset_timers"):
-            self.rollout_backend.reset_timers()
+        backend_reset_timers = getattr(self.rollout_backend, "reset_timers", None)
+        if callable(backend_reset_timers):
+            backend_reset_timers()
 
     def update_states(self, state_msgs: list) -> None:
         """Update states for all controllers.
@@ -693,8 +700,7 @@ def set_init_previous_actions(self, previous_actions_list: list[np.ndarray | Non
         else:
             pa_np = np.stack([pa for pa in previous_actions_list if pa is not None], axis=0)
             pa_broadcast = np.repeat(pa_np, self.rollout_backend.num_threads, axis=0)
-            import warp as wp  # noqa: PLC0415
-
+            import warp as wp  # pyright: ignore[reportMissingImports]  # noqa: PLC0415
             self._last_policy_output = wp.array(pa_broadcast, dtype=wp.float32, device=self.rollout_backend.device)
 
 
diff --git a/judo/utils/mjwarp_rollout_backend.py b/judo/utils/mjwarp_rollout_backend.py
index 9642ff26..085f4bac 100644
--- a/judo/utils/mjwarp_rollout_backend.py
+++ b/judo/utils/mjwarp_rollout_backend.py
@@ -10,7 +10,7 @@
 from mujoco import MjData, MjModel
 
 from judo.controller.batched_spot_locomotion import BatchedSpotLocomotion
-from judo.utils.rollout_backend import RolloutBackend
+from judo.utils.rollout_backend import BatchedRolloutBackend
 from judo.utils.timer import Timer
 
 logger = logging.getLogger(__name__)
@@ -56,7 +56,7 @@ def target_frequency(self) -> float:
         return float("inf")
 
 
-class MJWarpRolloutBackend(RolloutBackend):
+class MJWarpRolloutBackend(BatchedRolloutBackend):
     """GPU-accelerated rollout backend using mujoco_warp.
 
     Supports two modes:
diff --git a/judo/utils/rollout_backend.py b/judo/utils/rollout_backend.py
index a3a1f144..4a59df2d 100644
--- a/judo/utils/rollout_backend.py
+++ b/judo/utils/rollout_backend.py
@@ -45,3 +45,10 @@ def update(self, num_threads: int) -> None:
         Args:
             num_threads: New number of parallel threads.
         """
+
+
+class BatchedRolloutBackend(RolloutBackend, ABC):
+    """Rollout backend base class for multi-problem batched execution."""
+
+    num_problems: int
+    device: str

From 085b2ef97140e90102b4e0f628c9edc82cdfe742 Mon Sep 17 00:00:00 2001
From: Duy Ta <dta@rai-inst.com>
Date: Thu, 7 May 2026 17:50:20 -0400
Subject: [PATCH 18/19] precommit

---
 judo/controller/controller.py | 1 +
 run_mpc/mpc_batch.py          | 2 +-
 2 files changed, 2 insertions(+), 1 deletion(-)

diff --git a/judo/controller/controller.py b/judo/controller/controller.py
index 260b35af..e73541bf 100644
--- a/judo/controller/controller.py
+++ b/judo/controller/controller.py
@@ -701,6 +701,7 @@ def set_init_previous_actions(self, previous_actions_list: list[np.ndarray | Non
             pa_np = np.stack([pa for pa in previous_actions_list if pa is not None], axis=0)
             pa_broadcast = np.repeat(pa_np, self.rollout_backend.num_threads, axis=0)
             import warp as wp  # pyright: ignore[reportMissingImports]  # noqa: PLC0415
+
             self._last_policy_output = wp.array(pa_broadcast, dtype=wp.float32, device=self.rollout_backend.device)
 
 
diff --git a/run_mpc/mpc_batch.py b/run_mpc/mpc_batch.py
index 38a0bdda..5cca2856 100644
--- a/run_mpc/mpc_batch.py
+++ b/run_mpc/mpc_batch.py
@@ -9,7 +9,7 @@
 import numpy as np
 from tqdm import tqdm
 
-from judo.app.structs import MujocoState, RenderPose
+from judo.app.structs import RenderPose
 from judo.controller import BatchedControllers as JudoBatchedController
 from judo.controller import Controller as JudoController
 from judo.simulation.hierarchical_mj_simulation import HierarchicalMJSimulation

From a2da9f51804ba79881740b5de266991dfbfaed9f Mon Sep 17 00:00:00 2001
From: Duy Ta <dta@rai-inst.com>
Date: Thu, 7 May 2026 18:02:34 -0400
Subject: [PATCH 19/19] review

---
 judo/app/dora/controller_node.py              | 5 ++++-
 judo/controller/controller.py                 | 7 +++++--
 judo/simulation/hierarchical_mj_simulation.py | 5 ++++-
 3 files changed, 13 insertions(+), 4 deletions(-)

diff --git a/judo/app/dora/controller_node.py b/judo/app/dora/controller_node.py
index c97cb010..872b066e 100644
--- a/judo/app/dora/controller_node.py
+++ b/judo/app/dora/controller_node.py
@@ -57,7 +57,10 @@ def _build_controller(self, task_name: str, optimizer_name: str) -> Controller:
         )
 
     def _current_optimizer_name(self) -> str:
-        """Look up the name of the current optimizer from the registry."""
+        """Look up the name of the current optimizer from the registry.
+
+        Returns "cem" as a safe default if no registry entry matches the active optimizer instance.
+        """
         for name, (cls, _) in self.controller.available_optimizers.items():
             if isinstance(self.controller.optimizer, cls):
                 return name
diff --git a/judo/controller/controller.py b/judo/controller/controller.py
index e73541bf..6f2a5a6c 100644
--- a/judo/controller/controller.py
+++ b/judo/controller/controller.py
@@ -460,7 +460,8 @@ def _make_rollout_backend(
         final_kwargs.update(backend_kwargs or {})
 
         if backend_name == "mujoco_hierarchical":
-            # physics_substeps must come from task, cannot be overridden in kwargs
+            # For the built-in hierarchical backend, physics_substeps is task-owned
+            # and should not be overridden at controller construction time.
             if "physics_substeps" in final_kwargs:
                 raise ValueError(
                     f"Cannot specify 'physics_substeps' in rollout_backend_kwargs. "
@@ -469,7 +470,9 @@ def _make_rollout_backend(
                 )
             final_kwargs["physics_substeps"] = self.task.physics_substeps
 
-            # policy_path must come from task registry, cannot be overridden in kwargs
+            # policy_path is currently registry-owned for hierarchical backend wiring.
+            # This keeps policy/model assumptions centralized until hierarchical ONNX
+            # integration is generalized on the C++ side.
             if "policy_path" in final_kwargs:
                 raise ValueError(
                     f"Cannot specify 'policy_path' in rollout_backend_kwargs. "
diff --git a/judo/simulation/hierarchical_mj_simulation.py b/judo/simulation/hierarchical_mj_simulation.py
index e72c160a..483e8f1f 100644
--- a/judo/simulation/hierarchical_mj_simulation.py
+++ b/judo/simulation/hierarchical_mj_simulation.py
@@ -157,7 +157,10 @@ def set_task(self, task_name: str) -> None:
             self._init_cpp_systems(policy_path)
             self._last_policy_output = np.zeros(POLICY_OUTPUT_DIM)
         else:
-            self._systems = None
+            raise ValueError(
+                f"Task '{self.task.name}' does not use a locomotion policy. "
+                "Use MJSimulation instead of HierarchicalMJSimulation for this task."
+            )
 
     @property
     def last_policy_output(self) -> np.ndarray: