rai-opensource · dta-bdai · Mar 3, 2026 · Feb 17, 2026 · Feb 17, 2026 · Feb 17, 2026
diff --git a/.github/workflows/pypi_publish.yml b/.github/workflows/pypi_publish.yml
@@ -140,6 +140,8 @@ jobs:
 
     steps:
       - uses: actions/checkout@v6
+        with:
+          sparse-checkout: tests
 
       - uses: actions/setup-python@v5
         with:
@@ -149,10 +151,14 @@ jobs:
         run: |
           python -m venv /tmp/test-venv
           source /tmp/test-venv/bin/activate
-          pip install --index-url https://test.pypi.org/simple/ \
-                      --extra-index-url https://pypi.org/simple/ \
-                      judo-rai==${{ needs.publish-testpypi.outputs.dev_version }}
-          pip install pytest
+          for i in 1 2 3 4 5; do
+            pip install --index-url https://test.pypi.org/simple/ \
+                        --extra-index-url https://pypi.org/simple/ \
+                        judo-rai==${{ needs.publish-testpypi.outputs.dev_version }} && break
+            echo "Attempt $i failed, waiting 30s for TestPyPI index to update..."
+            sleep 30
+          done
+          pip install pytest pytest-xdist
 
       - name: Smoke test — imports and C++ extension
         run: |

@@ -8,7 +8,7 @@
     set_default_fr3_pick_overrides,
     set_default_leap_cube_down_overrides,
     set_default_leap_cube_overrides,
-    set_default_spot_tire_upright_overrides,
+    set_default_spot_overrides,
 )
 
 set_default_caltech_leap_cube_overrides()
@@ -17,7 +17,7 @@
 set_default_fr3_pick_overrides()
 set_default_leap_cube_overrides()
 set_default_leap_cube_down_overrides()
-set_default_spot_tire_upright_overrides()
+set_default_spot_overrides()
 
 __all__ = [
     "Controller",

@@ -67,15 +67,25 @@ def set_default_caltech_leap_cube_overrides() -> None:
     )
 
 
-def set_default_spot_tire_upright_overrides() -> None:
-    """Sets the default task-specific controller config overrides for the spot_tire_upright task."""
-    set_config_overrides(
-        "spot_tire_upright",
-        ControllerConfig,
-        {
-            "horizon": 2.0,
-        },
-    )
+_SPOT_TASK_NAMES = [
+    "spot_base",
+    "spot_box_push",
+    "spot_navigate",
+    "spot_tire_roll",
+    "spot_tire_upright",
+]
+
+
+def set_default_spot_overrides() -> None:
+    """Sets the default task-specific controller config overrides for all Spot tasks."""
+    for task_name in _SPOT_TASK_NAMES:
+        set_config_overrides(
+            task_name,
+            ControllerConfig,
+            {
+                "horizon": 2.0,
+            },
+        )
 
 
 def set_default_fr3_pick_overrides() -> None:

@@ -19,7 +19,7 @@
 
 DEFAULT_SLIDER_STEP_FLOAT = 0.01
 DEFAULT_SLIDER_STEP_INT = 1
