adding features and tasks for the Spot-related tasks by johnzhang3 · Pull Request #95 · rai-opensource/judo

johnzhang3 · 2025-10-21T13:25:38Z

this PR introduces several significant changes

implements mujoco rollouts in cpp
separates task-space and mujoco simulation number of controls @jbruedigam-bdai
custom cpp rollout functions for the Spot robot that interleaves mujoco simulation and an RL policy inference from Relic. also implements a simulation cutoff time if a single thread is taking significant time
other changes include the CMAES optimizer and some benchmarking scripts from @alberthli (maybe that belongs in a separate PR....)

let me know if you guys have any feedback or want me to break this down into multiple PRs

- Add return type annotation for micro_benchmark_onnx_inference function - Fix unused loop variables by renaming to _trial - Remove unused result variable assignments in benchmarking code - Fix docstring formatting and add missing type annotations

- Implement PersistentThreadPool class with worker threads - Add ThreadPoolManager singleton for lifecycle management - Create persistent_cpp_rollout and persistent_onnx_interleave_rollout functions - Update Python bindings to expose new persistent functions - Add comprehensive benchmarking and testing scripts - Achieve 1.43x speedup and 87% reduction in ONNX overhead - Maintain perfect correctness with identity verification

- Add missing type annotations for all public functions - Fix unused loop variables by replacing with underscores - Add proper docstrings for public functions - Use specific type hints for function parameters and return types - Apply ruff-format code formatting - All critical linting errors resolved while maintaining functionality

bhung-bdai

Mostly minor comments, but I'm wondering if we can reduce the code with a bit of lift. If not, I'm not opposed to applying tech debt and letting this get fixed in the future.

bhung-bdai · 2025-10-21T15:49:26Z

+    <joint name="yellow_chair_joint" type="free"/>
+    <inertial pos="0 0.06 0.26" mass="16.58" diaginertia="0.25 0.25 0.25"/>
+
+    <!-- <geom pos="0 0.06 0.26" name="z_axis" class="visual" type="cylinder" size="0.01 1"/>


Remove comments?

bhung-bdai · 2025-10-21T15:50:21Z

@@ -0,0 +1,20 @@
+<mujoco model="yellow_chair">


Since this is public, I think it would be nice to have a reference to the actual chair itself. Could we find a link to it for sale?

bhung-bdai · 2025-10-21T15:52:41Z

@@ -19,6 +20,7 @@ class OptimizerConfig(OverridableConfig):
    num_nodes: int = 4
    use_noise_ramp: bool = False
    noise_ramp: float = 2.5
+    cutoff_time: float = 0.2  # Default for general use, Spot tasks may override


A description of what this means, either here or in the README, and how it affects the performance would be very helpful for those who aren't familiar with the code.

there isn't really a place that documents other optimizer settings. for now, I provided more detailed inline comments and added it in the change log. if we write more detailed docs for all the settings this setting should be included as well.

bhung-bdai · 2025-10-21T16:06:13Z

@@ -161,3 +182,83 @@ def get_joint_velocity_start_index(self, joint_name: str) -> int:
            joint_name: The name of the joint to get the starting index in the state array of.
        """
        return self.model.nq + self.model.jnt_dofadr[self.model.joint(joint_name).id]
+
+    def success(self, model: MjModel, data: MjData, config: ConfigT, metadata: dict[str, Any] | None = None) -> bool:


I'd consider doing something like this instead:

raise NotImplementedError("The failure criteria needs to be implemented by the child task.")

bhung-bdai · 2025-10-21T16:06:44Z

+        """
+        return False
+
+    def failure(self, model: MjModel, data: MjData, config: ConfigT, metadata: dict[str, Any] | None = None) -> bool:


Same down here:

raise NotImplementedError("The episode failure criteria needs to be implemented by the child task.")

bhung-bdai · 2025-10-21T17:37:06Z

+        ).sum(axis=-1)
+
+        # Compute l2 distance from torso pos. to object pos.
+        torso_proximity_reward = config.w_torso_proximity * np.linalg.norm(body_pos - object_pos, axis=-1).mean(-1)


Just to check, should this be a positive or a negative reward?

ah. this should negative. fixed now.

bhung-bdai · 2025-10-21T17:40:30Z

-        assert controls.shape[-1] == nu
-        assert controls.shape[0] == full_states.shape[0]
+        assert processed_controls.ndim == 3
+        # assert processed_controls.shape[-1] == nu


bhung-bdai · 2025-10-21T17:51:15Z

