cosmetic changes

Author: dta-bdai

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:

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()

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:

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()

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

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