feat(vla_sim): collect demonstrations with the scripted oracle - #830
Open
danwahl wants to merge 2 commits into
Open
feat(vla_sim): collect demonstrations with the scripted oracle#830danwahl wants to merge 2 commits into
danwahl wants to merge 2 commits into
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughSummary by CodeRabbit
WalkthroughThe PR adds a cube-stack VLA simulation workflow with MuJoCo layouts, BehaviorTree recording objectives, MoveIt behavior plugins, joint-spline planning, gripper and recording controls, and a ROS 2 joint-state bridge. ChangesCube-stack VLA workflow
Possibly related PRs
Suggested reviewers: Caution Pre-merge checks failedPlease resolve all errors before merging. Addressing warnings is optional.
❌ Failed checks (1 error)
✅ Passed checks (3 passed)
Comment |
|
Consider whether the change should land upstream in Overlapping files
|
Author
|
[written by AI] Docs for this branch: PickNikRobotics/moveit_pro#21138 (stacked on the training-data guide, #20932). |
JWhitleyWork
force-pushed
the
20588-vla-sim
branch
from
August 6, 2026 16:48
417fdb3 to
011b7be
Compare
danwahl
force-pushed
the
vla-sim-oracle-collection
branch
from
August 7, 2026 01:45
41e62ae to
d2e550c
Compare
|
danwahl
commented
Aug 7, 2026
360 train + 150 eval randomized cube layouts, reachable by name through /mujoco_system/reset_keyframe. The eval set is held out from the shipped checkpoint's demonstrations, so it is the only fair set to score on.
danwahl
force-pushed
the
vla-sim-oracle-collection
branch
from
August 7, 2026 02:47
d2e550c to
42fdd29
Compare
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
|
danwahl
commented
Aug 7, 2026
danwahl
force-pushed
the
vla-sim-oracle-collection
branch
from
August 7, 2026 03:25
42fdd29 to
f33dd0d
Compare
Author
|
[written by AI] For anyone following the |
|
danwahl
force-pushed
the
vla-sim-oracle-collection
branch
from
August 7, 2026 04:00
f33dd0d to
331742d
Compare
|
danwahl
commented
Aug 7, 2026
danwahl
force-pushed
the
vla-sim-oracle-collection
branch
from
August 7, 2026 04:39
331742d to
a6e5f27
Compare
|
danwahl
marked this pull request as ready for review
August 7, 2026 05:02
There was a problem hiding this comment.
Actionable comments posted: 10
🧹 Nitpick comments (2)
src/vla_sim/description/mujoco/keyframes.xml (1)
22-33: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueConsider moving each prompt comment above the key it describes.
Each prompt comment follows its
<key>element. The comment on line 27 therefore belongs toeval_0, not toeval_1that starts on line 28. The header on lines 13-15 documents this, and the final comment on line 3081 aftertrain_359confirms it. A reader who assumes the usual leading-comment convention pairs every layout with the wrong prompt, and a wrong prompt is not detectable from the layout data.This file is generated, so the fix belongs in the generator. The current form is correct, so treat this as a readability change only.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/vla_sim/description/mujoco/keyframes.xml` around lines 22 - 33, Move each prompt comment in the keyframe generator so it appears immediately before the corresponding <key> element rather than after it. Preserve the generated keyframe content and ordering, including the existing association between each prompt and key such as eval_0 and eval_1; this is a readability-only change.src/vla_sim_behaviors/test/test_send_gripper_command.cpp (1)
45-48: 🩺 Stability & Availability | 🔵 Trivial | 💤 Low value
received_.set_valuethrows if a second goal arrives.The goal callback calls
set_valueon every accepted goal. A second goal makesstd::promise::set_valuethrowstd::future_errorinside an rclcpp callback. Only one goal is sent today, so this is latent. If you add a test that sends two goals, guard the promise with astd::once_flagor aboolundermutex_.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/vla_sim_behaviors/test/test_send_gripper_command.cpp` around lines 45 - 48, Update the goal callback’s received_ fulfillment so it records only the first accepted goal; guard set_value with the existing mutex_ and a bool or std::once_flag, preventing subsequent goals from calling std::promise::set_value again.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/vla_sim_behaviors/src/compute_top_down_keyposes.cpp`:
- Around line 225-240: Update the cost_of lambda to validate IK for every
keypose generated from heights, while retaining the existing jointDistanceCost
based on the first approach pose. Return std::nullopt if IK fails for any later
height, so yaw selection excludes candidates that PlanJointSplineThroughPoses
cannot execute. Add a regression test covering a yaw that succeeds at
heights.front() but fails at a subsequent height.
In `@src/vla_sim_behaviors/src/plan_joint_spline_through_poses.cpp`:
- Around line 360-367: Validate joint_velocity_scale before constructing
velocity_cap, requiring it to be finite and within the inclusive range (0.0,
1.0]. Reject invalid values before the bounded-joint velocity-cap calculation so
zero, NaN, and values above the model limit cannot reach the timing helpers.
- Around line 224-237: The duration calculation in splineDuration must retain
the minimum-duration floor without upper-clamping the required duration; update
the planner at src/vla_sim_behaviors/src/plan_joint_spline_through_poses.cpp
lines 224-237 and return FAILURE when that required duration exceeds
kMaximumDuration at lines 393-394. Update
src/vla_sim_behaviors/test/test_plan_joint_spline_through_poses.cpp lines
176-180 to assert planner failure instead of expecting a 60-second trajectory.
- Around line 393-400: Validate the duration-to-sample calculation in the
trajectory generation flow before converting it to std::size_t or reserving
points. Enforce a defined maximum point count (including the final sample) and
return FAILURE when the configured sampling_rate or resulting count exceeds that
limit; otherwise preserve the existing loop and trajectory construction
behavior.
In `@src/vla_sim_behaviors/src/send_gripper_command.cpp`:
- Around line 100-104: Update SendGripperCommand to use an asynchronous or
stateful execution model instead of SyncActionNode, retaining the
goal-acceptance future from client_->async_send_goal. Return RUNNING while
acceptance is pending, FAILURE when the resolved goal handle is null, and
SUCCESS only after a valid handle is received; do not wait for the action result
or gripper motion.
In `@src/vla_sim_behaviors/src/wait_for_episode_start.cpp`:
- Around line 102-136: Move deadline creation before client_->initialize in the
episode-start flow, then pass the remaining time until that deadline to
waitForServiceServer and each syncSendRequest call instead of fixed five-second
limits. Before sleeping, cap kPollPeriod to the remaining budget, and preserve
the existing timeout error behavior when the deadline is reached.
In `@src/vla_sim_behaviors/test/test_wait_for_episode_start.cpp`:
- Around line 84-91: Make executor shutdown scope-bound in both test fixtures:
in src/vla_sim_behaviors/test/test_wait_for_episode_start.cpp lines 84-91,
extract idempotent stopSpinning() logic from ~WaitForEpisodeStartTest() and call
it in every test before the local trainer is destroyed; in
src/vla_sim_behaviors/test/test_send_gripper_command.cpp lines 143-158, declare
a scope guard after server that invokes stopSpinning(), ensuring cleanup also
runs when ASSERT_EQ exits the test early.
In `@src/vla_sim/launch/simulated_extras.launch.py`:
- Around line 38-43: Update the JointCommandBridge Node configuration in
simulated_extras.launch.py to pass the MOVEIT_PRO_TRAIN_JOINT_STATES_TOPIC
environment variable as observation_state_topic, using EnvironmentVariable with
/observed_joint_states as the default value.
In `@src/vla_sim/objectives/record_cube_stack_episode.xml`:
- Around line 18-25: Ensure the episode sequence always invokes the idempotent
StopRecording action when WaitForEpisodeStart times out or the behavior tree is
halted, including when episode start fails and the normal Sequence path is
skipped. Update the RecordEpisode/WaitForEpisodeStart flow to attach cleanup
that runs on both timeout and halt while preserving the existing successful
episode path.
In `@src/vla_sim/script/joint_command_bridge.py`:
- Around line 65-66: Update is_reference_fresh to require the elapsed time (now
- stamp) to be non-negative as well as below timeout, so future timestamps are
rejected after simulated-clock resets; add a test covering a stamp later than
now.
---
Nitpick comments:
In `@src/vla_sim_behaviors/test/test_send_gripper_command.cpp`:
- Around line 45-48: Update the goal callback’s received_ fulfillment so it
records only the first accepted goal; guard set_value with the existing mutex_
and a bool or std::once_flag, preventing subsequent goals from calling
std::promise::set_value again.
In `@src/vla_sim/description/mujoco/keyframes.xml`:
- Around line 22-33: Move each prompt comment in the keyframe generator so it
appears immediately before the corresponding <key> element rather than after it.
Preserve the generated keyframe content and ordering, including the existing
association between each prompt and key such as eval_0 and eval_1; this is a
readability-only change.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: dddadb72-918e-4873-a90a-62f3be007c7e
📒 Files selected for processing (41)
docker-compose.yamlsrc/vla_sim/CMakeLists.txtsrc/vla_sim/README.mdsrc/vla_sim/config/config.yamlsrc/vla_sim/description/mujoco/cube_stack_scene.xmlsrc/vla_sim/description/mujoco/keyframes.xmlsrc/vla_sim/launch/simulated_extras.launch.pysrc/vla_sim/objectives/collect_cube_stack_demo.xmlsrc/vla_sim/objectives/command_cube_stack_gripper.xmlsrc/vla_sim/objectives/execute_cube_stack_oracle.xmlsrc/vla_sim/objectives/move_along_cube_stack_keyposes.xmlsrc/vla_sim/objectives/prepare_cube_stack_scene.xmlsrc/vla_sim/objectives/record_cube_stack_blue_on_green.xmlsrc/vla_sim/objectives/record_cube_stack_blue_on_red.xmlsrc/vla_sim/objectives/record_cube_stack_episode.xmlsrc/vla_sim/objectives/record_cube_stack_green_on_blue.xmlsrc/vla_sim/objectives/record_cube_stack_green_on_red.xmlsrc/vla_sim/objectives/record_cube_stack_red_on_blue.xmlsrc/vla_sim/objectives/record_cube_stack_red_on_green.xmlsrc/vla_sim/objectives/run_cube_stack_oracle.xmlsrc/vla_sim/package.xmlsrc/vla_sim/script/joint_command_bridge.pysrc/vla_sim/test/test_joint_command_bridge.pysrc/vla_sim_behaviors/CMakeLists.txtsrc/vla_sim_behaviors/include/vla_sim_behaviors/compute_top_down_keyposes.hppsrc/vla_sim_behaviors/include/vla_sim_behaviors/plan_joint_spline_through_poses.hppsrc/vla_sim_behaviors/include/vla_sim_behaviors/send_gripper_command.hppsrc/vla_sim_behaviors/include/vla_sim_behaviors/wait_for_episode_start.hppsrc/vla_sim_behaviors/package.xmlsrc/vla_sim_behaviors/src/compute_top_down_keyposes.cppsrc/vla_sim_behaviors/src/plan_joint_spline_through_poses.cppsrc/vla_sim_behaviors/src/register_behaviors.cppsrc/vla_sim_behaviors/src/send_gripper_command.cppsrc/vla_sim_behaviors/src/wait_for_episode_start.cppsrc/vla_sim_behaviors/test/CMakeLists.txtsrc/vla_sim_behaviors/test/test_behavior_plugins.cppsrc/vla_sim_behaviors/test/test_compute_top_down_keyposes.cppsrc/vla_sim_behaviors/test/test_plan_joint_spline_through_poses.cppsrc/vla_sim_behaviors/test/test_send_gripper_command.cppsrc/vla_sim_behaviors/test/test_wait_for_episode_start.cppsrc/vla_sim_behaviors/vla_sim_behaviors_plugin_description.xml
danwahl
force-pushed
the
vla-sim-oracle-collection
branch
from
August 10, 2026 20:47
a6e5f27 to
27575b7
Compare
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/vla_sim_behaviors/test/test_plan_joint_spline_through_poses.cpp`:
- Around line 152-159: Extend the SplineKnotParameters regression test with
durations {3.0, 0.0, 1.0} and verify the resulting parameters remain strictly
increasing and JointSpline does not throw. Update splineKnotParameters so that
when its epsilon-based increment fails to advance the cumulative value, it uses
a representably greater value such as std::nextafter.
In `@src/vla_sim/script/joint_command_bridge.py`:
- Around line 218-280: Update _publish_action to capture the current time once
and use it for freshness checks. Pass self._measured to assemble_joint_command
only when is_reference_fresh(now, self._measured_stamp, self._reference_timeout)
is true; otherwise pass no measured positions, allowing the action to be
suppressed when the controller reference is also stale. Reuse the captured now
for the existing controller-reference freshness check.
- Around line 135-144: The topic collision check in the bridge initialization
must compare resolved names rather than raw parameter strings. Use
self.resolve_topic_name() for both _observation_state_topic and
_joint_states_topic before the equality check, while preserving the existing
ValueError behavior and message context. Add coverage for a relative-topic
override resolving to the same absolute topic.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 6c10ff36-16e7-4718-bcfa-6c1f9138a430
📒 Files selected for processing (12)
src/vla_sim/launch/simulated_extras.launch.pysrc/vla_sim/objectives/record_cube_stack_episode.xmlsrc/vla_sim/script/joint_command_bridge.pysrc/vla_sim/test/test_joint_command_bridge.pysrc/vla_sim_behaviors/include/vla_sim_behaviors/compute_top_down_keyposes.hppsrc/vla_sim_behaviors/include/vla_sim_behaviors/plan_joint_spline_through_poses.hppsrc/vla_sim_behaviors/src/compute_top_down_keyposes.cppsrc/vla_sim_behaviors/src/plan_joint_spline_through_poses.cppsrc/vla_sim_behaviors/test/test_compute_top_down_keyposes.cppsrc/vla_sim_behaviors/test/test_plan_joint_spline_through_poses.cppsrc/vla_sim_behaviors/test/test_send_gripper_command.cppsrc/vla_sim_behaviors/test/test_wait_for_episode_start.cpp
🚧 Files skipped from review as they are similar to previous changes (7)
- src/vla_sim/objectives/record_cube_stack_episode.xml
- src/vla_sim_behaviors/include/vla_sim_behaviors/plan_joint_spline_through_poses.hpp
- src/vla_sim_behaviors/test/test_send_gripper_command.cpp
- src/vla_sim_behaviors/src/plan_joint_spline_through_poses.cpp
- src/vla_sim/test/test_joint_command_bridge.py
- src/vla_sim_behaviors/src/compute_top_down_keyposes.cpp
- src/vla_sim_behaviors/test/test_compute_top_down_keyposes.cpp
|
danwahl
force-pushed
the
vla-sim-oracle-collection
branch
from
August 10, 2026 21:53
27575b7 to
0b93b00
Compare
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/vla_sim_behaviors/src/plan_joint_spline_through_poses.cpp`:
- Around line 210-229: Update splineKnotParameters in
src/vla_sim_behaviors/src/plan_joint_spline_through_poses.cpp at lines 210-229
to accumulate raw nonnegative totals, normalize all parameters by total, then
enforce strict increase with std::nextafter in the normalized domain; handle
all-zero durations without division by zero. Add tests in
src/vla_sim_behaviors/test/test_plan_joint_spline_through_poses.cpp at lines
161-169 covering durations {3.9, 0.0, 3.9} and an all-zero duration case,
verifying the returned knot parameters remain valid and strictly increasing.
- Around line 362-376: Update PlanJointSplineThroughPoses before the setFromIK
loop to enforce the MoveIt robot model frame: transform each waypoint pose from
its header.frame_id into the model frame, or reject it when the frame cannot be
transformed. Ensure setFromIK receives only model-frame poses while preserving
the existing IK failure handling.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 77a403e8-04ae-4139-b428-6e05ddcc93eb
📒 Files selected for processing (3)
src/vla_sim/script/joint_command_bridge.pysrc/vla_sim_behaviors/src/plan_joint_spline_through_poses.cppsrc/vla_sim_behaviors/test/test_plan_joint_spline_through_poses.cpp
🚧 Files skipped from review as they are similar to previous changes (1)
- src/vla_sim/script/joint_command_bridge.py
|
danwahl
force-pushed
the
vla-sim-oracle-collection
branch
from
August 10, 2026 22:16
0b93b00 to
9520e3f
Compare
|
Adds the recording half of the config, so a replacement policy can be trained without leaving it: a scripted stacking oracle, per-prompt sweeps over the training layouts, and the joint-command bridge the Trainer needs to label `action` from commands rather than next states. vla_sim_behaviors carries the four Behaviors the oracle needs. Its SendGripperCommand test aborted the whole binary about one run in three: the stalling action server it stands up was destroyed while the executor could still dispatch its callbacks. The test now stops the executor first.
danwahl
force-pushed
the
vla-sim-oracle-collection
branch
from
August 11, 2026 19:09
9520e3f to
eb11687
Compare
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
[written (mostly) by AI]
Motivation
#815 ships
vla_simwith a checkpoint you can run, but no way to record the demonstrations behind one. This adds a scripted oracle that stacks the cubes and records itself, plus the randomized layouts it sweeps.Training stays out of scope. Epic PickNikRobotics/moveit_pro#20583 assigns it to the
train-modelClaude skill (PickNikRobotics/moveit_pro#20586), not to Pro, so the recipe behind the shipped checkpoint is parked onreference/vla-sim-train-recipefor that skill to build against.Part of PickNikRobotics/moveit_pro#20907, and the worked example that PickNikRobotics/moveit_pro#21138 documents. Targets 10.0.0.
How it works
Layouts.
keyframes.xmlcarries 360 train and 150 eval cube layouts, reachable by name through/mujoco_system/reset_keyframe. Every Objective resets to one before it runs, so the scene varies without anyone touching the simulator. The shipped checkpoint saw only the train layouts, which is what makes the eval set fair to score on.The oracle. Three ways in:
Run Cube-Stack Oracleperforms one stack with no recording. Quickest way to see whether a change to the scene or the planner still produces a clean demonstration.Collect Cube-Stack Demonstrationrecords one episode.Record Cube-Stack <held> On <target>Objectives each sweep the 60 training layouts drawn for their prompt, producing one dataset per prompt.All three are built from four new Behaviors in
vla_sim_behaviors:ComputeTopDownKeyposesderives the approach, grasp, lift, and place poses from where the cubes actually are, picking whichever of the cube's four equivalent yaws costs the arm least.PlanJointSplineThroughPosesfits one joint-space spline through them, so the recorded motion flows through the waypoints instead of stopping at each.SendGripperCommandsends a goal and succeeds as soon as it is dispatched, without waiting for the server to accept it, matching howExecutePolicydrives the gripper at deploy time.WaitForEpisodeStartholds the arm until the Trainer's recording marker lands, so the reset motion stays out of the episode.What gets recorded.
joint_command_bridge.pyfixes both halves of the recorded pair. Both defaults fail silently: the recording succeeds either way, and the damage surfaces only in the trained policy.action: the Trainer labels it from/joint_commandsand falls back to next-state labels when that topic is silent. Onlyquest_oculus_teleoppublishes it, so an Objective-driven recording would take the fallback without saying so. The bridge republishes the controller's reference trajectory there, so the datasets carry real commanded actions.observation.state: the default/joint_statescarries all 15 joints, 8 of them passive Robotiq linkage, so a dataset recorded from it trains against a state vector the deployed policy never sees. The bridge also publishes the 8 policy joints on/observed_joint_states, anddocker-compose.yamlpointsMOVEIT_PRO_TRAIN_JOINT_STATES_TOPICat it.config.yamlhosts the bridge throughadditional_agent_launch_file, so it comes up on both the dev and runtime paths.Manual verification
Collect Cube-Stack Demonstrationproduces an MCAP with aligned command, state, and three camera streams, and conversion labelsactionfrom commands rather than falling back to next states.pi05_kinova_gen3_cube_stack_simcheckpoint was trained on datasets recorded this way.colcon buildandcolcon testforvla_sim_behaviorsgreen against the Pro dev image: 51 C++ tests, 0 failures, plus 12 Python tests for the bridge. TheSendGripperCommandtest used to abort the whole binary about one run in three, because the stalling action server it stands up was destroyed while the executor could still dispatch its callbacks. Both gtest fixtures now stop the executor through a scope guard, so a failing assertion cannot skip it.pre-commit run --from-ref origin/main --to-ref HEADclean.Two commits, below.
mainhas moved ahead since the branch point, but only across SAM2/SAM3 and other config packages, so nothing it touches overlapssrc/vla_sim*and the PR stays mergeable.A note for reviewers on CI coverage
CI does not build or test these packages. The gating jobs run
colcon build --packages-up-to lab_sim(andhangar_sim) andcolcon test --packages-select lab_sim;vla_sim_behaviorsis a dependency of neither, so a green check on this PR says nothing about it. The build and test results above were produced by hand in the Pro dev image. Worth deciding separately whethervla_simshould join the integration-test matrix.