-GOAL_RADIUS = 0.05
+GOAL_RADIUS = 0.15
 
 
 def slider(

@@ -0,0 +1,51 @@
+<mujoco model="spot_box">
+<!-- Copyright (c) 2025 Robotics and AI Institute LLC. All rights reserved. -->
+<!-- Scene: Spot robot with box object for push tasks -->
+
+  <compiler angle="radian" meshdir="../../meshes/" assetdir="../../meshes/" autolimits="true"/>
+  <option timestep="0.01" solver="Newton" integrator="implicitfast" density="1"/>
+
+  <include file="../spot_primitive/default.xml" />
+  <include file="../spot_primitive/assets.xml" />
+
+  <worldbody>
+    <light directional="true" diffuse=".4 .4 .4" specular="0.1 0.1 0.1" pos="0 0 5" dir="0 0 -1" castshadow="false"/>
+    <light directional="true" diffuse=".6 .6 .6" specular="0.2 0.2 0.2" pos="0 0 4" dir="0 0 -1"/>
+
+    <geom name="ground" type="plane" size="10 10 0.01" class="collision" priority="5" friction="0.8" material="blue_grid"/>
+
+    <body name="body" pos="0 0 0.52">
+      <include file="../spot_primitive/body.xml" />
+      <include file="../spot_primitive/legs.xml" />
+      <include file="../spot_primitive/arm.xml" />
+    </body>
+
+    <!-- Box object -->
+    <body name="box_body" pos="2 0 0.254">
+      <joint name="box_joint" type="free"/>
+      <inertial pos="0 0 0" mass="1.5" diaginertia="0.1445 0.1445 0.1445"/>
+      <geom name="box_collision" type="box" size="0.254 0.254 0.254" class="collision" priority="6"/>
+      <geom name="box_visual" type="box" size="0.254 0.254 0.254" class="visual" rgba="0.6 0.3 0.1 1"/>
+      <site name="site_object" pos="0 0 0" size="0.01"/>
+    </body>
+  </worldbody>
+
+  <include file="../spot_primitive/actuator.xml" />
+  <include file="../spot_primitive/contact.xml" />
+
+  <sensor>
+    <framepos name="sensor_body" objtype="site" objname="site_body"/>
+    <framexaxis name="body_x_axis" objtype="site" objname="site_body"/>
+    <framexaxis name="object_x_axis" objtype="site" objname="site_object"/>
+    <frameyaxis name="object_y_axis" objtype="site" objname="site_object"/>
+    <framezaxis name="object_z_axis" objtype="site" objname="site_object"/>
+    <framepos name="trace_fngr_site" objtype="site" objname="site_arm_link_fngr"/>
+    <framexaxis name="gripper_x_axis" objtype="site" objname="site_arm_link_wr1"/>
+    <frameyaxis name="gripper_y_axis" objtype="site" objname="site_arm_link_wr1"/>
+    <framexaxis name="finger_x_axis" objtype="site" objname="site_arm_link_fngr"/>
+    <framepos name="fl_pos" objtype="site" objname="site_front_left"/>
+    <framepos name="fr_pos" objtype="site" objname="site_front_right"/>
+    <framepos name="hl_pos" objtype="site" objname="site_rear_left"/>
+    <framepos name="hr_pos" objtype="site" objname="site_rear_right"/>
+  </sensor>
+</mujoco>
diff --git a/judo/models/xml/spot_primitive/robot.xml b/judo/models/xml/spot_primitive/robot.xml
@@ -7,7 +7,7 @@
     <light directional="true" diffuse=".4 .4 .4" specular="0.1 0.1 0.1" pos="0 0 5" dir="0 0 -1" castshadow="false"/>
     <light directional="true" diffuse=".6 .6 .6" specular="0.2 0.2 0.2" pos="0 0 4" dir="0 0 -1"/>
 
-    <geom name="ground" type="plane" size="10 10 0.01" class="collision" priority="5" friction="0.7" material="blue_grid"/>
+    <geom name="ground" type="plane" size="10 10 0.01" class="collision" priority="5" friction="0.8" material="blue_grid"/>
 
     <body name="body" pos="0 0 0.7">
       <include file="body.xml" />

@@ -4,8 +4,14 @@
 <!-- Additional sensors should be added to dexterity/sene_generation/__init__.py -->
   <sensor>
     <framepos name="sensor_body" objtype="site" objname="site_body" reftype="site" refname="site_object" />
+    <framexaxis name="body_x_axis" objtype="site" objname="site_body"/>
     <frameyaxis name="object_y_axis" objtype="site" objname="site_object"/>
     <framezaxis name="object_z_axis" objtype="site" objname="site_object"/>
+    <framepos name="trace_fngr_site" objtype="site" objname="site_arm_link_fngr"/>
+    <framepos name="fl_pos" objtype="site" objname="site_front_left"/>
+    <framepos name="fr_pos" objtype="site" objname="site_front_right"/>
+    <framepos name="hl_pos" objtype="site" objname="site_rear_left"/>
+    <framepos name="hr_pos" objtype="site" objname="site_rear_right"/>
     <framepos name="sensor_arm_link_sh0" objtype="site" objname="site_arm_link_sh0" reftype="site" refname="site_object" />
     <framepos name="sensor_arm_link_sh1" objtype="site" objname="site_arm_link_sh1" reftype="site" refname="site_object" />
     <framepos name="sensor_arm_link_el0" objtype="site" objname="site_arm_link_el0" reftype="site" refname="site_object" />

@@ -73,7 +73,7 @@
     <light directional="true" diffuse=".4 .4 .4" specular="0.1 0.1 0.1" pos="0 0 5" dir="0 0 -1" castshadow="false"/>
     <light directional="true" diffuse=".6 .6 .6" specular="0.2 0.2 0.2" pos="0 0 4" dir="0 0 -1"/>
 
-    <geom name="ground" type="plane" size="10 10 0.01" class="collision" priority="5" friction="0.7" material="blue_grid"/>
+    <geom name="ground" type="plane" size="10 10 0.01" class="collision" priority="5" friction="0.8" material="blue_grid"/>
 
     <!-- Spot robot body at standing height -->
     <body name="body" pos="0 0 0.52">
@@ -94,9 +94,14 @@
   <!-- Sensors for task rewards -->
   <sensor>
     <framepos name="sensor_body" objtype="site" objname="site_body"/>
+    <framexaxis name="body_x_axis" objtype="site" objname="site_body"/>
+    <framexaxis name="object_x_axis" objtype="site" objname="site_object"/>
     <frameyaxis name="object_y_axis" objtype="site" objname="site_object"/>
     <framezaxis name="object_z_axis" objtype="site" objname="site_object"/>
     <framepos name="trace_fngr_site" objtype="site" objname="site_arm_link_fngr"/>
+    <framexaxis name="gripper_x_axis" objtype="site" objname="site_arm_link_wr1"/>
+    <frameyaxis name="gripper_y_axis" objtype="site" objname="site_arm_link_wr1"/>
+    <framexaxis name="finger_x_axis" objtype="site" objname="site_arm_link_fngr"/>
     <framepos name="fl_pos" objtype="site" objname="site_front_left"/>
     <framepos name="fr_pos" objtype="site" objname="site_front_right"/>
     <framepos name="hl_pos" objtype="site" objname="site_rear_left"/>

@@ -12,7 +12,7 @@
     set_default_fr3_pick_overrides,
     set_default_leap_cube_down_overrides,
     set_default_leap_cube_overrides,
-    set_default_spot_tire_upright_overrides,
+    set_default_spot_overrides,
 )
 from judo.optimizers.ps import PredictiveSampling, PredictiveSamplingConfig
 
@@ -22,7 +22,7 @@
 set_default_fr3_pick_overrides()
 set_default_leap_cube_overrides()
 set_default_leap_cube_down_overrides()
-set_default_spot_tire_upright_overrides()
+set_default_spot_overrides()
 
 
 _registered_optimizers: dict[str, tuple[Type[Optimizer], Type[OptimizerConfig]]] = {

@@ -199,9 +199,19 @@ def _set_spot_optimizer_overrides(task_name: str) -> None:
     set_config_overrides(task_name, MPPIConfig, _spot_base)
 
 
-def set_default_spot_tire_upright_overrides() -> None:
-    """Sets the default task-specific optimizer config overrides for the spot_tire_upright task."""
-    _set_spot_optimizer_overrides("spot_tire_upright")
+_SPOT_TASK_NAMES = [
+    "spot_base",
+    "spot_box_push",
+    "spot_navigate",
+    "spot_tire_roll",
+    "spot_tire_upright",
+]
+
+
+def set_default_spot_overrides() -> None:
+    """Sets the default task-specific optimizer config overrides for all Spot tasks."""
+    for task_name in _SPOT_TASK_NAMES:
+        _set_spot_optimizer_overrides(task_name)
 
 
 def set_default_fr3_pick_overrides() -> None:

@@ -3,14 +3,29 @@
 """MuJoCo Simulation with locomotion policy support."""
 
 from pathlib import Path
+from typing import NoReturn
 
 import numpy as np
 from mujoco import mj_forward
 from omegaconf import DictConfig
 
 from judo.simulation.mj_simulation import MJSimulation
 from judo.tasks.spot.spot_constants import DEFAULT_SPOT_ROLLOUT_CUTOFF_TIME, POLICY_OUTPUT_DIM
-from mujoco_extensions.policy_rollout import create_systems_vector, threaded_rollout  # type: ignore
+
+try:
+    from mujoco_extensions.policy_rollout import create_systems_vector, threaded_rollout  # type: ignore
+except ImportError as e:
+    _import_error = e
+
+    def _raise(*args: object, **kwargs: object) -> NoReturn:
+        raise ImportError(
+            "mujoco_extensions is not built. Spot locomotion tasks require the C++ extension.\n"
+            "Build it with:  pixi run build\n"
+            "See README.md for details."
+        ) from _import_error
+
+    create_systems_vector = _raise  # type: ignore[assignment]
+    threaded_rollout = _raise  # type: ignore[assignment]
 
 
 class PolicyMJSimulation(MJSimulation):

@@ -9,7 +9,18 @@
 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, SpotTireUpright, SpotTireUprightConfig
+from judo.tasks.spot import (
+    SpotBase,
+    SpotBaseConfig,
+    SpotBoxPush,
+    SpotBoxPushConfig,
+    SpotNavigate,
+    SpotNavigateConfig,
+    SpotTireRoll,
+    SpotTireRollConfig,
+    SpotTireUpright,
+    SpotTireUprightConfig,
+)
 
 _registered_tasks: Dict[str, Tuple[Type[Task], Type[TaskConfig]]] = {
     CylinderPush.name: (CylinderPush, CylinderPushConfig),
@@ -19,6 +30,9 @@
     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),
 }
 
