Skip to content

Add full-horizon waypoint APG training foundation - #570

Draft
yuecideng wants to merge 3 commits into
feature/newton-physics-backendfrom
codex/nmg-training-foundation
Draft

Add full-horizon waypoint APG training foundation#570
yuecideng wants to merge 3 commits into
feature/newton-physics-backendfrom
codex/nmg-training-foundation

Conversation

@yuecideng

@yuecideng yuecideng commented Aug 31, 2026

Copy link
Copy Markdown
Contributor

Description

Adds the reusable RL foundation needed to reproduce the full-horizon NMG training contract from neural_motion_generator PR #6.

  • Adds scheduled variable-horizon differentiable rollouts with full K-dependent horizons, reciprocal-K objective scaling, first-done masking, gradient accumulation, and exact optimizer-update budgeting.
  • Adds overflow-safe per-environment action-adjoint clipping, semantic-mask-aware Welford observation normalization, deterministic lightweight-runtime seeding, and checkpoint support for the new runtime state.
  • Extends the lightweight training CLI, evaluation path, public API docs, project context, and focused tests.

The waypoint Transformer model and Franka task implementation are intentionally excluded from this PR. They live together on the pushed downstream branch codex/franka-waypoint-nmg; no PR has been opened for that branch. This branch depends on the Newton integration already present in the target branch and adds no package dependencies.

Type of change

  • Enhancement (non-breaking change which improves an existing functionality)
  • New feature (non-breaking change which adds functionality)
  • Documentation update

Screenshots

Not applicable.

Validation

  • black . — 869 files unchanged
  • python docs/scripts/check_api_docs.py — 1876/1876 public exports documented
  • pytest -q tests/learning tests/docs/test_check_api_docs.py — 77 passed, 3 deselected
  • Downstream Franka parity branch compared PR Update default physics mode #6 on CPU and CUDA: joint, Cartesian, and mixed task tensors/one-step gradients had zero max error; a complete 30-step APG update also had zero error for the action trace, loss, gradient norm, every updated actor parameter, and observation-normalizer state

Checklist

  • I have run the black . command to format the code base.
  • I have made corresponding changes to the documentation.
  • Public API changes are reflected in the API docs (python docs/scripts/check_api_docs.py).
  • I have added tests that prove my fix is effective or that my feature works.
  • Dependencies have been updated, if applicable (no new dependencies).

@yuecideng yuecideng added enhancement New feature or request rl Features related to reinforcement learning model neual model related functionality motion gen Things related to motion generation for robot labels Aug 31, 2026
@yuecideng yuecideng removed the model neual model related functionality label Aug 31, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

enhancement New feature or request motion gen Things related to motion generation for robot rl Features related to reinforcement learning

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant