Add full-horizon waypoint APG training foundation

agent_context/MAP.yaml

@@ -490,6 +490,13 @@ topics:
       - reward
       - RolloutKind
       - DifferentiableTrainer
+      - ScheduledDifferentiableVecEnv
+      - DifferentiableRolloutSpec
+      - complete rollout
+      - full horizon
+      - gradient accumulation
+      - action adjoint
+      - observation normalization
       - SyncCollector
       - learning_env
       - gym_config
@@ -503,6 +510,8 @@ topics:
       - embodichain/learning/rl/train.py
       - embodichain/learning/rl/env.py
       - embodichain/learning/rl/evaluation.py
+      - embodichain/learning/rl/gradients.py
+      - embodichain/learning/rl/normalization.py
       - embodichain/learning/rl/routing.py
       - embodichain/learning/rl/differentiable_trainer.py
       - embodichain/learning/rl/utils/config.py

agent_context/topics/rl-learning/rl-learning.md

@@ -83,8 +83,16 @@ as a mapping with `name` and `cfg`.
 `DifferentiableVecEnv.detach_state()` as the truncated-backpropagation
 boundary.
 
+Variable-horizon APG environments may also implement
+`ScheduledDifferentiableVecEnv.prepare_differentiable_rollout()`. The returned
+`DifferentiableRolloutSpec` selects the next reset's complete horizon,
+per-environment objective scale, and scalar rollout metadata.
+
 This path supports both standard algorithms and differentiable algorithms,
 but currently rejects distributed training and environment profiling.
+`trainer.seed` seeds Python, NumPy, Torch, CUDA, and Warp before environment
+construction. Set `trainer.torch_deterministic: true` when a reference run
+requires PyTorch deterministic algorithms in addition to seeded sampling.
 
 ## Rollout and Trainer Routing
 
@@ -114,9 +122,24 @@ use it as padding. The collector writes into the preallocated rollout and the
 algorithm consumes it after collection.
 
 The differentiable path does not copy transitions into the standard buffer.
-It preserves the action-to-reward autograd graph across short segments.
-`segment_length` sets TBPTT boundaries, while `update_horizon` controls
-how many environment steps contribute to one optimizer update.
+It has two explicitly configured modes:
+
+- `rollout_mode: segmented` preserves the action-to-reward graph within TBPTT
+  segments. `segment_length` sets detach boundaries and `update_horizon`
+  controls the optimizer budget.
+- `rollout_mode: complete` resets before every independent microbatch and
+  preserves one graph across the entire environment-provided horizon. It masks
+  rewards after the first done, applies the rollout's objective scale, and
+  averages `gradient_accumulation_steps` full trajectories before one optimizer
+  step. It never shortens a scheduled horizon to satisfy `total_timesteps`.
+  When no explicit timestep budget is configured, CLI `iterations` maps to an
+  exact optimizer-update budget so changing the K distribution does not change
+  the number of gradient steps.
+
+Complete mode can clamp actions to the environment space and install a
+per-environment action-adjoint norm hook. Non-finite adjoint rows are zeroed;
+finite rows are clipped independently using an overflow-safe norm. APG also
+supports a pre-clip policy-gradient safety limit that skips unsafe updates.
 
 ## Component Ownership
 
@@ -127,6 +150,8 @@ how many environment steps contribute to one optimizer update.
 | Standard rollout storage and views | `buffer/` |
 | Standard and differentiable collection | `collector/` |
 | Policy interface, actor-critic, actor-only, MLP builder | `models/` |
+| Running observation statistics | `normalization.py` |
+| Batched action-adjoint stabilization | `gradients.py` |
 | Standard collect/update loop | `utils/trainer.py` |
 | Differentiable TBPTT/update loop | `differentiable_trainer.py` |
 | Shared completed-episode evaluation | `evaluation.py` |
@@ -153,8 +178,10 @@ Evaluation uses an independent environment and
 terminal metrics, temporarily switches the policy to evaluation mode, and
 restores its prior mode.
 