@@ -52,6 +66,12 @@ def register_task(name: str, task_type: Type[Task], task_config_type: Type[TaskC
     "LeapCubeDownConfig",
     "SpotBase",
     "SpotBaseConfig",
+    "SpotBoxPush",
+    "SpotBoxPushConfig",
+    "SpotNavigate",
+    "SpotNavigateConfig",
+    "SpotTireRoll",
+    "SpotTireRollConfig",
     "SpotTireUpright",
     "SpotTireUprightConfig",
 ]
@@ -3,11 +3,20 @@
 """Spot locomotion and manipulation tasks."""
 
 from judo.tasks.spot.spot_base import SpotBase, SpotBaseConfig
+from judo.tasks.spot.spot_box_push import SpotBoxPush, SpotBoxPushConfig
+from judo.tasks.spot.spot_navigate import SpotNavigate, SpotNavigateConfig
+from judo.tasks.spot.spot_tire_roll import SpotTireRoll, SpotTireRollConfig
 from judo.tasks.spot.spot_tire_upright import SpotTireUpright, SpotTireUprightConfig
 
 __all__ = [
     "SpotBase",
     "SpotBaseConfig",
+    "SpotBoxPush",
+    "SpotBoxPushConfig",
+    "SpotNavigate",
+    "SpotNavigateConfig",
+    "SpotTireRoll",
+    "SpotTireRollConfig",
     "SpotTireUpright",
     "SpotTireUprightConfig",
 ]
@@ -88,15 +88,25 @@ class SpotBase(Task[ConfigT], Generic[ConfigT]):
     config_t: type[SpotBaseConfig] = SpotBaseConfig  # type: ignore[assignment]
 
     def _process_spec(self) -> None:
-        """Replace Spot mesh and texture paths with mujoco_menagerie assets."""
+        """Replace mesh and texture paths with correct resolved paths.
+
+        MjSpec resolves mesh file paths relative to the included file's directory,
+        ignoring the parent model's meshdir. We fix this by rewriting:
+        - Spot robot meshes → mujoco_menagerie assets (from robot_descriptions)
+        - Object meshes (tire, wheel_rim, etc.) → MODEL_PATH/meshes/objects/...
+        """
         from robot_descriptions import spot_mj_description  # noqa: PLC0415
 
         menagerie_dir = Path(spot_mj_description.PACKAGE_PATH)
         menagerie_assets = menagerie_dir / "assets"
+        meshes_root = MODEL_PATH / "meshes"
         for mesh in self.spec.meshes:
             if "spot/meshes/" in mesh.file:
                 basename = Path(mesh.file).name
                 mesh.file = str(menagerie_assets / basename)
+            elif "objects/" in mesh.file:
+                idx = mesh.file.index("objects/")
+                mesh.file = str(meshes_root / mesh.file[idx:])
         for texture in self.spec.textures:
             if "spot/textures/" in texture.file:
                 texture.file = str(menagerie_dir / "spot.png")
@@ -341,7 +351,9 @@ def task_to_sim_ctrl(self, controls: np.ndarray) -> np.ndarray:
             controls = controls[:, None, :]
             T = 1
 
-        out = np.zeros((controls.shape[0], controls.shape[1], 25), dtype=controls.dtype)
+        # Initialize from default_policy_command so uncontrolled dimensions
+        # keep their defaults (e.g. arm stays at ARM_STOWED_POS when use_arm=False)
+        out = np.broadcast_to(self.default_policy_command, (controls.shape[0], controls.shape[1], 25)).copy()
 
         # Index calculations after selection mask removal
         base_end = 3
@@ -352,9 +364,6 @@ def task_to_sim_ctrl(self, controls: np.ndarray) -> np.ndarray:
         # Base velocity
         out[..., 0:3] = controls[..., 0:base_end]
 
-        # Default torso height
-        out[..., 24] = STANDING_HEIGHT
-
         # Arm commands
         if self.use_arm:
             out[..., 3:10] = controls[..., base_end:arm_end]