Improve PPO training stability

[acrlw]

Description

This PR improves EmbodiChain's PPO implementation while preserving the existing default observation behavior.

• Add separate actor and critic observation groups with independently checkpointed running normalization.
• Add adaptive learning rates from Gaussian KL, clipped value loss, action-mean bounds, and finite-gradient rejection.
• Apply timeout bootstrapping at episode boundaries.
• Preserve policy distribution parameters and critic observations through rollout collection and evaluation.
• Add update-based checkpoint intervals and keep the update counter in checkpoints and training summaries.
• Cover normalization, grouped observations, minibatch partitioning, timeout bootstrapping, checkpoint restoration, and PPO update stability with focused tests.

No new dependencies are required.

Validation results

• pytest tests/learning -q: 78 passed, 3 deselected.
• Local CLI smoke test: completed 2 PPO updates and 64 environment steps, including evaluation and checkpoint creation; checkpoint counters, normalization state, and tensors were valid.
• black --check --diff --color ./: passed with Black 26.3.1.
• python docs/scripts/check_api_docs.py: 1863/1863 public exports documented.
• Changed-file structure checks passed for headers, future annotations, and public-module exports across 18 Python files.

Type of change

• Enhancement (non-breaking change which improves an existing functionality)

Screenshots

Not applicable.

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), if applicable
• I have added tests that prove my fix is effective or that my feature works
• Dependencies have been updated, if applicable.

[greptile-apps]

Confidence Score: 5/5

The PR appears safe to merge.

Summary

• Adds separate actor and critic observation groups with independently checkpointed normalization.
• Adds adaptive KL-based learning rates, clipped value loss, action-mean penalties, finite-gradient rejection, and complete minibatch partitioning.
• Adds timeout bootstrapping, update-based checkpoint intervals, restored update counters, and focused stability tests.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  E[Environment observations] --> G[Resolve actor and critic groups]
  G --> A[Actor observations]
  G --> C[Critic observations]
  A --> AN[Actor normalizer]
  C --> CN[Critic normalizer]
  AN --> P[Gaussian policy]
  CN --> V[Value function]
  P --> R[Rollout buffer]
  V --> R
  R --> PPO[PPO minibatch updates]
  PPO --> KL[Adaptive KL learning rate]
  PPO --> CKPT[Checkpoint policy, optimizer, normalizers, and counters]

Loading

_{Reviews (4) · Last reviewed commit: "feat(rl): improve PPO training stability"}