-Checkpoints include policy parameters, trainer counters, best-evaluation
-state, and optimizer or LR-scheduler state when present.
+Checkpoints include policy parameters, trainer and complete-rollout counters,
+best-evaluation state, observation-normalizer state when enabled, and optimizer
+or LR-scheduler state when present. Evaluation reuses the frozen training
+normalizer without updating its statistics.
 
 On the simulator path, distributed mode initializes NCCL, assigns one CUDA
 device per local rank, wraps the policy in
@@ -199,7 +226,11 @@ example is an experimental gradient reference, not a general simulator task.
 2. Register the factory with `@register_learning_env`.
 3. Ensure finished rows auto-reset while returning terminal reward/done with
    the next initial observation.
-4. Add an official config under
+4. For variable complete APG rollouts, implement
+   `ScheduledDifferentiableVecEnv` and return the full non-truncated horizon.
+5. Expose `observation_normalize_mask` when semantic dimensions must remain
+   raw during normalization.
+6. Add an official config under
    `embodichain_tasks/configs/tasks/<domain>/<task>/agents/` when it is a
    bundled task.
 
@@ -217,6 +248,10 @@ components.
 - The standard buffer holds at most one unconsumed rollout.
 - APG must retain differentiable rewards until its optimizer boundary;
   `detach_state()` must not reset or resample the task.
+- Complete APG mode must reset once per independent rollout, detach only after
+  backward, and exclude post-terminal auto-reset rewards from its objective.
+- Observation normalization statistics stay frozen throughout each complete
+  rollout and semantic mask/type fields remain unnormalized.
 - GRPO environment count must satisfy its grouping contract.
 - Evaluation must use completed episodes and an independent environment.
 - Only rank zero owns external logging and checkpoints in distributed runs.
@@ -232,6 +267,8 @@ components.
 | Policy dimension mismatch | Policy config disagrees with the built environment's observation or action space |
 | Standard buffer is already full | A rollout was started before the previous one was consumed with `get()` |
 | APG gradients disappear | Actions were sampled under `no_grad`, transitions were copied/detached, or the state was detached too early |
+| Long-horizon APG accuracy is lower than the reference | `rollout_mode` is still `segmented`, the scheduled horizon was truncated, return scaling is missing, or observation normalization differs |
+| One environment poisons every APG row | Action-adjoint clipping is disabled or non-finite row filtering is bypassed |
 | GRPO reshape or grouping fails | `num_envs` is not divisible by `group_size` |
 | Evaluation never completes | The environment does not emit completed asynchronous episodes or terminal metrics correctly |
 | Output/checkpoint directories diverge across ranks | Distributed run metadata was not coordinated through rank zero |

docs/source/api_reference/embodichain/embodichain.learning.rl.algo.rst

@@ -33,6 +33,7 @@ policy, while :func:`compute_gae` provides generalized advantage estimation.
       build_algo
       get_registered_algo_names
       compute_gae
+      complete_discounted_return
       segmented_discounted_return
 
 .. automodule:: embodichain.learning.rl.algo

docs/source/api_reference/embodichain/embodichain.learning.rl.models.rst

@@ -36,4 +36,3 @@ through :func:`register_policy` / :func:`get_policy_class`.
    :members:
    :undoc-members:
    :show-inheritance:
-

docs/source/api_reference/embodichain/embodichain.learning.rl.rst

@@ -27,12 +27,18 @@ collection logic, policy/model builders, and training entry points.
 
       DifferentiableTrainer
       DifferentiableTrainerCfg
+      DifferentiableRolloutSpec
       DifferentiableVecEnv
       LearningVecEnv
+      ScheduledDifferentiableVecEnv
+      RunningObservationNormalizer
+      BatchedGradientNormStats
       build_learning_env
+      clip_batched_gradient_norm
       evaluate_episodes
       get_trainer_class
       register_learning_env
