adding features and tasks for the Spot-related tasks

.gitignore

@@ -188,5 +188,8 @@ outputs/
 MUJOCO_LOG.TXT
 
 # assets
-*.obj
-*.stl
+# *.obj
+# *.stl
+
+.vscode/
+*.h5

CHANGELOG.md

@@ -1,3 +1,34 @@
+# Unreleased
+
+## Added
+* C++ rollout implementations for MuJoCo simulations (@johnzhang3, @jbruedigam-bdai, #95)
+    * Implemented high-performance C++ backend for MuJoCo rollouts with significant speedup
+    * Added custom C++ rollout functions for Spot robot tasks with ONNX policy inference
+    * Separated task-space control dimensions from simulation control dimensions for better flexibility
+* Added `cutoff_time` parameter to `OptimizerConfig` (@johnzhang3, #95)
+    * Allows specifying maximum wall-clock time (in seconds) per rollout thread before termination
+    * Prevents single threads from blocking optimization when rollouts take too long
+    * Default value: 0.2 seconds
+* New Spot robot tasks (@johnzhang3, #95)
+    * `spot_yellow_chair`: Manipulation task for pushing a yellow office chair
+    * `spot_yellow_chair_ramp`: Chair manipulation with ramp navigation
+    * `spot_box`: Box manipulation task for Spot
+* CMAES (Covariance Matrix Adaptation Evolution Strategy) optimizer (@alberthli, #95)
+
+## Fixed
+* Fixed reward sign in `spot_yellow_chair_ramp` task where torso proximity was incorrectly penalizing closeness (@johnzhang-rai, #95)
+* Updated `Task.success()` and `Task.failure()` methods to raise `NotImplementedError` instead of returning `False` for clearer API expectations (@johnzhang-rai, #95)
+
+## Changed
+* Renamed policy file from `xinghao_policy_v1.onnx` to `spot_relic_policy.onnx` for better clarity (@johnzhang-rai, #95)
+
+## Documentation
+* Added detailed documentation for `cutoff_time` parameter in optimizer configuration (@johnzhang-rai, #95)
+
+## Dev
+* Removed commented-out code and cleaned up task implementations (@johnzhang-rai, #95)
+* Added C++ build instructions to README for development setup (@johnzhang-rai, #95)
+
 # v0.0.5
 
 ## Added

README.md

@@ -35,6 +35,8 @@ We recommend installing `judo` using `pip` as follows:
 pip install judo-rai  # if you want dev dependencies, use judo-rai[dev]
 ```
 
+Prebuilt wheels with C++ extensions are provided for Linux and macOS (both Intel and Apple Silicon) for Python 3.10+.
+
 ### Developers
 #### Conda
 For developers, run the following commands from this folder after cloning:
@@ -58,6 +60,9 @@ To create our environment (and activate it each time later), run the following i
 # every time you want to activate
 pixi shell -e dev
 
+# build judo cpp bindings
+pixi run build
+
 # first time only
 pre-commit install
 pybind11-stubgen mujoco -o typings/
@@ -68,6 +73,12 @@ To start the simulator, you can simply run:
 ```bash
 judo
 ```
+
+To start with a specific config file
+```bash
+pixi run -e dev judo --config-name judo_default
+```
+
 This will start the stack and print a link in the terminal that will open the app in your browser, e.g.,
 ```
 http://localhost:8080

example_configs/__init__.py

@@ -0,0 +1,3 @@
+# Copyright (c) 2025 Robotics and AI Institute LLC. All rights reserved.
+
+# Example configurations for judo

judo/app/dora/controller.py

@@ -45,11 +45,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)
+                optimizer_config = self.controller.optimizer_config_cls()
+                optimizer = self.controller.optimizer_cls(optimizer_config, task.nu)
                 self.controller = Controller(
                     controller_config=self.controller.controller_cfg,
                     task=task,
                     optimizer=optimizer,
+                    optimizer_config=optimizer_config,
                 )
                 self.write_controls()
         else:

judo/controller/__init__.py

@@ -9,6 +9,8 @@
     set_default_fr3_pick_overrides,
     set_default_leap_cube_down_overrides,
     set_default_leap_cube_overrides,
+    set_default_spot_yellow_chair_overrides,
+    set_default_spot_yellow_chair_ramp_overrides,
 )
 from judo.utils.assets import download_and_extract_meshes
 
@@ -24,6 +26,8 @@
 set_default_fr3_pick_overrides()
 set_default_leap_cube_overrides()
 set_default_leap_cube_down_overrides()
+set_default_spot_yellow_chair_overrides()
+set_default_spot_yellow_chair_ramp_overrides()
 
 __all__ = [
     "Controller",

judo/controller/controller.py

@@ -14,7 +14,7 @@
 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.utils.mujoco import RolloutBackend, make_model_data_pairs
+from judo.utils.mujoco import make_model_data_pairs
 from judo.utils.normalization import (
     IdentityNormalizer,
     Normalizer,
@@ -37,6 +37,7 @@ class ControllerConfig(OverridableConfig):
     max_opt_iters: int = 1
     max_num_traces: int = 5
     action_normalizer: Literal["none", "min_max", "running"] = "none"
+    rollout_backend: Literal["mujoco", "mujoco_spot", "mujoco_cpp"] | None = None
 
 
 class Controller:
@@ -47,15 +48,18 @@ def __init__(
         controller_config: ControllerConfig,
         task: Task,
         optimizer: Optimizer,
-        rollout_backend: Literal["mujoco"] = "mujoco",
+        optimizer_config: OptimizerConfig,
+        rollout_backend: Literal["mujoco", "mujoco_spot", "mujoco_cpp"] | None = None,
     ) -> None:
         """Initialize the controller.
 
         Args:
-            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.
+            controller_config: The configuration for the controller.
+            task: The Task object that specifies the environment.
+            task_config: The configuration for the task.
+            optimizer: The optimizer object that will be used for optimization.
+            optimizer_config: The configuration for the optimizer.
+            rollout_backend: The backend to use for rollouts. If None, uses task's default or "mujoco".
         """
         self._controller_cfg = controller_config
         self.task = task
@@ -67,7 +71,18 @@ def __init__(
         self.model = self.task.model
         self.model_data_pairs = make_model_data_pairs(self.model, self.optimizer_cfg.num_rollouts)
 
-        self.rollout_backend = RolloutBackend(num_threads=self.optimizer_cfg.num_rollouts, backend=rollout_backend)
+        # Determine backend: config override > task default > "mujoco"
+        backend: Literal["mujoco", "mujoco_spot", "mujoco_cpp"] = (
+            controller_config.rollout_backend or rollout_backend or getattr(task, "default_backend", "mujoco")
+        )
+
+        self.rollout_backend = task.RolloutBackend(
+            num_threads=self.optimizer_cfg.num_rollouts,
+            backend=backend,
+            task_to_sim_ctrl=task.task_to_sim_ctrl,
+            cutoff_time=self.optimizer_cfg.cutoff_time,
+        )
+
         self.action_normalizer = self._init_action_normalizer()
 
         # a container for any metadata from the system that we want to pass to the task
@@ -76,7 +91,7 @@ def __init__(
         self.states = np.zeros((self.optimizer_cfg.num_rollouts, self.num_timesteps, self.model.nq + self.model.nv))
         self.current_state = np.concatenate([self.task.data.qpos, self.task.data.qvel])
         self.sensors = np.zeros((self.optimizer_cfg.num_rollouts, self.num_timesteps, self.model.nsensordata))
-        self.rollout_controls = np.zeros((self.optimizer_cfg.num_rollouts, self.num_timesteps, self.model.nu))
+        self.rollout_controls = np.zeros((self.optimizer_cfg.num_rollouts, self.num_timesteps, self.task.nu))
         self.rewards = np.zeros((self.optimizer_cfg.num_rollouts,))
         self.reset()
 
@@ -205,7 +220,12 @@ def update_action(self) -> None:
         # resizing any variables due to changes in the GUI
         if len(self.model_data_pairs) != self.optimizer_cfg.num_rollouts:
             self.model_data_pairs = make_model_data_pairs(self.model, self.optimizer_cfg.num_rollouts)
-            self.rollout_backend.update(self.optimizer_cfg.num_rollouts)
+            self.rollout_backend.update(self.optimizer_cfg.num_rollouts, cutoff_time=self.optimizer_cfg.cutoff_time)
+
+        # Update rollout backend when cutoff_time changes
+        if self.rollout_backend.cutoff_time != self.optimizer_cfg.cutoff_time:
+            self.rollout_backend.cutoff_time = self.optimizer_cfg.cutoff_time
+            self.rollout_backend.update(self.optimizer_cfg.num_rollouts, cutoff_time=self.optimizer_cfg.cutoff_time)
 
         normalizer_cls = normalizer_registry.get(self.action_normalizer_type)
         if normalizer_cls is None:
@@ -232,8 +252,8 @@ def update_action(self) -> None:
             candidate_knots_normalized = self.optimizer.sample_control_knots(nominal_knots_normalized)
             candidate_knots_normalized = np.clip(
                 candidate_knots_normalized,
-                self.action_normalizer.normalize(self.task.actuator_ctrlrange[:, 0]),
-                self.action_normalizer.normalize(self.task.actuator_ctrlrange[:, 1]),
+                self.action_normalizer.normalize(self.task.ctrlrange[:, 0]),
+                self.action_normalizer.normalize(self.task.ctrlrange[:, 1]),
             )
             self.candidate_knots = self.action_normalizer.denormalize(candidate_knots_normalized)
 
@@ -346,11 +366,11 @@ def _init_action_normalizer(self) -> Normalizer:
         """Initialize the action normalizer."""
         action_normalizer_kwargs = {}
         if self.action_normalizer_type == "min_max":
-            action_normalizer_kwargs["min"] = self.task.actuator_ctrlrange[:, 0]
-            action_normalizer_kwargs["max"] = self.task.actuator_ctrlrange[:, 1]
+            action_normalizer_kwargs["min"] = self.task.ctrlrange[:, 0]
+            action_normalizer_kwargs["max"] = self.task.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:
@@ -401,7 +421,8 @@ def make_controller(
     task = task_cls()
 
     optimizer_cls, optimizer_config_cls = optimizer_entry
-    optimizer = optimizer_cls(optimizer_config_cls(), task.nu)
+    optimizer_config = optimizer_config_cls()
+    optimizer = optimizer_cls(optimizer_config, task.nu)
 
     controller_cfg = ControllerConfig()
     controller_cfg.set_override(init_task)
@@ -410,5 +431,6 @@ def make_controller(
         controller_config=controller_cfg,
         task=task,
         optimizer=optimizer,
+        optimizer_config=optimizer_config,
         rollout_backend=rollout_backend,
     )

judo/controller/overrides.py

@@ -79,3 +79,25 @@ def set_default_fr3_pick_overrides() -> None:
             "control_freq": 20.0,
         },
     )
+
+
+def set_default_spot_yellow_chair_overrides() -> None:
+    """Sets the default task-specific controller config overrides for the spot locomotion task."""
+    set_config_overrides(
+        "spot_yellow_chair",
+        ControllerConfig,
+        {
+            "horizon": 3.0,
+        },
+    )
+
+
+def set_default_spot_yellow_chair_ramp_overrides() -> None:
+    """Sets the default task-specific controller config overrides for the spot yellow chair ramp task."""
+    set_config_overrides(
+        "spot_yellow_chair_ramp",
+        ControllerConfig,
+        {
+            "horizon": 3.0,
+        },
+    )