Refactor: decouple core libraries from the default task set and make judo reusable as a third-party library

.github/workflows/pypi_publish.yml

@@ -164,7 +164,7 @@ jobs:
         run: |
           source /tmp/test-venv/bin/activate
           python -c "
-          from mujoco_extensions.policy_rollout import create_systems_vector, threaded_rollout
+          from judo.mujoco_extensions.policy_rollout import create_systems_vector, threaded_rollout
           print('mujoco_extensions OK')
           from judo.tasks.spot.spot_base import SpotBase, XML_PATH
           from judo.tasks.spot.spot_constants import SPOT_LOCOMOTION_POLICY_PATH

.github/workflows/test.yml

@@ -42,9 +42,9 @@ jobs:
         uses: actions/cache@v5
         with:
           path: |
-            mujoco_extensions/build
-            mujoco_extensions/policy_rollout/*.so
-          key: mujoco-ext-${{ matrix.os }}-${{ matrix.pixi-env }}-${{ hashFiles('pixi.lock', 'mujoco_extensions/**/*.cpp', 'mujoco_extensions/**/*.h', 'mujoco_extensions/**/CMakeLists.txt') }}
+            judo/mujoco_extensions/build
+            judo/mujoco_extensions/policy_rollout/*.so
+          key: mujoco-ext-${{ matrix.os }}-${{ matrix.pixi-env }}-${{ hashFiles('pixi.lock', 'judo/mujoco_extensions/**/*.cpp', 'judo/mujoco_extensions/**/*.h', 'judo/mujoco_extensions/**/CMakeLists.txt') }}
 
       - name: Build mujoco_extensions (cache miss)
         if: steps.ext-cache.outputs.cache-hit != 'true'

.github/workflows/test_gpu.yml

@@ -29,9 +29,9 @@ jobs:
         uses: actions/cache@v5
         with:
           path: |
-            mujoco_extensions/build
-            mujoco_extensions/policy_rollout/*.so
-          key: mujoco-ext-gpu-${{ hashFiles('pixi.lock', 'mujoco_extensions/**/*.cpp', 'mujoco_extensions/**/*.h', 'mujoco_extensions/**/CMakeLists.txt') }}
+            judo/mujoco_extensions/build
+            judo/mujoco_extensions/policy_rollout/*.so
+          key: mujoco-ext-gpu-${{ hashFiles('pixi.lock', 'judo/mujoco_extensions/**/*.cpp', 'judo/mujoco_extensions/**/*.h', 'judo/mujoco_extensions/**/CMakeLists.txt') }}
 
       - name: Build mujoco_extensions (cache miss)
         if: steps.ext-cache.outputs.cache-hit != 'true'

.gitignore

@@ -188,7 +188,7 @@ outputs/
 MUJOCO_LOG.TXT
 
 # bundled shared libraries (auditwheel handles bundling into wheel)
-mujoco_extensions/policy_rollout/libonnxruntime*
+judo/mujoco_extensions/policy_rollout/libonnxruntime*
 wheelhouse/
 
 # assets

README.md

@@ -71,6 +71,18 @@ pixi run clean
 pixi run build
 ```
 
+If you change the `policy_rollout_pybind` C++ API (signatures in
+`judo/mujoco_extensions/policy_rollout/pybind/`), regenerate its type stubs so
+type checkers stay in sync (this builds the module first):
+```bash
+pixi run stubgen-extension
+```
+The generated `.pyi` files live next to the compiled module in
+`judo/mujoco_extensions/policy_rollout/policy_rollout_pybind/` and are committed
+to the repo. `pybind11-stubgen` occasionally emits raw C++ types (e.g.
+`SystemClass::System`, `Ort::Session`) as invalid `...` annotations — review the
+diff and replace those with the proper Python types before committing.
+
 ## 2. Run the `judo` app!
 To start the simulator, from within the pixi shell, you can simply run:
 ```bash

judo/app/__init__.py

@@ -1 +1,81 @@
 # Copyright (c) 2025 Robotics and AI Institute LLC. All rights reserved.