+      stratified_rollout_value
 
 Algorithms
 ----------
@@ -58,6 +64,22 @@ Evaluation
    :undoc-members:
    :show-inheritance:
 
+Gradient Stabilization
+----------------------
+
+.. automodule:: embodichain.learning.rl.gradients
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+Observation Normalization
+-------------------------
+
+.. automodule:: embodichain.learning.rl.normalization
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
 Routing
 -------
 

docs/source/api_reference/public_api.rst

@@ -1748,6 +1748,7 @@ embodichain.learning.rl.algo.apg
 
    APG
    APGCfg
+   complete_discounted_return
    segmented_discounted_return
 
 embodichain.learning.rl.algo.base
@@ -1848,6 +1849,18 @@ embodichain.learning.rl.experimental.newton.train_planar_reach
    NewtonPlanarReachTrainingCfg
    train_planar_reach
 
+embodichain.learning.rl.gradients
+---------------------------------
+
+Row-wise action-adjoint clipping and its rollout-level diagnostics.
+
+.. currentmodule:: embodichain.learning.rl.gradients
+
+.. autosummary::
+
+   BatchedGradientNormStats
+   clip_batched_gradient_norm
+
 embodichain.learning.rl.models.actor_critic
 -------------------------------------------
 
@@ -1875,6 +1888,15 @@ embodichain.learning.rl.models.policy
 
    Policy
 
+embodichain.learning.rl.normalization
+-------------------------------------
+
+.. currentmodule:: embodichain.learning.rl.normalization
+
+.. autosummary::
+
+   RunningObservationNormalizer
+
 embodichain.learning.rl.utils.optimizer
 ---------------------------------------
 

embodichain/learning/rl/__init__.py

@@ -19,6 +19,8 @@
 Algorithms (PPO/GRPO), rollout buffers, collectors, policy/model builders, and the training entry point; rollout data flows as ``TensorDict`` objects.
 """
 
+from __future__ import annotations
+
 from . import algo
 from . import buffer
 from . import models
@@ -29,26 +31,37 @@
 )
 from .env import (
     DifferentiableObservation,
+    DifferentiableRolloutSpec,
     DifferentiableVecEnv,
     LearningVecEnv,
     build_learning_env,
     get_registered_learning_env_names,
     register_learning_env,
+    ScheduledDifferentiableVecEnv,
+    stratified_rollout_value,
 )
 from .evaluation import evaluate_episodes
+from .gradients import BatchedGradientNormStats, clip_batched_gradient_norm
+from .normalization import RunningObservationNormalizer
 from .routing import get_trainer_class
 
 __all__ = [
     "DifferentiableObservation",
     "DifferentiableTrainer",
     "DifferentiableTrainerCfg",
+    "DifferentiableRolloutSpec",
     "DifferentiableVecEnv",
+    "BatchedGradientNormStats",
     "LearningVecEnv",
+    "RunningObservationNormalizer",
+    "ScheduledDifferentiableVecEnv",
     "build_learning_env",
+    "clip_batched_gradient_norm",
     "evaluate_episodes",
     "get_registered_learning_env_names",
     "get_trainer_class",
     "register_learning_env",
+    "stratified_rollout_value",
     "algo",
     "buffer",
     "models",

embodichain/learning/rl/algo/__init__.py

@@ -28,7 +28,12 @@
     coerce_optimizer_cfg,
 )
 
-from .apg import APG, APGCfg, segmented_discounted_return
+from .apg import (
+    APG,
+    APGCfg,
+    complete_discounted_return,
+    segmented_discounted_return,
+)
 from .base import BaseAlgorithm, RolloutKind
 from .common import compute_gae
 from .grpo import GRPO, GRPOCfg
@@ -95,6 +100,7 @@ def build_algo(
     "RolloutKind",
     "APGCfg",
     "APG",
+    "complete_discounted_return",
     "segmented_discounted_return",
     "PPOCfg",
     "PPO",