Changes for Sumo

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
@@ -23,10 +24,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: Callable | None = 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,
@@ -36,25 +41,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()
-                optimizer = self.controller.optimizer_cls(self.controller.optimizer_config_cls(), task.nu)
-                self.controller = Controller(
-                    controller_config=self.controller.controller_cfg,
-                    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."""
@@ -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

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,41 @@ 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  # 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."""
-        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:
@@ -85,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:

judo/app/dora/visualization.py

@@ -8,7 +8,8 @@
 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
 
 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:
@@ -85,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:
@@ -96,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

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

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

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
@@ -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: Any,
 ) -> 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,
     )

judo/simulation/__init__.py

@@ -4,19 +4,36 @@
 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() -> 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,
 }
 
 
 def get_simulation_backend(simulation_backend: str) -> type:
-    """Get the simulation class for a given backend."""
-    return simulation_registry[simulation_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]()
 
 
 __all__ = [
     "Simulation",
     "MJSimulation",
     "PolicyMJSimulation",
+    "get_simulation_backend",
 ]

judo/simulation/base.py

@@ -5,6 +5,7 @@
 import numpy as np
 from omegaconf import DictConfig
 
+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
@@ -54,6 +55,26 @@ 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,  # type: ignore
+            qpos=self.task.data.qpos,  # type: ignore
+            qvel=self.task.data.qvel,  # 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."""
+        return self.task.dt