+
+"""Judo application defaults.
+
+This package is judo's *batteries-included* task set. It is intentionally kept out of the
+individual libraries (tasks, optimizers, controllers, simulation) so those libraries can be reused
+by third-party packages that want to register their own tasks. Its sole job is to declare judo's
+default tasks and how each one is wired (rollout backend, simulation backend, and low-level
+locomotion policy).
+
+Importing ``judo.app`` registers judo's built-in tasks via :func:`register_default_tasks`.
+Third-party applications that only want the *mechanism* (not judo's default tasks) should avoid
+importing this module and instead call :func:`judo.tasks.register_task` themselves. Applications
+that want judo's defaults *and* their own can import this module (or call
+:func:`register_default_tasks`) and then register/override additional tasks by name.
+"""
+
+from judo.tasks import register_task
+from judo.tasks.caltech_leap_cube import CaltechLeapCube, CaltechLeapCubeConfig
+from judo.tasks.cartpole import Cartpole, CartpoleConfig
+from judo.tasks.cylinder_push import CylinderPush, CylinderPushConfig
+from judo.tasks.fr3_pick import FR3Pick, FR3PickConfig
+from judo.tasks.leap_cube import LeapCube, LeapCubeConfig
+from judo.tasks.leap_cube_down import LeapCubeDown, LeapCubeDownConfig
+from judo.tasks.spot import (
+    SpotBase,
+    SpotBaseConfig,
+    SpotBoxPush,
+    SpotBoxPushConfig,
+    SpotNavigate,
+    SpotNavigateConfig,
+    SpotTireRoll,
+    SpotTireRollConfig,
+    SpotTireUpright,
+    SpotTireUprightConfig,
+)
+from judo.tasks.spot.spot_constants import SPOT_LOCOMOTION_POLICY_PATH
+
+
+def register_default_tasks() -> None:
+    """Register judo's built-in tasks with their default backends and locomotion policies.
+
+    This is the application-level source of truth for which tasks are available and how each one
+    is wired (rollout backend, simulation backend, and low-level locomotion policy). Third-party
+    applications can call :func:`judo.tasks.register_task` (or
+    :func:`judo.registration.register_tasks_from_cfg`) to register their own tasks instead of, or
+    in addition to, these defaults. Re-registering a name overrides the previous entry, so callers
+    may register these defaults and then override individual tasks by name.
+    """
+    register_task(CylinderPush.name, CylinderPush, CylinderPushConfig)
+    register_task(Cartpole.name, Cartpole, CartpoleConfig)
+    register_task(FR3Pick.name, FR3Pick, FR3PickConfig)
+    register_task(LeapCube.name, LeapCube, LeapCubeConfig)
+    register_task(LeapCubeDown.name, LeapCubeDown, LeapCubeDownConfig)
+    register_task(CaltechLeapCube.name, CaltechLeapCube, CaltechLeapCubeConfig)
+
+    spot_policy_path = str(SPOT_LOCOMOTION_POLICY_PATH)
+    for spot_task, spot_config in (
+        (SpotBase, SpotBaseConfig),
+        (SpotBoxPush, SpotBoxPushConfig),
+        (SpotNavigate, SpotNavigateConfig),
+        (SpotTireRoll, SpotTireRollConfig),
+        (SpotTireUpright, SpotTireUprightConfig),
+    ):
+        register_task(
+            spot_task.name,
+            spot_task,
+            spot_config,
+            rollout_backend="mujoco_hierarchical",
+            simulation_backend="mujoco_hierarchical",
+            locomotion_policy_path=spot_policy_path,
+        )
+
+
+# Register judo's built-in tasks when the application defaults are imported.
+register_default_tasks()
+
+
+__all__ = [
+    "register_default_tasks",
+]

judo/app/dora/controller_node.py

@@ -9,8 +9,10 @@
 from dora_utils.node import DoraNode, on_event
 from omegaconf import DictConfig
 
-from judo.app.structs import MujocoState
 from judo.controller import Controller, make_controller
+from judo.optimizers import get_registered_optimizers
+from judo.structs import MujocoState
+from judo.tasks import get_registered_tasks
 
 
 class ControllerNode(DoraNode):
