Integrate K=1–5 NMG benchmarking and differentiable rollout support - #623
Integrate K=1–5 NMG benchmarking and differentiable rollout support#623yangchen73 wants to merge 10 commits into
Conversation
yangchen73
changed the title
 |
There was a problem hiding this comment.
🟡 Changes recommended
Two critical issues remain: stale K=4 batched observation fixtures and unsupported joint-space cases being scheduled through NMG.
Get a fresh assessment by requesting another Copilot review.
Pull request overview
Adds K=1–5 NMG benchmarking, differentiable recurrent-state gradients, and NeuralPlanner alignment with K=5 ONNX policies.
Changes:
- Updates NeuralPlanner observation layouts and defaults.
- Adds functional-state gradient propagation through differentiable rollouts.
- Adds NMG capacity filtering, waypoint budgets, and benchmark configurations.
- Updates atomic scenarios, tests, and documentation.
File summaries
| File | Summary |
|---|---|
tests/sim/motion/planners/test_neural_planner.py |
Adds NeuralPlanner observation-dimension coverage. |
tests/gym/envs/test_differentiable_embodied_env.py |
Tests recurrent functional-state gradients. |
tests/benchmark/motion_generation/test_motion_generation_benchmark.py |
Tests NMG capacity and rollout budgets. |
tests/benchmark/motion_generation/test_atomic_task_benchmark.py |
Tests atomic benchmark configuration. |
scripts/benchmark/motion_generation/suites/smoke.yaml |
Updates smoke benchmark settings. |
scripts/benchmark/motion_generation/suites/coverage.yaml |
Updates coverage benchmark settings. |
scripts/benchmark/motion_generation/suites/atomic_franka_pgi_curobo.yaml |
Updates atomic benchmark configuration. |
scripts/benchmark/motion_generation/suites/atomic_franka_pgi_curobo_randomized.yaml |
Updates randomized atomic cases, including joint-waypoint cases. |
scripts/benchmark/motion_generation/suites/atomic_franka_pgi_curobo_pose_batch.yaml |
Updates pose-batch benchmark configuration. |
scripts/benchmark/motion_generation/scenarios/atomic_objects.py |
Updates collision and physics configuration. |
scripts/benchmark/motion_generation/runner.py |
Records unsupported benchmark cases. |
scripts/benchmark/motion_generation/planners/nmg_onnx.py |
Adds capacity and waypoint-budget handling; joint-space cases must be excluded before scheduling. |
scripts/benchmark/motion_generation/planners/base.py |
Updates planner base interfaces. |
embodichain/lab/sim/motion/planners/neural_planner.py |
Aligns observation layouts with K=5 policies; existing K=4 batched fixtures still expect width 101 instead of 97. |
embodichain/lab/sim/diff/bridge.py |
Bridges functional-state gradients through Warp and PyTorch. |
embodichain/lab/gym/envs/differentiable_env.py |
Forwards recurrent state through environment steps. |
embodichain_tasks/embodichain_tasks/special/franka_reach_apg.py |
Updates the differentiable action-kernel contract. |
agent_context/topics/motion-planning/motion-planning.md |
Documents updated motion-planning contracts. |
agent_context/topics/differentiable-env/differentiable-env.md |
Documents differentiable environment changes. |
Review details
- Files reviewed: 19/19 changed files
- Comments generated: 2
- Review effort level: Lite
💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
There was a problem hiding this comment.
🔵 Needs a closer look
Three moderate validation issues remain in functional-state handling and rollout configuration.
Review details
Suppressed comments (3)
Previously missed (1) — in code that hasn't changed since the last review.
embodichain/lab/sim/diff/bridge.py:137
- Passing a non-tensor functional state (for example,
Nonefrom_functional_state_tensors()) fails here while evaluating.shape, so the explicitTypeErrorbelow is unreachable and callers receive anAttributeErrorinstead. Validate all state inputs before derivingsaved_state_shapesso invalid state-hook implementations fail with the documented, actionable error.
scripts/benchmark/motion_generation/planners/nmg_onnx.py:101
- [P2] Reject non-positive rollout budgets before passing them to the planner. With
steps_per_waypoint: 0(or a non-positivemax_steps), this returns 0/negative, butNeuralPlanner.plan()usesoptions.max_steps or self._max_steps, so zero silently falls back to the configured 150-step rollout while a negative value produces no rollout. Validate both config values as positive here (or in suite validation) so the per-waypoint budget cannot be bypassed.
steps_per_waypoint = int(self.spec.config.get("steps_per_waypoint", 30))
configured_max = int(self.spec.config.get("max_steps", 150))
return min(configured_max, int(case.num_waypoints) * steps_per_waypoint)
scripts/benchmark/motion_generation/suites/smoke.yaml:46
- [P2] Validate the new NMG rollout controls at suite-load time.
SuiteCfg.validate_benchmark()currently checks onlypos_eps/rot_eps, sosteps_per_waypoint: 0or a non-positivemax_stepsis accepted and reaches_case_max_steps(); this either silently falls back throughNeuralPlanOptions(max_steps=0)or produces an empty rollout, making the benchmark budget invalid instead of failing configuration validation. Add positive-integer validation fornum_waypoints,steps_per_waypoint, andmax_steps(and apply it to the other updated suites).
num_waypoints: 5
steps_per_waypoint: 30
max_steps: 150
- Files reviewed: 20/20 changed files
- Comments generated: 0 new
- Review effort level: Lite
|
The NMG integration direction is good, but I recommend the following changes before merge:
|
Support Joint case, and the missing validation and contract tests are now covered. |
| """Return the unified NMG constraint-observation width.""" | ||
| n = int(num_waypoints) | ||
| dim = 7 + 7 + n * (3 + 4 + 7 + 5) + 7 + n | ||
| dim = 7 + 7 + n * (3 + 4 + 7 + 5) + 7 |
There was a problem hiding this comment.
It woould be better to use config file for dim definition. Since we may change it later
There was a problem hiding this comment.
I stored the obs layout when export to onnx file. Then read the layout from EmbodiChain's side.
There was a problem hiding this comment.
🟡 Changes recommended
Unresolved critical and moderate findings remain.
Get a fresh assessment by requesting another Copilot review.
Review details
Suppressed comments (4)
scripts/benchmark/motion_generation/config.py:255
- The new NMG settings validation covers waypoint and rollout budgets plus
pos_eps/rot_eps, but it still acceptsjoint_eps <= 0. This PR addsordered_joint_waypoints; with the planner's strictjoint_dist < self._joint_epscheck, zero or negative values make every joint waypoint unreachable (even an exact match), so benchmark configuration should reject this value just like the other convergence thresholds.
for nmg in (spec for spec in self.planners if spec.adapter == "nmg_onnx"):
for field, default in (
("num_waypoints", 5),
("steps_per_waypoint", 30),
("max_steps", 150),
):
_positive_int(
nmg.config.get(field),
name=f"NMG {field}",
default=default,
)
scripts/benchmark/motion_generation/planners/nmg_onnx.py:151
- The new JOINT_MOVE branch contradicts the benchmark design, which still says NMG supports only EEF_MOVE and that joint-space cases are cuRobo-only (
scripts/benchmark/motion_generation/BENCHMARK_DESIGN.md:370-372). Update that contract and any leaderboard-track guidance, or avoid advertising joint NMG support; otherwise the documented eligibility rules disagree with the runner.
if motion_validity == "ordered_joint_waypoints":
targets = [
PlanState.from_qpos(
case.reference_qpos[:, index], move_type=MoveType.JOINT_MOVE
)
for index in range(case.num_waypoints)
]
elif motion_validity == "ordered_cartesian_waypoints":
scripts/benchmark/motion_generation/planners/nmg_onnx.py:91
validate_benchmark()treats an explicitnullvalue as “use the default” via_positive_int, but this adapter casts the raw mapping value directly. A valid suite containingnum_waypoints: nulltherefore passes validation and then raisesTypeErrorhere (and the same mismatch exists forsteps_per_waypoint/max_stepsin_case_max_steps()andbuild()). ResolveNoneto the same defaults before converting these fields so the validated configuration can actually run.
capacity = int(self.spec.config.get("num_waypoints", 5))
scripts/benchmark/motion_generation/runner.py:422
- This introduces a new
unsupported_capacityfailure code, but the benchmark's documented stable failure taxonomy inscripts/benchmark/motion_generation/BENCHMARK_DESIGN.md:822-841does not include it. Either use the documentedunsupported_capabilitycode or update the taxonomy and downstream artifact/report contract so consumers do not encounter an undocumented value.
self._record_unavailable(
writer,
metadata,
case,
reason or "case is outside planner capacity",
failure_code="unsupported_capacity",
)
- Files reviewed: 21/21 changed files
- Comments generated: 2
- Review effort level: Lite
There was a problem hiding this comment.
🟡 Changes recommended
Unresolved critical and moderate findings block approval.
Get a fresh assessment by requesting another Copilot review.
Review details
Suppressed comments (7)
embodichain/lab/gym/envs/differentiable_env.py:198
- [P2] Preserve
kernel_argswhen adding recurrent state.NewtonStepFunccalls this wrapper asaction_kernel(action, tape, *kernel_args, *state_wps), but_innertreats every argument aftertapeas a state array. AnyDifferentiableEnvsubclass that overrides the existingkernel_argsentry and also returns functional state will therefore pass the static kernel arguments into_apply_action_kernelas state (and shift the real state positions), causing incorrect Warp inputs or an arity failure. Keep the kernel-argument count/closure separate from the recurrent state arguments, or explicitly reject non-emptykernel_argsfor this environment.
def _inner(action_wp: Any, tape: Any, *state_wps: Any) -> None:
env._apply_action_kernel(action_wp, *state_wps, tape=tape)
scripts/benchmark/motion_generation/config.py:255
- [P2] Apply NMG CLI overrides to adapter-selected IDs — validation now recognizes every planner with
adapter == "nmg_onnx"(including the renamed ID used by the new validation test), but_apply_overrides()still searches only forspec.id == "nmg". Selecting a valid renamed NMG planner therefore leaves--nmg-onnx-pathand the NMG epsilon overrides unapplied; the planner is then skipped for a missing model path or runs with stale thresholds. Apply these overrides by adapter, or consistently require the canonical ID.
for nmg in (spec for spec in self.planners if spec.adapter == "nmg_onnx"):
for field, default in (
("num_waypoints", 5),
("steps_per_waypoint", 30),
("max_steps", 150),
):
_positive_int(
nmg.config.get(field),
name=f"NMG {field}",
default=default,
)
scripts/benchmark/motion_generation/planners/nmg_onnx.py:92
- [P2] Resolve nullable NMG defaults before casting —
validate_benchmark()treatsNoneas the default fornum_waypoints,steps_per_waypoint, andmax_steps, so a YAML value such asnum_waypoints: nullis accepted. This call then executesint(None)and crashes before capability filtering; the same pattern exists in_case_max_steps()andbuild(). NormalizeNoneto the documented defaults in the adapter or reject null values during validation.
capacity = int(self.spec.config.get("num_waypoints", 5))
if case.num_waypoints > capacity:
scripts/benchmark/motion_generation/planners/nmg_onnx.py:165
- This per-case rollout cap is applied only when the runner calls
NmgOnnxAdapter.plan. Atomic-task cases take a different path:AtomicTaskScenario.plan_casecompiles anAtomicActionEngine, whose primitives passMotionPolicy.plan_opts=Noneto the adapter-ownedMotionGenerator, soNeuralPlannerfalls back to its configured 150 steps. A one-waypoint atomic case therefore still gets 150 policy steps instead of the 30-per-waypoint budget asserted by the new test. Thread a case-specificNeuralPlanOptions(max_steps=...)through the atomic invocation path, or explicitly scope this budget to direct planner cases and test that scope.
plan_opts=NeuralPlanOptions(max_steps=self._case_max_steps(case)),
scripts/benchmark/motion_generation/planners/nmg_onnx.py:147
- This new
JOINT_MOVEroute makes the benchmark contract inscripts/benchmark/motion_generation/BENCHMARK_DESIGN.mdstale: its capability boundary and primary-track text still say current NMG supports only EEF waypoints and that joint-space cases are cuRobo-only (lines 370-372). Update the design/track capability documentation in this change so the newly eligible joint cases and resulting coverage are not misinterpreted.
PlanState.from_qpos(
case.reference_qpos[:, index], move_type=MoveType.JOINT_MOVE
scripts/benchmark/motion_generation/runner.py:422
- [P2] Keep unsupported-capacity cases out of the success denominator. This new path records an availability row but leaves the case in the full
caseslist;aggregation.pythen treats the missing measured outcomes as zero in_case_macro_rate/_case_macro_primary_rateand includes the case in the expected coverage count. Selecting a suite with a waypoint count above NMG capacity therefore lowers NMG's success rates and makes it ineligible instead of applying the benchmark contract that unsupported cases are excluded from the track denominator (scripts/benchmark/motion_generation/BENCHMARK_DESIGN.md:672-680). Carry an explicit supported-case set into aggregation, while retaining separate coverage reporting.
self._record_unavailable(
writer,
metadata,
case,
reason or "case is outside planner capacity",
failure_code="unsupported_capacity",
)
scripts/benchmark/motion_generation/scenarios/atomic_objects.py:186
- [P2] Do not truncate invalid mesh hull counts. A YAML value such as
max_convex_hull_num: 1.5is silently converted to1here and selects a single convex hull, while the underlyingMeshCollisionCfg.max_hullscontract requires an integer (and rejects booleans/non-integers). Validate the raw value before converting it so malformed benchmark scene configuration fails instead of changing the collision approximation.
max_hulls = int(config.get("max_convex_hull_num", 16))
- Files reviewed: 21/21 changed files
- Comments generated: 1
- Review effort level: Lite
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
There was a problem hiding this comment.
🔵 Needs a closer look
Seven unresolved moderate findings remain across planner, benchmark, and differentiable-environment behavior.
Review details
Suppressed comments (7)
Previously missed (1) — in code that hasn't changed since the last review.
scripts/benchmark/motion_generation/runner.py:422
- [P2] Exclude unsupported cases from the planner's success denominator. These availability records are not excluded by aggregation:
aggregate_results()still receives the fullself.cases, and_case_macro_rate()treats every case without a measured outcome as 0.0. Thus an NMG case over the five-waypoint capacity (for example the eight-waypoint twist case) is reported as unsupported but also lowers success rates and prevents coverage/eligibility, contrary to the benchmark contract that unsupported cases are excluded from the track denominator; carry the per-planner supported-case set into aggregation (while retaining coverage reporting) or otherwise filter unsupported case IDs there.
embodichain/lab/sim/motion/planners/neural_planner.py:284
- Changing the planner default to five slots leaves the existing example CLI inconsistent:
examples/sim/motion/planners/neural_planner.pystill accepts--num-waypointsvalues up to eight, but itsNeuralPlannerCfgconstruction does not pass that argument, so requesting six or more targets leaves the planner at capacity five and_parse_waypointsraises. Forward the CLI capacity into the config or reject values above the supported policy capacity.
num_waypoints: int = 5
"""Number of constraint slots encoded by the ONNX policy."""
scripts/benchmark/motion_generation/config.py:259
joint_epsis passed toNeuralPlannerand convergence usesjoint_dist < self._joint_eps, but this validation only rejects non-positive position and rotation tolerances. A zero or negative joint tolerance is accepted and causes every joint waypoint to miss until the rollout budget is exhausted; validate it here with the other NMG tolerances.
if float(nmg.config.get("pos_eps", 0.01)) <= 0.0:
raise ValueError("NMG pos_eps must be > 0.")
if float(nmg.config.get("rot_eps", 0.1)) <= 0.0:
raise ValueError("NMG rot_eps must be > 0.")
scripts/benchmark/motion_generation/planners/base.py:108
- This new generic capability gate treats
capabilitiesas proof that a case is representable, butCuroboAdapteradvertisesjoint_waypointwhile itsplan()still converts every case throughPlanState.from_xpos(case.target_waypoints[...])and prepares only the EEF backend. Atomicordered_joint_waypointscases will therefore be counted as supported while being planned in Cartesian space; either add the JOINT_MOVE path/backend preparation or remove that capability until it is implemented.
def supports_case(self, case: BenchmarkCase) -> tuple[bool, str | None]:
"""Return whether one manifest case is representable without mutation."""
capability, reason = _case_motion_capability(case)
if capability is None:
return False, reason
scripts/benchmark/motion_generation/planners/nmg_onnx.py:97
- [P2] Normalize nullable NMG defaults before consuming them.
SuiteCfg.validate_benchmark()deliberately acceptsNonefornum_waypoints,steps_per_waypoint, andmax_stepsas requests for their defaults, but this adapter callsint(self.spec.config.get(...))insupports_case,_case_max_steps, andbuild; a YAML entry such asmax_steps: nulltherefore passes validation and then crashes withTypeErrorinstead of using 150. Resolve these fields once with the same defaults used by validation and reuse the normalized values.
capacity = int(self.spec.config.get("num_waypoints", 5))
if case.num_waypoints > capacity:
return (
False,
f"NMG policy supports at most {capacity} waypoints; case has "
+ f"{case.num_waypoints} and was not truncated or split.",
)
scripts/benchmark/motion_generation/planners/nmg_onnx.py:104
- This per-case budget is only applied by
NmgOnnxAdapter.plan(). Atomic benchmark cases use the adapter-ownedMotionGeneratorthroughAtomicActionEngine; theirMotionPolicyleavesplan_opts=None, soNeuralPlannerfalls back to the fixedcfg.max_steps(150) instead ofmin(max_steps, num_waypoints * steps_per_waypoint). A one-waypoint atomic case can therefore run 150 steps rather than the configured 30, changing its success and latency measurement. Propagate a case-specificNeuralPlanOptionsinto that path or expose an adapter hook for the budget.
def _case_max_steps(self, case: BenchmarkCase) -> int:
"""Match NMG's training/evaluation budget of 30 steps per waypoint."""
steps_per_waypoint = int(self.spec.config.get("steps_per_waypoint", 30))
configured_max = int(self.spec.config.get("max_steps", 150))
return min(configured_max, int(case.num_waypoints) * steps_per_waypoint)
tests/gym/envs/test_differentiable_embodied_env.py:175
- [P1] Make the default-hook regression test use the keyword-only tape contract. The wrapper now invokes the hook as
env._apply_action_kernel(..., tape=tape), but this test callback still declarestapepositionally. The call at line 187 therefore raisesTypeError: unexpected keyword argument 'tape'before the assertion; declare the callback with*, tape(or otherwise accept the keyword).
def test_environment_action_hook_receives_no_state_by_default() -> None:
"""The default bridge adapter supplies only action and keyword tape."""
- Files reviewed: 22/22 changed files
- Comments generated: 0 new
- Review effort level: Lite
There was a problem hiding this comment.
🔵 Needs a closer look
Moderate issues remain in backend preparation, null/default handling, fixed-batch filtering, and unsupported-case aggregation.
Review details
Suppressed comments (5)
scripts/benchmark/motion_generation/planners/curobo.py:130
- This prepares only the move type of the single case passed by the runner. The atomic benchmark contains both EEF and joint cases, while the runner calls
prepare()only forsupported_cases[0]; with the current ordering the joint backend is first created and warmed lazily inside the first joint plan. That charges backend construction/warmup to planning phases and makes the mixed-modality latency results invalid. Prepare every distinct(batch_size, move_type)used by the supported cases before entering the cold/warm/measured loops.
return planner.prepare_backend(
control_part=self.context.control_part,
batch_size=case.batch_size,
move_type=self._case_move_type(case),
scripts/benchmark/motion_generation/planners/nmg_onnx.py:92
- Suite validation treats an explicitly present
nullas the default (config.py:_positive_int), but this.getreturnsNonewhen the key is present andint(None)raises. The same failure occurs forsteps_per_waypoint/max_stepsin_case_max_stepsand formax_steps/num_waypointsinbuild, so a valid YAML override such asnum_waypoints: nullfails during case filtering instead of using the documented default. Normalize present-but-null values through one resolver before all these reads.
capacity = int(self.spec.config.get("num_waypoints", 5))
if case.num_waypoints > capacity:
scripts/benchmark/motion_generation/planners/nmg_onnx.py:98
- [P2] Reject fixed-batch exports before timed planning.
_OnnxPolicyaccepts static-batch ONNX graphs but raises whenever a case batch differs from that fixed size. Becausesupports_case()runs beforebuild()and only checks waypoint capacity, a fixed-B=1 model is marked supported for the B=8/64 suites and fails later as a timed planner exception. Make case filtering honor the model's fixed batch size, or classify such models as non-batched before execution.
def supports_case(self, case: BenchmarkCase) -> tuple[bool, str | None]:
"""Reject unsupported modalities or sequences beyond policy capacity."""
supported, reason = super().supports_case(case)
if not supported:
return supported, reason
capacity = int(self.spec.config.get("num_waypoints", 5))
if case.num_waypoints > capacity:
return (
False,
f"NMG policy supports at most {capacity} waypoints; case has "
+ f"{case.num_waypoints} and was not truncated or split.",
)
return True, None
scripts/benchmark/motion_generation/runner.py:422
- [P1] Exclude unsupported cases from success denominators
When a manifest contains a case beyond the NMG capacity, this records an unsupported_capacity availability record but leaves the case in the manifest passed to aggregation. The _case_macro_* aggregators treat missing measured outcomes as 0.0, so the newly supported partial-capacity path penalizes NMG's success metrics as if the case had failed, contrary to BENCHMARK_DESIGN.md:672-680; planner errors should remain failures, but unsupported cases must be filtered from the success denominator. Propagate the unsupported case IDs/status into aggregation and retain them only for coverage/reporting.
self._record_unavailable(
writer,
metadata,
case,
reason or "case is outside planner capacity",
failure_code="unsupported_capacity",
)
scripts/benchmark/motion_generation/runner.py:415
- [P2] Prepare every backend modality before measured plans
CuroboAdapter now has distinct cached EEF and JOINT backends, but this lifecycle later prepares only first_case. In a supported batch containing Cartesian cases followed by joint cases, the first joint plan() lazily builds and warms its backend inside the measured phase, so backend construction/CUDA graph setup contaminates planning latency. Prepare each distinct supported case shape/modality before the cold and measured phases (or expose an adapter-level multi-backend prepare hook).
supported, reason = adapter.supports_case(case)
if supported:
supported_cases.append(case)
continue
- Files reviewed: 22/22 changed files
- Comments generated: 0 new
- Review effort level: Lite
There was a problem hiding this comment.
🔵 Needs a closer look
Resolve joint-epsilon validation, adapter-based NMG overrides, and outdated benchmark design documentation.
Review details
Suppressed comments (3)
scripts/benchmark/motion_generation/config.py:259
- The new benchmark config forwards
joint_epsintoNeuralPlanner, but this validation block only rejects non-positivepos_epsandrot_eps. A suite can therefore setjoint_eps <= 0; because_is_active_reached()uses the strict comparisonjoint_dist < self._joint_eps, every joint-waypoint case will fail to converge and be recorded as a planner failure instead of being rejected during suite validation. Validatejoint_epshere with the same> 0rule (and add it to the precision-validation test).
if float(nmg.config.get("rot_eps", 0.1)) <= 0.0:
raise ValueError("NMG rot_eps must be > 0.")
scripts/benchmark/motion_generation/config.py:245
- [P2] Apply NMG overrides by adapter, not the literal planner id. This validation now intentionally accepts every
adapter == "nmg_onnx"spec (including the renamed IDs exercised by the tests), butrun_benchmark._apply_overrides()still searches only forspec.id == "nmg"; a valid suite using another ID silently ignores--nmg-onnx-path,--nmg-pos-eps, and--nmg-rot-eps, so the selected NMG planner can remain unavailable or use stale settings. Use the same adapter-based lookup (and handle duplicate NMG adapters explicitly if needed).
for nmg in (spec for spec in self.planners if spec.adapter == "nmg_onnx"):
scripts/benchmark/motion_generation/planners/base.py:40
- [P2] Update the benchmark design text with the new joint capability. This mapping makes NMG joint-waypoint cases supported, but
BENCHMARK_DESIGN.mdstill says current NMG supports onlyEEF_MOVEand that joint-space cases are cuRobo-only, so the authoritative protocol now contradicts the shipped adapter and suites.
"ordered_joint_waypoints": "joint_waypoint",
- Files reviewed: 22/22 changed files
- Comments generated: 0 new
- Review effort level: Lite
Description
Add K=1–5 NMG benchmark support, complete functional-state gradients for differentiable rollouts, and align
NeuralPlannerwith K=5 ONNX policies.Type of change
Checklist
black .command to format the code base.python docs/scripts/check_api_docs.py), if applicable.