+    auto heap_buf = new std::vector<double>(std::move(buf));
+    // Create a capsule that will delete the vector when the array is gone:
+    py::capsule free_when_done(heap_buf, [](void *p) {
+        delete reinterpret_cast<std::vector<double>*>(p);


Nit: I kind of wish we just used smart pointers here instead of constant castings

bhung-bdai · 2025-10-21T17:56:56Z

+// SpotThreadPool Implementation
+// =============================================================================
+
+SpotThreadPool::SpotThreadPool(int num_threads)


Do we need a separate spot thread pool vs a normal thread pool?

bhung-bdai · 2025-10-21T18:58:40Z

+    }
+}
+
+void SpotThreadPool::worker_thread() {


I'm really unsure about why we can't use the PersistentThreadPool object for this. The only main difference seems to be that you must spin down the PersistentThreadPool after you run the rollout.

bhung-bdai · 2025-10-21T19:25:23Z

We also likely need to think about how to manage the assets in a more lightweight way, but that's not totally related to this work.

jbruedigam-bdai

Agree with Brandon's comments and have a few additional ones

jbruedigam-bdai · 2025-10-22T08:19:49Z

Could we rename this to spot_relic_policy
Since we may have different policies later on

yep. this is now fixed

jbruedigam-bdai · 2025-10-22T08:24:23Z

High-level questions: is the rollout specific to Spot or does it work for any mj model and policy?

for now they are specific to the relic policy for the spot robot. we do have plans for making this more general in the future.

Resolved conflicts: - CHANGELOG.md: Merged unreleased changes with v0.0.5 release - judo/app/simulation.py: Accepted deletion (refactored into simulation/ module) - judo/controller/controller.py: Kept C++ backend support and optimizer_config - judo/tasks/__init__.py: Updated to use Task.name pattern for all tasks - judo/tasks/base.py: Kept task_to_sim_ctrl method and default_backend

- Added name class attribute to SpotBox, SpotYellowChair, and SpotYellowChairRamp - Required by main's task registration refactoring that uses Task.name - Fixes AttributeError when importing judo.app.dora modules

- Updated make_controller() to pass optimizer_config parameter - Updated ControllerNode.update_task() to pass optimizer_config when recreating controller - Fixes Controller instantiation after merge with main that kept spot_cpp's modified __init__ signature

- Updated version to 0.0.5 - Kept C++ build configuration and pixi tasks - Resolved conflicts in pyproject.toml and pixi.lock

Remove build-wheels.yml workflow and pyproject-build.toml as we'll revisit the wheel building setup later. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com>

Add type: ignore comments for MuJoCo data array attributes since they return ndarray[float64] which pyright doesn't accept as ndarray[_AnyShape]. The types are correct at runtime. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com>

Change from mypy-style type: ignore to pyright: ignore[reportArgumentType] to suppress the covariance type errors in MujocoState instantiation. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com>

Replace pyright: ignore comment with explicit typing.cast() calls to convert MuJoCo data arrays to generic np.ndarray type for MujocoState constructor. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com>

johnzhang-rai and others added 30 commits August 8, 2025 16:36

add spot files

ba3a98f

spot standing runs

6c0988f

add boucing ball model to fr3

ba595b5

reorg

930b60e

add spor door and box task

f22d65d

add backend options

eb02a82

onnx inference working with judo app

00683d1

s

da6f6d5

fixes for arm mac

e129a58

Merge commit 'e129a58b9a8bf4ba9dcd8083e3406a111843cc1a' into spot

6b7fa7c

cpp on mac

c812a1c

add spot model tracking

f370baf

more cpp changes for mac

2c88222

reorganize

84a945e

debug the policy rollout

6a1701d

add cpp build in readme

82662b0

basic version working with spot on policy onnx

0ecf5aa

working dummy policy setup

67758f8

rm unused stuff

b940d0b

clean up

efaa14b

add xinghao policy

9b6bf81

update gitignore

67c69cb

wip policy wrapper

11a2130

add onnx

e684e1f

policy wrapper

eb7cacf

wip integrated wrapped policy with judo cpp

cee9471

wip onnx inference runs in judo app

92e2d07

johnzhang3 requested a review from bhung-bdai as a code owner October 21, 2025 13:25

johnzhang-rai added 6 commits October 21, 2025 10:52

Fix pre-commit linting issues

5510b29

Clean up task registrations and default config

ec9167f

Fix pyright type errors

6e655a8

Fix C++ import errors for CI environments

e138187

Use dummy implementations for C++ imports when module not available

2ea51d2

Add type stubs for C++ imports to fix pyright in CI

73033e8

bhung-bdai reviewed Oct 21, 2025

View reviewed changes

jbruedigam-bdai reviewed Oct 22, 2025

View reviewed changes

johnzhang3 and others added 20 commits October 23, 2025 10:41

update readme

1a50263

address comments from brandon and jan

3e9ea03

update change log

c7b503e

add yellow chair link

1afacd0

Merge branch 'bdaiinstitute:main' into main

641fb59

trying prebuilt binary

d866cb4

update change log

a924245

Fix: Add name attribute to Spot task classes

02deb67

- Added name class attribute to SpotBox, SpotYellowChair, and SpotYellowChairRamp - Required by main's task registration refactoring that uses Task.name - Fixes AttributeError when importing judo.app.dora modules

bug fixes with new code

debd75a

Merge upstream/main into spot_cpp

27b939d

- Updated version to 0.0.5 - Kept C++ build configuration and pixi tasks - Resolved conflicts in pyproject.toml and pixi.lock

Regenerate pixi.lock after merge with upstream/main

e925330

Add mujoco and onnxruntime to CI build dependencies

b06ad85

Fix CI: Use CIBW_BEFORE_BUILD instead of invalid CIBW_BUILD_REQUIRES

7275d3e

Fix Eigen download URL in CI (use GitLab instead of old GitHub mirror)

cbbb870

Remove wheel building infrastructure

c4a1d26

Remove build-wheels.yml workflow and pyproject-build.toml as we'll revisit the wheel building setup later. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com>

Conversation

johnzhang3 commented Oct 21, 2025

Uh oh!

bhung-bdai left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

bhung-bdai commented Oct 21, 2025

Uh oh!

jbruedigam-bdai left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants