Skip to content

Commit 1b959ba

Browse files
dta-bdaidta-bdai
authored
Add all spot tasks (#111)
Add other spot tasks
1 parent f2a472d commit 1b959ba

19 files changed

Lines changed: 685 additions & 49 deletions

‎judo/controller/__init__.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@
88
set_default_fr3_pick_overrides,
99
set_default_leap_cube_down_overrides,
1010
set_default_leap_cube_overrides,
11-
set_default_spot_tire_upright_overrides,
11+
set_default_spot_overrides,
1212
)
1313

1414
set_default_caltech_leap_cube_overrides()
@@ -17,7 +17,7 @@
1717
set_default_fr3_pick_overrides()
1818
set_default_leap_cube_overrides()
1919
set_default_leap_cube_down_overrides()
20-
set_default_spot_tire_upright_overrides()
20+
set_default_spot_overrides()
2121

2222
__all__ = [
2323
"Controller",

‎judo/controller/overrides.py‎

Lines changed: 19 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -67,15 +67,25 @@ def set_default_caltech_leap_cube_overrides() -> None:
6767
)
6868

6969

70-
def set_default_spot_tire_upright_overrides() -> None:
71-
"""Sets the default task-specific controller config overrides for the spot_tire_upright task."""
72-
set_config_overrides(
73-
"spot_tire_upright",
74-
ControllerConfig,
75-
{
76-
"horizon": 2.0,
77-
},
78-
)
70+
_SPOT_TASK_NAMES = [
71+
"spot_base",
72+
"spot_box_push",
73+
"spot_navigate",
74+
"spot_tire_roll",
75+
"spot_tire_upright",
76+
]
77+
78+
79+
def set_default_spot_overrides() -> None:
80+
"""Sets the default task-specific controller config overrides for all Spot tasks."""
81+
for task_name in _SPOT_TASK_NAMES:
82+
set_config_overrides(
83+
task_name,
84+
ControllerConfig,
85+
{
86+
"horizon": 2.0,
87+
},
88+
)
7989

8090

8191
def set_default_fr3_pick_overrides() -> None:

‎judo/gui.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@
1919

2020
DEFAULT_SLIDER_STEP_FLOAT = 0.01
2121
DEFAULT_SLIDER_STEP_INT = 1
22-
GOAL_RADIUS = 0.05
22+
GOAL_RADIUS = 0.15
2323

2424

2525
def slider(

‎judo/models/xml/spot_box/robot.xml‎

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,51 @@
1+
<mujoco model="spot_box">
2+
<!-- Copyright (c) 2025 Robotics and AI Institute LLC. All rights reserved. -->
3+
<!-- Scene: Spot robot with box object for push tasks -->
4+
5+
<compiler angle="radian" meshdir="../../meshes/" assetdir="../../meshes/" autolimits="true"/>
6+
<option timestep="0.01" solver="Newton" integrator="implicitfast" density="1"/>
7+
8+
<include file="../spot_primitive/default.xml" />
9+
<include file="../spot_primitive/assets.xml" />
10+
11+
<worldbody>
12+
<light directional="true" diffuse=".4 .4 .4" specular="0.1 0.1 0.1" pos="0 0 5" dir="0 0 -1" castshadow="false"/>
13+
<light directional="true" diffuse=".6 .6 .6" specular="0.2 0.2 0.2" pos="0 0 4" dir="0 0 -1"/>
14+
15+
<geom name="ground" type="plane" size="10 10 0.01" class="collision" priority="5" friction="0.7" material="blue_grid"/>
16+
17+
<body name="body" pos="0 0 0.52">
18+
<include file="../spot_primitive/body.xml" />
19+
<include file="../spot_primitive/legs.xml" />
20+
<include file="../spot_primitive/arm.xml" />
21+
</body>
22+
23+
<!-- Box object -->
24+
<body name="box_body" pos="2 0 0.254">
25+
<joint name="box_joint" type="free"/>
26+
<inertial pos="0 0 0" mass="1.5" diaginertia="0.1445 0.1445 0.1445"/>
27+
<geom name="box_collision" type="box" size="0.254 0.254 0.254" class="collision" priority="4"/>
28+
<geom name="box_visual" type="box" size="0.254 0.254 0.254" class="visual" rgba="0.6 0.3 0.1 1"/>
29+
<site name="site_object" pos="0 0 0" size="0.01"/>
30+
</body>
31+
</worldbody>
32+
33+
<include file="../spot_primitive/actuator.xml" />
34+
<include file="../spot_primitive/contact.xml" />
35+
36+
<sensor>
37+
<framepos name="sensor_body" objtype="site" objname="site_body"/>
38+
<framexaxis name="body_x_axis" objtype="site" objname="site_body"/>
39+
<framexaxis name="object_x_axis" objtype="site" objname="site_object"/>
40+
<frameyaxis name="object_y_axis" objtype="site" objname="site_object"/>
41+
<framezaxis name="object_z_axis" objtype="site" objname="site_object"/>
42+
<framepos name="trace_fngr_site" objtype="site" objname="site_arm_link_fngr"/>
43+
<framexaxis name="gripper_x_axis" objtype="site" objname="site_arm_link_wr1"/>
44+
<frameyaxis name="gripper_y_axis" objtype="site" objname="site_arm_link_wr1"/>
45+
<framexaxis name="finger_x_axis" objtype="site" objname="site_arm_link_fngr"/>
46+
<framepos name="fl_pos" objtype="site" objname="site_front_left"/>
47+
<framepos name="fr_pos" objtype="site" objname="site_front_right"/>
48+
<framepos name="hl_pos" objtype="site" objname="site_rear_left"/>
49+
<framepos name="hr_pos" objtype="site" objname="site_rear_right"/>
50+
</sensor>
51+
</mujoco>

‎judo/models/xml/spot_primitive/sensor.xml‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,8 +4,14 @@
44
<!-- Additional sensors should be added to dexterity/sene_generation/__init__.py -->
55
<sensor>
66
<framepos name="sensor_body" objtype="site" objname="site_body" reftype="site" refname="site_object" />
7+
<framexaxis name="body_x_axis" objtype="site" objname="site_body"/>
78
<frameyaxis name="object_y_axis" objtype="site" objname="site_object"/>
89
<framezaxis name="object_z_axis" objtype="site" objname="site_object"/>
10+
<framepos name="trace_fngr_site" objtype="site" objname="site_arm_link_fngr"/>
11+
<framepos name="fl_pos" objtype="site" objname="site_front_left"/>
12+
<framepos name="fr_pos" objtype="site" objname="site_front_right"/>
13+
<framepos name="hl_pos" objtype="site" objname="site_rear_left"/>
14+
<framepos name="hr_pos" objtype="site" objname="site_rear_right"/>
915
<framepos name="sensor_arm_link_sh0" objtype="site" objname="site_arm_link_sh0" reftype="site" refname="site_object" />
1016
<framepos name="sensor_arm_link_sh1" objtype="site" objname="site_arm_link_sh1" reftype="site" refname="site_object" />
1117
<framepos name="sensor_arm_link_el0" objtype="site" objname="site_arm_link_el0" reftype="site" refname="site_object" />

‎judo/models/xml/spot_tire/robot.xml‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -94,9 +94,14 @@
9494
<!-- Sensors for task rewards -->
9595
<sensor>
9696
<framepos name="sensor_body" objtype="site" objname="site_body"/>
97+
<framexaxis name="body_x_axis" objtype="site" objname="site_body"/>
98+
<framexaxis name="object_x_axis" objtype="site" objname="site_object"/>
9799
<frameyaxis name="object_y_axis" objtype="site" objname="site_object"/>
98100
<framezaxis name="object_z_axis" objtype="site" objname="site_object"/>
99101
<framepos name="trace_fngr_site" objtype="site" objname="site_arm_link_fngr"/>
102+
<framexaxis name="gripper_x_axis" objtype="site" objname="site_arm_link_wr1"/>
103+
<frameyaxis name="gripper_y_axis" objtype="site" objname="site_arm_link_wr1"/>
104+
<framexaxis name="finger_x_axis" objtype="site" objname="site_arm_link_fngr"/>
100105
<framepos name="fl_pos" objtype="site" objname="site_front_left"/>
101106
<framepos name="fr_pos" objtype="site" objname="site_front_right"/>
102107
<framepos name="hl_pos" objtype="site" objname="site_rear_left"/>

‎judo/optimizers/__init__.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@
1212
set_default_fr3_pick_overrides,
1313
set_default_leap_cube_down_overrides,
1414
set_default_leap_cube_overrides,
15-
set_default_spot_tire_upright_overrides,
15+
set_default_spot_overrides,
1616
)
1717
from judo.optimizers.ps import PredictiveSampling, PredictiveSamplingConfig
1818

@@ -22,7 +22,7 @@
2222
set_default_fr3_pick_overrides()
2323
set_default_leap_cube_overrides()
2424
set_default_leap_cube_down_overrides()
25-
set_default_spot_tire_upright_overrides()
25+
set_default_spot_overrides()
2626

2727

2828
_registered_optimizers: dict[str, tuple[Type[Optimizer], Type[OptimizerConfig]]] = {

‎judo/optimizers/overrides.py‎

Lines changed: 13 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -199,9 +199,19 @@ def _set_spot_optimizer_overrides(task_name: str) -> None:
199199
set_config_overrides(task_name, MPPIConfig, _spot_base)
200200

201201

202-
def set_default_spot_tire_upright_overrides() -> None:
203-
"""Sets the default task-specific optimizer config overrides for the spot_tire_upright task."""
204-
_set_spot_optimizer_overrides("spot_tire_upright")
202+
_SPOT_TASK_NAMES = [
203+
"spot_base",
204+
"spot_box_push",
205+
"spot_navigate",
206+
"spot_tire_roll",
207+
"spot_tire_upright",
208+
]
209+
210+
211+
def set_default_spot_overrides() -> None:
212+
"""Sets the default task-specific optimizer config overrides for all Spot tasks."""
213+
for task_name in _SPOT_TASK_NAMES:
214+
_set_spot_optimizer_overrides(task_name)
205215

206216

207217
def set_default_fr3_pick_overrides() -> None:

‎judo/simulation/policy_mj_simulation.py‎

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,15 @@
1010

1111
from judo.simulation.mj_simulation import MJSimulation
1212
from judo.tasks.spot.spot_constants import DEFAULT_SPOT_ROLLOUT_CUTOFF_TIME, POLICY_OUTPUT_DIM
13-
from mujoco_extensions.policy_rollout import create_systems_vector, threaded_rollout # type: ignore
13+
14+
try:
15+
from mujoco_extensions.policy_rollout import create_systems_vector, threaded_rollout # type: ignore
16+
except ImportError as e:
17+
raise ImportError(
18+
"mujoco_extensions is not built. Spot locomotion tasks require the C++ extension.\n"
19+
"Build it with: pixi run build\n"
20+
"See README.md for details."
21+
) from e
1422

1523

1624
class PolicyMJSimulation(MJSimulation):

‎judo/tasks/__init__.py‎

Lines changed: 21 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,18 @@
99
from judo.tasks.fr3_pick import FR3Pick, FR3PickConfig
1010
from judo.tasks.leap_cube import LeapCube, LeapCubeConfig
1111
from judo.tasks.leap_cube_down import LeapCubeDown, LeapCubeDownConfig
12-
from judo.tasks.spot import SpotBase, SpotBaseConfig, SpotTireUpright, SpotTireUprightConfig
12+
from judo.tasks.spot import (
13+
SpotBase,
14+
SpotBaseConfig,
15+
SpotBoxPush,
16+
SpotBoxPushConfig,
17+
SpotNavigate,
18+
SpotNavigateConfig,
19+
SpotTireRoll,
20+
SpotTireRollConfig,
21+
SpotTireUpright,
22+
SpotTireUprightConfig,
23+
)
1324

1425
_registered_tasks: Dict[str, Tuple[Type[Task], Type[TaskConfig]]] = {
1526
CylinderPush.name: (CylinderPush, CylinderPushConfig),
@@ -19,6 +30,9 @@
1930
LeapCubeDown.name: (LeapCubeDown, LeapCubeDownConfig),
2031
CaltechLeapCube.name: (CaltechLeapCube, CaltechLeapCubeConfig),
2132
SpotBase.name: (SpotBase, SpotBaseConfig),
33+
SpotBoxPush.name: (SpotBoxPush, SpotBoxPushConfig),
34+
SpotNavigate.name: (SpotNavigate, SpotNavigateConfig),
35+
SpotTireRoll.name: (SpotTireRoll, SpotTireRollConfig),
2236
SpotTireUpright.name: (SpotTireUpright, SpotTireUprightConfig),
2337
}
2438

@@ -52,6 +66,12 @@ def register_task(name: str, task_type: Type[Task], task_config_type: Type[TaskC
5266
"LeapCubeDownConfig",
5367
"SpotBase",
5468
"SpotBaseConfig",
69+
"SpotBoxPush",
70+
"SpotBoxPushConfig",
71+
"SpotNavigate",
72+
"SpotNavigateConfig",
73+
"SpotTireRoll",
74+
"SpotTireRollConfig",
5575
"SpotTireUpright",
5676
"SpotTireUprightConfig",
5777
]

0 commit comments

Comments
 (0)