@@ -44,6 +46,7 @@ def __init__(
         self._optimizer_registration_cfg = optimizer_registration_cfg
         self.controller = self._build_controller(init_task, init_optimizer)
         self._paused = False
+        self._running = True
         self.write_controls()
         self.lock = Lock()
 
@@ -61,7 +64,7 @@ def _current_optimizer_name(self) -> str:
 
         Returns "cem" as a safe default if no registry entry matches the active optimizer instance.
         """
-        for name, (cls, _) in self.controller.available_optimizers.items():
+        for name, (cls, _) in get_registered_optimizers().items():
             if isinstance(self.controller.optimizer, cls):
                 return name
         return "cem"
@@ -70,7 +73,7 @@ def _current_optimizer_name(self) -> str:
     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)
+        task_entry = get_registered_tasks().get(new_task)
         if task_entry is None:
             raise ValueError(f"Task {new_task} not found in task registry.")
 
@@ -94,7 +97,7 @@ def set_paused_status(self, event: dict) -> None:
     def update_optimizer(self, event: dict) -> None:
         """Updates the optimizer type."""
         new_optimizer = event["value"].to_numpy(zero_copy_only=False)[0]
-        optimizer_entry = self.controller.available_optimizers.get(new_optimizer)
+        optimizer_entry = get_registered_optimizers().get(new_optimizer)
         if optimizer_entry is not None:
             optimizer_cls, optimizer_config_cls = optimizer_entry
             optimizer_config = optimizer_config_cls()
@@ -120,19 +123,33 @@ def update_task_config(self, event: dict) -> None:
         """Callback to update optimizer task config on receiving a new config message."""
         self.controller.task_config = from_event(event, type(self.controller.task_config))
 
+    def _send_output(self, output_id: str, data: pa.Array, metadata: dict | None = None) -> None:
+        """Publish an output, tolerating event-stream closure during shutdown.
+
+        During Ctrl-C / dataflow teardown the dora event stream can close while a step is in
+        progress, which makes ``send_output`` raise ``RuntimeError``. Treat that as a stop signal
+        so shutdown stays quiet instead of surfacing a traceback.
+        """
+        if not self._running:
+            return
+        try:
+            self.node.send_output(output_id, data, metadata or {})
+        except RuntimeError:
+            self._running = False
+
     def write_controls(self) -> None:
         """Util that publishes the current controller spline."""
         # send control action
         arr, metadata = to_arrow(self.controller.spline_data)
-        self.node.send_output("controls", arr, metadata)
+        self._send_output("controls", arr, metadata)
 
         # send traces
         if self.controller.traces is not None and len(self.controller.traces) > 0:
             metadata = {
                 "all_traces_rollout_size": str(self.controller.all_traces_rollout_size),
                 "shape": self.controller.traces.shape,
             }
-            self.node.send_output("traces", pa.array(self.controller.traces.flatten()), metadata=metadata)
+            self._send_output("traces", pa.array(self.controller.traces.flatten()), metadata=metadata)
 
     @on_event("INPUT", "states")
     def update_states(self, event: dict) -> None:
@@ -163,16 +180,18 @@ def step(self) -> None:
             self.controller.update_action()
             end = time.perf_counter()
 
-        self.node.send_output("plan_time", pa.array([end - start]))
+        self._send_output("plan_time", pa.array([end - start]))
         self.write_controls()
 
     def spin(self) -> None:
         """Spin logic for the controller node."""
         try:
-            while True:
+            while self._running:
                 start_time = time.time()
                 self.parse_messages()
                 self.step()
+                if not self._running:
+                    break
 
                 # Force controller to run at fixed rate specified by control_freq.
                 sleep_dt = 1 / self.controller.controller_cfg.control_freq - (time.time() - start_time)

judo/app/dora/simulation_node.py

@@ -8,9 +8,9 @@
 from dora_utils.node import DoraNode, on_event
 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.structs import SplineData
 from judo.tasks import get_registered_tasks
 
 
@@ -56,7 +56,12 @@ def _init_sim(self, task_name: str) -> None:
             raise ValueError(f"Task {task_name} not found in task registry.")
 
         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)
+        sim_kwargs: dict = {"init_task": task_name, "task_registration_cfg": self._task_registration_cfg}
+        # The hierarchical backend needs the low-level policy path passed explicitly (the
+        # Simulation classes are decoupled from the task registry).
+        if task_entry.simulation_backend == "mujoco_hierarchical":
+            sim_kwargs["locomotion_policy_path"] = task_entry.locomotion_policy_path
+        self.sim = sim_backend_cls(**sim_kwargs)
 
     @on_event("INPUT", "task")
     def update_task(self, event: dict) -> None:

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 RenderPose
+from judo.structs import RenderPose
 from judo.tasks import TaskRegistration
 from judo.visualizers.visualizer import Visualizer
 

judo/cli.py

@@ -119,7 +119,11 @@ def main_app(cfg: DictConfig) -> None:
     """Main function to run judo via a hydra configuration yaml file."""
     try:
         run(cfg)
-    except (KeyboardInterrupt, SystemExit):
+    except (KeyboardInterrupt, SystemExit, subprocess.CalledProcessError):
+        # On shutdown (e.g. Ctrl+C), dora_utils' cleanup (run.py) calls `dora destroy` with
+        # check=True; if the dataflow is already being torn down it exits non-zero and raises
+        # CalledProcessError. Our own _force_cleanup() already handles teardown, so all of these
+        # are expected shutdown paths rather than real failures.
         _force_cleanup()
 
 
@@ -144,7 +148,7 @@ def _warm_caches() -> None:
 def _require_mujoco_extensions() -> None:
     """Fail fast if mujoco_extensions is unavailable in the current environment."""
     try:
-        import mujoco_extensions  # noqa: F401, PLC0415
+        import judo.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"

judo/controller/__init__.py

@@ -4,8 +4,8 @@
     BatchedControllers,
     Controller,
     ControllerConfig,
-    make_controller,
 )
+from judo.controller.factory import make_controller
 from judo.controller.overrides import (
     set_default_caltech_leap_cube_overrides,
     set_default_cartpole_overrides,