From 1a1f4fb0c04b9076de1fcc3d3af5fea668b5fcd8 Mon Sep 17 00:00:00 2001 From: yuecideng Date: Mon, 31 Aug 2026 13:24:17 +0800 Subject: [PATCH 1/3] feat(learning): add full-horizon waypoint APG support --- agent_context/MAP.yaml | 10 + .../topics/rl-learning/rl-learning.md | 57 +- .../embodichain.learning.rl.algo.rst | 1 + .../embodichain.learning.rl.models.rst | 14 +- .../embodichain/embodichain.learning.rl.rst | 22 + docs/source/api_reference/public_api.rst | 39 ++ embodichain/learning/rl/__init__.py | 13 + embodichain/learning/rl/algo/__init__.py | 8 +- embodichain/learning/rl/algo/apg.py | 200 ++++++- .../learning/rl/collector/differentiable.py | 98 +++- .../learning/rl/differentiable_trainer.py | 229 ++++++-- embodichain/learning/rl/env.py | 71 +++ embodichain/learning/rl/evaluation.py | 16 +- embodichain/learning/rl/gradients.py | 127 +++++ embodichain/learning/rl/models/__init__.py | 71 +++ .../rl/models/waypoint_transformer.py | 538 ++++++++++++++++++ embodichain/learning/rl/normalization.py | 180 ++++++ embodichain/learning/rl/train.py | 45 +- tests/learning/test_apg.py | 109 ++++ tests/learning/test_differentiable_trainer.py | 101 ++++ .../test_observation_normalization.py | 76 +++ tests/learning/test_waypoint_transformer.py | 145 +++++ 22 files changed, 2100 insertions(+), 70 deletions(-) create mode 100644 embodichain/learning/rl/gradients.py create mode 100644 embodichain/learning/rl/models/waypoint_transformer.py create mode 100644 embodichain/learning/rl/normalization.py create mode 100644 tests/learning/test_observation_normalization.py create mode 100644 tests/learning/test_waypoint_transformer.py diff --git a/agent_context/MAP.yaml b/agent_context/MAP.yaml index 1dc082947..89d7eb26a 100644 --- a/agent_context/MAP.yaml +++ b/agent_context/MAP.yaml @@ -490,6 +490,14 @@ topics: - reward - RolloutKind - DifferentiableTrainer + - ScheduledDifferentiableVecEnv + - DifferentiableRolloutSpec + - complete rollout + - full horizon + - gradient accumulation + - action adjoint + - observation normalization + - waypoint transformer - SyncCollector - learning_env - gym_config @@ -503,6 +511,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 diff --git a/agent_context/topics/rl-learning/rl-learning.md b/agent_context/topics/rl-learning/rl-learning.md index b1a19ee6b..b57daff79 100644 --- a/agent_context/topics/rl-learning/rl-learning.md +++ b/agent_context/topics/rl-learning/rl-learning.md @@ -39,7 +39,7 @@ CLI --config → trainer / policy / algorithm blocks → choose trainer.learning_env or trainer.gym_config → build environment - → build policy and optional MLP modules + → build policy and optional MLP or waypoint-Transformer modules → build algorithm config from the registry → route by algorithm.rollout_kind → train, evaluate, log, and checkpoint @@ -83,6 +83,11 @@ 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. @@ -114,9 +119,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 @@ -126,7 +146,9 @@ how many environment steps contribute to one optimizer update. | PPO, GRPO, APG implementations | `algo/ppo.py`, `algo/grpo.py`, `algo/apg.py` | | Standard rollout storage and views | `buffer/` | | Standard and differentiable collection | `collector/` | -| Policy interface, actor-critic, actor-only, MLP builder | `models/` | +| Policy interface, actor-critic, actor-only, MLP/waypoint Transformer builders | `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 +175,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 @@ -193,13 +217,24 @@ example is an experimental gradient reference, not a general simulator task. 3. Ensure its outputs satisfy every intended algorithm. 4. Provide graph-preserving sampling if used with differentiable rollouts. +The built-in `waypoint_transformer` module consumes the unified ordered +constraint layout: joint state, end-effector pose, absolute pose/joint targets, +active/valid and modality masks, last action, optional relative pose/joint +errors, and waypoint type. Its token sequence is +`[ACTION, STATE, ACTIVE_GOAL, WP_1..WP_K]`; attention is bidirectional so the +action can depend on future valid waypoints. + ### Add a Lightweight Environment 1. Implement `LearningVecEnv`, or `DifferentiableVecEnv` for APG. 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///agents/` when it is a bundled task. @@ -217,6 +252,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 +271,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 | diff --git a/docs/source/api_reference/embodichain/embodichain.learning.rl.algo.rst b/docs/source/api_reference/embodichain/embodichain.learning.rl.algo.rst index 0db256501..8410924e8 100644 --- a/docs/source/api_reference/embodichain/embodichain.learning.rl.algo.rst +++ b/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 diff --git a/docs/source/api_reference/embodichain/embodichain.learning.rl.models.rst b/docs/source/api_reference/embodichain/embodichain.learning.rl.models.rst index dba6979ed..3eb79fefe 100644 --- a/docs/source/api_reference/embodichain/embodichain.learning.rl.models.rst +++ b/docs/source/api_reference/embodichain/embodichain.learning.rl.models.rst @@ -8,8 +8,10 @@ Overview Policy-network registration and model construction APIs for RL agents. Policies implement the :class:`Policy` ABC; the built-in actor-critic variants are -:class:`ActorCritic` and :class:`ActorOnly`, both built on the :class:`MLP` -backbone. :func:`build_policy` constructs a policy from a config block, with +:class:`ActorCritic` and :class:`ActorOnly`. Backbones include :class:`MLP` and +the full-context :class:`WaypointTransformerActor` / +:class:`WaypointTransformerCritic` pair for ordered mixed-modality constraints. +:func:`build_policy` constructs a policy from a config block, with :class:`~embodichain.learning.rl.utils.config.AlgorithmCfg`-style registration through :func:`register_policy` / :func:`get_policy_class`. @@ -21,19 +23,25 @@ through :func:`register_policy` / :func:`get_policy_class`. ActorCritic ActorOnly MLP + WaypointTransformerEncoder + WaypointTransformerActor + WaypointTransformerCritic .. rubric:: Functions .. autosummary:: + build_model_from_cfg build_mlp_from_cfg build_policy get_policy_class get_registered_policy_names register_policy + parse_waypoint_observation + waypoint_observation_dim + waypoint_observation_normalize_mask .. automodule:: embodichain.learning.rl.models :members: :undoc-members: :show-inheritance: - \ No newline at end of file diff --git a/docs/source/api_reference/embodichain/embodichain.learning.rl.rst b/docs/source/api_reference/embodichain/embodichain.learning.rl.rst index bf1b5b01e..289e176d4 100644 --- a/docs/source/api_reference/embodichain/embodichain.learning.rl.rst +++ b/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 ------- diff --git a/docs/source/api_reference/public_api.rst b/docs/source/api_reference/public_api.rst index 840c9c24c..418c20f74 100644 --- a/docs/source/api_reference/public_api.rst +++ b/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,32 @@ embodichain.learning.rl.models.policy Policy +embodichain.learning.rl.models.waypoint_transformer +--------------------------------------------------- + +Unified tokenization and full-context models for ordered Cartesian, joint, and +mixed-modality waypoint policies. + +.. currentmodule:: embodichain.learning.rl.models.waypoint_transformer + +.. autosummary:: + + WaypointTransformerActor + WaypointTransformerCritic + WaypointTransformerEncoder + parse_waypoint_observation + waypoint_observation_dim + waypoint_observation_normalize_mask + +embodichain.learning.rl.normalization +------------------------------------- + +.. currentmodule:: embodichain.learning.rl.normalization + +.. autosummary:: + + RunningObservationNormalizer + embodichain.learning.rl.utils.optimizer --------------------------------------- diff --git a/embodichain/learning/rl/__init__.py b/embodichain/learning/rl/__init__.py index 9c97766b1..4808e9caa 100644 --- a/embodichain/learning/rl/__init__.py +++ b/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", diff --git a/embodichain/learning/rl/algo/__init__.py b/embodichain/learning/rl/algo/__init__.py index 7551d2b0c..cededb485 100644 --- a/embodichain/learning/rl/algo/__init__.py +++ b/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", diff --git a/embodichain/learning/rl/algo/apg.py b/embodichain/learning/rl/algo/apg.py index 279320593..9d40b09d0 100644 --- a/embodichain/learning/rl/algo/apg.py +++ b/embodichain/learning/rl/algo/apg.py @@ -23,12 +23,18 @@ import torch from embodichain.learning.rl.collector import DifferentiableRollout +from embodichain.learning.rl.gradients import BatchedGradientNormStats from embodichain.learning.rl.utils import AlgorithmCfg from embodichain.utils import configclass from .base import BaseAlgorithm, RolloutKind -__all__ = ["APG", "APGCfg", "segmented_discounted_return"] +__all__ = [ + "APG", + "APGCfg", + "complete_discounted_return", + "segmented_discounted_return", +] def segmented_discounted_return( @@ -51,6 +57,37 @@ def segmented_discounted_return( return returns +def complete_discounted_return( + rollout: DifferentiableRollout, + gamma: float, +) -> torch.Tensor: + """Compute per-environment returns up to the first terminal transition. + + Unlike segmented TBPTT returns, complete-rollout discounting never restarts + after ``done``. This excludes rewards emitted by any automatic reset during + the fixed full horizon. + + Args: + rollout: Complete graph-preserving trajectory. + gamma: Per-step discount factor. + + Returns: + One masked discounted return per environment. + + Raises: + ValueError: If the rollout is empty. + """ + if rollout.num_steps == 0: + raise ValueError("Cannot compute returns for an empty rollout.") + rewards = rollout.rewards + discounts = torch.as_tensor(gamma, device=rewards.device, dtype=rewards.dtype) ** ( + torch.arange(rollout.num_steps, device=rewards.device, dtype=rewards.dtype) + ) + return ( + rewards * rollout.alive_mask.to(rewards.dtype) * discounts.unsqueeze(-1) + ).sum(dim=0) + + @configclass class APGCfg(AlgorithmCfg): """Analytic policy-gradient config. @@ -60,6 +97,7 @@ class APGCfg(AlgorithmCfg): ent_coef: float = 0.0 skip_nonfinite_updates: bool = True + max_grad_norm_before_clip: float = 0.0 class APG(BaseAlgorithm[DifferentiableRollout]): @@ -68,6 +106,8 @@ class APG(BaseAlgorithm[DifferentiableRollout]): rollout_kind = RolloutKind.DIFFERENTIABLE def __init__(self, cfg: APGCfg, policy: torch.nn.Module) -> None: + if cfg.max_grad_norm_before_clip < 0.0: + raise ValueError("max_grad_norm_before_clip cannot be negative.") self.cfg = cfg self.policy = policy self.device = torch.device(cfg.device) @@ -79,6 +119,7 @@ def __init__(self, cfg: APGCfg, policy: torch.nn.Module) -> None: self._objective_total = 0.0 self._entropy_total = 0.0 self._num_accumulated_steps = 0 + self._action_gradient_stats: list[BatchedGradientNormStats] = [] def update(self, rollout: DifferentiableRollout) -> Dict[str, float]: """Apply one pathwise-gradient update from a rollout segment.""" @@ -101,6 +142,7 @@ def begin_update(self) -> None: self._objective_total = 0.0 self._entropy_total = 0.0 self._num_accumulated_steps = 0 + self._action_gradient_stats = [] def accumulate_segment(self, rollout: DifferentiableRollout) -> None: """Accumulate gradients from one TBPTT segment without stepping the optimizer.""" @@ -116,27 +158,105 @@ def accumulate_segment(self, rollout: DifferentiableRollout) -> None: objective = returns.mean() entropy = entropy_returns.mean() loss = -objective - self.cfg.ent_coef * entropy + self._accumulate_loss( + rollout, + loss=loss, + objective=objective, + entropy=entropy, + ) + + def accumulate_complete_rollout( + self, + rollout: DifferentiableRollout, + *, + objective_scale: float | torch.Tensor = 1.0, + accumulation_scale: float = 1.0, + ) -> None: + """Accumulate one independent full-horizon rollout. + + Args: + rollout: Complete graph-preserving trajectory. + objective_scale: Scalar or per-environment return multiplier. + accumulation_scale: Loss multiplier, normally reciprocal to the + number of independent rollout microbatches in one update. + + Raises: + RuntimeError: If no APG update is active. + ValueError: If the rollout, scaling, or accumulation factor is invalid. + """ + if not self._update_active: + raise RuntimeError("Call begin_update() before accumulating a rollout.") + if rollout.num_steps == 0: + raise ValueError("APG requires a non-empty differentiable rollout.") + if accumulation_scale <= 0.0: + raise ValueError("accumulation_scale must be positive.") + + returns = complete_discounted_return(rollout, self.cfg.gamma) + scale = torch.as_tensor( + objective_scale, + dtype=returns.dtype, + device=returns.device, + ).detach() + if scale.ndim > 1 or (scale.ndim == 1 and scale.shape != returns.shape): + raise ValueError( + "objective_scale must be scalar or have one value per environment." + ) + if not bool(torch.isfinite(scale).all()): + raise ValueError("objective_scale must contain only finite values.") + + entropy_returns = torch.zeros_like(returns) + discount = torch.ones_like(returns) + alive = torch.ones_like(returns, dtype=torch.bool) + for transition in rollout.transitions: + if "entropy" in transition.policy_output.keys(): + entropy_returns = entropy_returns + ( + discount + * transition.policy_output["entropy"] + * alive.to(discount.dtype) + ) + alive = alive & ~transition.done + discount = discount * self.cfg.gamma + + objective = (returns * scale).mean() + entropy = (entropy_returns * scale).mean() + loss = (-objective - self.cfg.ent_coef * entropy) * accumulation_scale + self._accumulate_loss( + rollout, + loss=loss, + objective=objective * accumulation_scale, + entropy=entropy * accumulation_scale, + ) + + def _accumulate_loss( + self, + rollout: DifferentiableRollout, + *, + loss: torch.Tensor, + objective: torch.Tensor, + entropy: torch.Tensor, + ) -> None: + """Validate and backpropagate one contribution to the active update.""" self._loss_total += float(loss.detach()) self._objective_total += float(objective.detach()) self._entropy_total += float(entropy.detach()) self._num_accumulated_steps += rollout.num_steps + if rollout.action_gradient_stats is not None: + self._action_gradient_stats.append(rollout.action_gradient_stats) - if not bool(torch.isfinite(loss)): - if not self.cfg.skip_nonfinite_updates: - raise FloatingPointError("APG produced a non-finite loss.") - self._update_valid = False - self.optimizer.zero_grad(set_to_none=True) - return - if not self._update_valid: - return + loss_is_finite = bool(torch.isfinite(loss)) + if not loss_is_finite and not self.cfg.skip_nonfinite_updates: + raise FloatingPointError("APG produced a non-finite loss.") + # Backward is also the lifecycle boundary for custom differentiable + # simulator steps. Invoke it even when this accumulation window is + # already invalid so every rollout can release its retained graph/tape. loss.backward() parameters = tuple(self.policy.parameters()) gradients_are_finite = all( parameter.grad is None or bool(torch.isfinite(parameter.grad).all()) for parameter in parameters ) - if not gradients_are_finite: + if not loss_is_finite or not gradients_are_finite or not self._update_valid: if not self.cfg.skip_nonfinite_updates: raise FloatingPointError("APG produced a non-finite policy gradient.") self._update_valid = False @@ -155,6 +275,7 @@ def finish_update(self) -> Dict[str, float]: metrics = self._accumulated_metrics( grad_norm=float("nan"), skipped_update=1.0, + skipped_excessive_gradient=0.0, ) self._update_active = False return metrics @@ -162,11 +283,25 @@ def finish_update(self) -> Dict[str, float]: parameters, self.cfg.max_grad_norm, ) + excessive_gradient = not bool(torch.isfinite(grad_norm)) or ( + self.cfg.max_grad_norm_before_clip > 0.0 + and float(grad_norm) > self.cfg.max_grad_norm_before_clip + ) + if excessive_gradient: + self.optimizer.zero_grad(set_to_none=True) + metrics = self._accumulated_metrics( + grad_norm=float(grad_norm.detach()), + skipped_update=1.0, + skipped_excessive_gradient=1.0, + ) + self._update_active = False + return metrics self.optimizer.step() self._step_scheduler() metrics = self._accumulated_metrics( grad_norm=float(grad_norm.detach()), skipped_update=0.0, + skipped_excessive_gradient=0.0, ) self._update_active = False return metrics @@ -174,6 +309,7 @@ def finish_update(self) -> Dict[str, float]: def cancel_update(self) -> None: self.optimizer.zero_grad(set_to_none=True) self._update_active = False + self._action_gradient_stats = [] def _discounted_terms( self, @@ -205,12 +341,54 @@ def _accumulated_metrics( *, grad_norm: float, skipped_update: float, + skipped_excessive_gradient: float, ) -> Dict[str, float]: - return { + metrics = { "loss": self._loss_total, "objective": self._objective_total, "entropy": self._entropy_total, "grad_norm": grad_norm, "skipped_update": skipped_update, + "skipped_excessive_gradient": skipped_excessive_gradient, "learning_rate": self.current_learning_rate(), } + if self._action_gradient_stats: + rows = sum(float(stats.rows) for stats in self._action_gradient_stats) + finite_rows = sum( + float(stats.finite_rows) for stats in self._action_gradient_stats + ) + metrics.update( + { + "action_adjoint_preclip_mean_norm": ( + sum( + float(stats.norm_sum) + for stats in self._action_gradient_stats + ) + / finite_rows + if finite_rows > 0.0 + else 0.0 + ), + "action_adjoint_preclip_max_norm": max( + float(stats.norm_max) for stats in self._action_gradient_stats + ), + "action_adjoint_clipped_fraction": ( + sum( + float(stats.clipped_rows) + for stats in self._action_gradient_stats + ) + / rows + if rows > 0.0 + else 0.0 + ), + "action_adjoint_nonfinite_fraction": ( + sum( + float(stats.nonfinite_rows) + for stats in self._action_gradient_stats + ) + / rows + if rows > 0.0 + else 0.0 + ), + } + ) + return metrics diff --git a/embodichain/learning/rl/collector/differentiable.py b/embodichain/learning/rl/collector/differentiable.py index 787c7c2d4..f8550633b 100644 --- a/embodichain/learning/rl/collector/differentiable.py +++ b/embodichain/learning/rl/collector/differentiable.py @@ -28,7 +28,12 @@ DifferentiableObservation, DifferentiableVecEnv, ) +from embodichain.learning.rl.gradients import ( + BatchedGradientNormStats, + clip_batched_gradient_norm, +) from embodichain.learning.rl.models import Policy +from embodichain.learning.rl.normalization import RunningObservationNormalizer from embodichain.learning.rl.utils import flatten_dict_observation __all__ = [ @@ -67,6 +72,7 @@ class DifferentiableRollout: initial_observation: torch.Tensor transitions: tuple[DifferentiableTransition, ...] + action_gradient_stats: BatchedGradientNormStats | None = None @property def num_steps(self) -> int: @@ -89,6 +95,39 @@ def rewards(self) -> torch.Tensor: ) return torch.stack([transition.reward for transition in self.transitions]) + @property + def observations(self) -> torch.Tensor: + """Stack the raw policy observations. + + Returns: + Tensor shaped ``[time, num_envs, features]``. + """ + if not self.transitions: + return self.initial_observation.new_empty( + (0,) + tuple(self.initial_observation.shape) + ) + return torch.stack([transition.observation for transition in self.transitions]) + + @property + def alive_mask(self) -> torch.Tensor: + """Return rows active before each step, stopping after the first done. + + Returns: + Boolean tensor shaped ``[time, num_envs]``. + """ + if not self.transitions: + return torch.empty( + (0, self.initial_observation.shape[0]), + dtype=torch.bool, + device=self.initial_observation.device, + ) + alive = torch.ones_like(self.transitions[0].done, dtype=torch.bool) + masks = [] + for transition in self.transitions: + masks.append(alive) + alive = alive & ~transition.done + return torch.stack(masks) + class DifferentiableCollector: """Collect graph-preserving rollouts without a preallocated buffer.""" @@ -98,11 +137,38 @@ def __init__( env: DifferentiableVecEnv, policy: Policy, device: torch.device, + *, + observation_normalizer: RunningObservationNormalizer | None = None, + clip_actions_to_space: bool = False, + action_adjoint_max_norm: float = 0.0, ) -> None: + if action_adjoint_max_norm < 0.0: + raise ValueError("action_adjoint_max_norm cannot be negative.") self.env = env self.policy = policy self.device = device + self.observation_normalizer = observation_normalizer + self.clip_actions_to_space = bool(clip_actions_to_space) + self.action_adjoint_max_norm = float(action_adjoint_max_norm) self._observation: DifferentiableObservation | None = None + self._action_lower: torch.Tensor | None = None + self._action_upper: torch.Tensor | None = None + if self.clip_actions_to_space: + action_space = self.env.single_action_space + if not hasattr(action_space, "low") or not hasattr(action_space, "high"): + raise TypeError( + "clip_actions_to_space requires an action space with low/high bounds." + ) + self._action_lower = torch.as_tensor( + action_space.low, + device=self.device, + dtype=torch.float32, + ) + self._action_upper = torch.as_tensor( + action_space.high, + device=self.device, + dtype=torch.float32, + ) def reset( self, *, seed: int | None = None @@ -138,11 +204,21 @@ def collect( initial_observation = self._flatten_observation(self._observation) transitions: list[DifferentiableTransition] = [] + gradient_stats = ( + BatchedGradientNormStats(self.device) + if self.action_adjoint_max_norm > 0.0 + else None + ) for _ in range(num_steps): observation = self._flatten_observation(self._observation) + policy_observation = ( + self.observation_normalizer.normalize(observation) + if self.observation_normalizer is not None + else observation + ) policy_input = TensorDict( - {"obs": observation}, + {"obs": policy_observation}, batch_size=[self.env.num_envs], device=self.device, ) @@ -150,8 +226,25 @@ def collect( policy_input, deterministic=deterministic, ) + action = policy_output["action"] + if self.clip_actions_to_space: + assert self._action_lower is not None + assert self._action_upper is not None + action = torch.maximum( + torch.minimum(action, self._action_upper), + self._action_lower, + ) + policy_output["action"] = action + if gradient_stats is not None and action.requires_grad: + action.register_hook( + lambda gradient, stats=gradient_stats: clip_batched_gradient_norm( + gradient, + self.action_adjoint_max_norm, + stats, + ) + ) next_observation, reward, terminated, truncated, info = self.env.step( - policy_output["action"] + action ) transition = DifferentiableTransition( observation=observation, @@ -171,6 +264,7 @@ def collect( return DifferentiableRollout( initial_observation=initial_observation, transitions=tuple(transitions), + action_gradient_stats=gradient_stats, ) def detach_state(self) -> torch.Tensor: diff --git a/embodichain/learning/rl/differentiable_trainer.py b/embodichain/learning/rl/differentiable_trainer.py index a69cee762..24e487f74 100644 --- a/embodichain/learning/rl/differentiable_trainer.py +++ b/embodichain/learning/rl/differentiable_trainer.py @@ -28,10 +28,18 @@ import wandb from embodichain.learning.rl.algo import APG -from embodichain.learning.rl.collector import DifferentiableCollector -from embodichain.learning.rl.env import DifferentiableVecEnv +from embodichain.learning.rl.collector import ( + DifferentiableCollector, + DifferentiableRollout, +) +from embodichain.learning.rl.env import ( + DifferentiableRolloutSpec, + DifferentiableVecEnv, + ScheduledDifferentiableVecEnv, +) from embodichain.learning.rl.evaluation import evaluate_episodes from embodichain.learning.rl.models import Policy +from embodichain.learning.rl.normalization import RunningObservationNormalizer from embodichain.learning.rl.utils import LRSchedulerCfg, build_lr_scheduler from embodichain.utils import configclass @@ -45,11 +53,17 @@ @configclass class DifferentiableTrainerCfg: - """Configuration for graph-preserving segmented training.""" + """Configuration for segmented or complete graph-preserving training.""" segment_length: int = 16 update_horizon: int | None = None + rollout_mode: str = "segmented" + gradient_accumulation_steps: int = 1 deterministic_actions: bool = False + clip_actions_to_space: bool = False + action_adjoint_max_norm: float = 0.0 + normalize_observations: bool = False + rollout_seed: int | None = None checkpoint_dir: str = "outputs/checkpoints" experiment_name: str = "apg" save_frequency_updates: int = 0 @@ -62,7 +76,7 @@ class DifferentiableTrainerCfg: class DifferentiableTrainer: - """Coordinate APG updates and truncated-backpropagation boundaries.""" + """Coordinate APG updates, full rollouts, and optional TBPTT boundaries.""" def __init__( self, @@ -72,16 +86,25 @@ def __init__( algorithm: APG, writer: SummaryWriter | None = None, eval_env: DifferentiableVecEnv | None = None, + observation_normalizer: RunningObservationNormalizer | None = None, ) -> None: if cfg.segment_length <= 0: raise ValueError("segment_length must be positive.") update_horizon = ( cfg.segment_length if cfg.update_horizon is None else cfg.update_horizon ) - if update_horizon < cfg.segment_length: + if cfg.rollout_mode not in {"segmented", "complete"}: + raise ValueError("rollout_mode must be 'segmented' or 'complete'.") + if cfg.rollout_mode == "segmented" and update_horizon < cfg.segment_length: raise ValueError("update_horizon must be at least segment_length.") - if update_horizon % cfg.segment_length != 0: + if cfg.rollout_mode == "segmented" and update_horizon % cfg.segment_length != 0: raise ValueError("update_horizon must be divisible by segment_length.") + if update_horizon <= 0: + raise ValueError("update_horizon must be positive.") + if cfg.gradient_accumulation_steps <= 0: + raise ValueError("gradient_accumulation_steps must be positive.") + if cfg.action_adjoint_max_norm < 0.0: + raise ValueError("action_adjoint_max_norm cannot be negative.") if cfg.save_frequency_updates < 0: raise ValueError("save_frequency_updates cannot be negative.") if cfg.eval_frequency_steps < 0: @@ -100,10 +123,22 @@ def __init__( self.algorithm = algorithm self.writer = writer self.eval_env = eval_env + if observation_normalizer is None and cfg.normalize_observations: + observation_dim = int(env.single_observation_space.shape[-1]) + normalize_mask = getattr(env, "observation_normalize_mask", None) + observation_normalizer = RunningObservationNormalizer( + observation_dim, + algorithm.device, + normalize_mask=normalize_mask, + ) + self.observation_normalizer = observation_normalizer self.collector = DifferentiableCollector( env=env, policy=policy, device=algorithm.device, + observation_normalizer=observation_normalizer, + clip_actions_to_space=cfg.clip_actions_to_space, + action_adjoint_max_norm=cfg.action_adjoint_max_norm, ) self.global_step = 0 self.num_updates = 0 @@ -125,43 +160,72 @@ def __init__( self._next_eval_step = ( cfg.eval_frequency_steps if cfg.eval_frequency_steps > 0 else None ) + self._rollout_index = 0 + + def train( + self, + total_timesteps: int | None = None, + *, + total_updates: int | None = None, + ) -> dict[str, Any]: + """Train to a vector-transition or optimizer-update budget. + + ``total_updates`` is the stable budget for scheduled variable-horizon + training. Exactly one budget may be supplied. - def train(self, total_timesteps: int) -> dict[str, Any]: - """Train until at least ``total_timesteps`` vector transitions exist.""" - if total_timesteps < 0: + Args: + total_timesteps: Optional absolute vector-transition budget. + total_updates: Optional absolute optimizer-update budget. + + Returns: + Current training counters, metrics, histories, and checkpoint paths. + + Raises: + ValueError: If neither/both budgets are supplied or one is negative. + """ + if total_timesteps is None and total_updates is None: + raise ValueError("Provide total_timesteps or total_updates.") + if total_timesteps is not None and total_updates is not None: + raise ValueError("Provide only one training budget.") + if total_timesteps is not None and total_timesteps < 0: raise ValueError("total_timesteps cannot be negative.") + if total_updates is not None and total_updates < 0: + raise ValueError("total_updates cannot be negative.") - steps_per_update = self.update_horizon * self.env.num_envs - if total_timesteps > 0 and steps_per_update > 0: - total_updates = math.ceil(total_timesteps / steps_per_update) - self.algorithm.bind_schedule(total_updates=total_updates) + if total_updates is not None: + if total_updates > 0: + self.algorithm.bind_schedule(total_updates=total_updates) + else: + assert total_timesteps is not None + steps_per_update = ( + self.update_horizon + * self.env.num_envs + * ( + self.cfg.gradient_accumulation_steps + if self.cfg.rollout_mode == "complete" + else 1 + ) + ) + if total_timesteps > 0 and steps_per_update > 0: + estimated_updates = math.ceil(total_timesteps / steps_per_update) + self.algorithm.bind_schedule(total_updates=estimated_updates) self.policy.train() - while self.global_step < total_timesteps: - remaining_vector_steps = math.ceil( - (total_timesteps - self.global_step) / self.env.num_envs - ) - update_steps = min(self.update_horizon, remaining_vector_steps) - collected_steps = 0 - self.algorithm.begin_update() - try: - while collected_steps < update_steps: - segment_steps = min( - self.cfg.segment_length, - update_steps - collected_steps, - ) - rollout = self.collector.collect( - segment_steps, - deterministic=self.cfg.deterministic_actions, - on_step_callback=self._on_step, + while ( + self.num_updates < total_updates + if total_updates is not None + else self.global_step < total_timesteps + ): + if self.cfg.rollout_mode == "complete": + collected_steps, metrics = self._complete_rollout_update() + else: + update_steps = self.update_horizon + if total_timesteps is not None: + remaining_vector_steps = math.ceil( + (total_timesteps - self.global_step) / self.env.num_envs ) - self.algorithm.accumulate_segment(rollout) - self.collector.detach_state() - collected_steps += rollout.num_steps - metrics = self.algorithm.finish_update() - except Exception: - self.algorithm.cancel_update() - raise + update_steps = min(self.update_horizon, remaining_vector_steps) + collected_steps, metrics = self._segmented_update(update_steps) self.global_step += collected_steps * self.env.num_envs self.num_updates += 1 @@ -202,6 +266,80 @@ def train(self, total_timesteps: int) -> dict[str, Any]: return self.get_summary() + def _segmented_update(self, update_steps: int) -> tuple[int, dict[str, float]]: + """Run one backward-compatible TBPTT optimizer update.""" + collected_steps = 0 + self.algorithm.begin_update() + try: + while collected_steps < update_steps: + segment_steps = min( + self.cfg.segment_length, + update_steps - collected_steps, + ) + rollout = self.collector.collect( + segment_steps, + deterministic=self.cfg.deterministic_actions, + on_step_callback=self._on_step, + ) + self.algorithm.accumulate_segment(rollout) + self._update_observation_normalizer(rollout) + self.collector.detach_state() + collected_steps += rollout.num_steps + return collected_steps, self.algorithm.finish_update() + except Exception: + self.algorithm.cancel_update() + raise + + def _complete_rollout_update(self) -> tuple[int, dict[str, float]]: + """Accumulate independent full trajectories and apply one APG step.""" + accumulation_steps = self.cfg.gradient_accumulation_steps + collected_steps = 0 + metadata: dict[str, list[float]] = {} + self.algorithm.begin_update() + try: + for _ in range(accumulation_steps): + spec = self._prepare_complete_rollout() + seed = self.cfg.rollout_seed if self._rollout_index == 0 else None + self.collector.reset(seed=seed) + rollout = self.collector.collect( + spec.num_steps, + deterministic=self.cfg.deterministic_actions, + on_step_callback=self._on_step, + ) + self.algorithm.accumulate_complete_rollout( + rollout, + objective_scale=spec.objective_scale, + accumulation_scale=1.0 / accumulation_steps, + ) + self._update_observation_normalizer(rollout) + self.collector.detach_state() + collected_steps += rollout.num_steps + for key, value in spec.metadata.items(): + metadata.setdefault(str(key), []).append(float(value)) + self._rollout_index += 1 + metrics = self.algorithm.finish_update() + except Exception: + self.algorithm.cancel_update() + raise + for key, values in metadata.items(): + metrics[f"rollout_{key}_mean"] = sum(values) / len(values) + return collected_steps, metrics + + def _prepare_complete_rollout(self) -> DifferentiableRolloutSpec: + if isinstance(self.env, ScheduledDifferentiableVecEnv): + return self.env.prepare_differentiable_rollout(self._rollout_index) + return DifferentiableRolloutSpec(num_steps=self.update_horizon) + + def _update_observation_normalizer( + self, + rollout: DifferentiableRollout, + ) -> None: + if self.observation_normalizer is None: + return + observations = rollout.observations + alive_observations = observations[rollout.alive_mask] + self.observation_normalizer.update(alive_observations.detach()) + def save_checkpoint(self, path: str | Path | None = None) -> str: """Save policy, optimizer, and trainer counters.""" if path is None: @@ -218,7 +356,10 @@ def save_checkpoint(self, path: str | Path | None = None) -> str: "policy": self.policy.state_dict(), "optimizer": self.algorithm.optimizer.state_dict(), "best_eval_value": self.best_eval_value, + "rollout_index": self._rollout_index, } + if self.observation_normalizer is not None: + payload["observation_normalizer"] = self.observation_normalizer.state_dict() if self.algorithm.lr_scheduler is not None: payload["lr_scheduler"] = self.algorithm.lr_scheduler.state_dict() payload["lr_scheduler_cfg"] = { @@ -259,6 +400,15 @@ def load_checkpoint(self, path: str | Path) -> None: self.global_step = int(checkpoint["global_step"]) self.num_updates = int(checkpoint["num_updates"]) self.best_eval_value = checkpoint.get("best_eval_value") + self._rollout_index = int(checkpoint.get("rollout_index", self.num_updates)) + normalizer_state = checkpoint.get("observation_normalizer") + if normalizer_state is not None: + if self.observation_normalizer is None: + raise ValueError( + "Checkpoint contains observation normalization state, but the " + "trainer has normalization disabled." + ) + self.observation_normalizer.load_state_dict(normalizer_state) self.latest_checkpoint_path = str(path) def get_summary(self) -> dict[str, Any]: @@ -305,6 +455,11 @@ def _evaluate(self) -> dict[str, float]: num_episodes=self.cfg.num_eval_episodes, device=self.algorithm.device, seed=self.cfg.eval_seed, + observation_transform=( + self.observation_normalizer.normalize + if self.observation_normalizer is not None + else None + ), ) entry = {"global_step": float(self.global_step), **metrics} self.eval_history.append(entry) diff --git a/embodichain/learning/rl/env.py b/embodichain/learning/rl/env.py index 718c3df7f..60bdeaff7 100644 --- a/embodichain/learning/rl/env.py +++ b/embodichain/learning/rl/env.py @@ -19,6 +19,7 @@ from __future__ import annotations from collections.abc import Callable +from dataclasses import dataclass, field from typing import Any, Mapping, Protocol, TypeAlias, runtime_checkable import torch @@ -27,11 +28,14 @@ __all__ = [ "DifferentiableObservation", + "DifferentiableRolloutSpec", "DifferentiableVecEnv", "LearningVecEnv", "build_learning_env", "get_registered_learning_env_names", "register_learning_env", + "ScheduledDifferentiableVecEnv", + "stratified_rollout_value", ] DifferentiableObservation: TypeAlias = torch.Tensor | TensorDict @@ -40,6 +44,48 @@ _LEARNING_ENV_REGISTRY: dict[str, LearningEnvFactory] = {} +@dataclass(frozen=True) +class DifferentiableRolloutSpec: + """Describe one independent complete rollout used for a gradient microbatch. + + Args: + num_steps: Full rollout horizon. The trainer must not detach or truncate it. + objective_scale: Per-environment or scalar multiplier applied to returns. + metadata: Scalar labels recorded with training metrics. + """ + + num_steps: int + objective_scale: float | torch.Tensor = 1.0 + metadata: Mapping[str, float] = field(default_factory=dict) + + def __post_init__(self) -> None: + if isinstance(self.num_steps, bool) or int(self.num_steps) != self.num_steps: + raise TypeError("num_steps must be a positive integer.") + if self.num_steps <= 0: + raise ValueError("num_steps must be a positive integer.") + + +def stratified_rollout_value(index: int, minimum: int, maximum: int) -> int: + """Cycle uniformly through an integer range and rotate each cycle's order. + + Args: + index: Zero-based rollout index. + minimum: Inclusive minimum scheduled value. + maximum: Inclusive maximum scheduled value. + + Returns: + Scheduled integer for ``index``. + + Raises: + ValueError: If the inclusive range is empty. + """ + count = int(maximum) - int(minimum) + 1 + if count < 1: + raise ValueError("minimum must be less than or equal to maximum.") + cycle, position = divmod(int(index), count) + return int(minimum) + ((position + cycle) % count) + + @runtime_checkable class LearningVecEnv(Protocol): """Structural interface shared by lightweight vector environments.""" @@ -88,6 +134,31 @@ def detach_state(self) -> DifferentiableObservation: ... +@runtime_checkable +class ScheduledDifferentiableVecEnv(DifferentiableVecEnv, Protocol): + """Differentiable env that schedules variable independent rollouts. + + The trainer calls :meth:`prepare_differentiable_rollout` immediately before + resetting the environment. Implementations may select task difficulty for + the next reset, such as an ordered-waypoint count, and must return the full + horizon and objective scaling for that selection. + """ + + def prepare_differentiable_rollout( + self, + rollout_index: int, + ) -> DifferentiableRolloutSpec: + """Configure the next reset and return its complete-rollout contract. + + Args: + rollout_index: Zero-based independent-rollout index. + + Returns: + Full horizon, objective scaling, and optional metric metadata. + """ + ... + + def _is_nested_package_shadow(existing: Any, candidate: Any) -> bool: """Return True when ``candidate`` is an editable-install nested duplicate. diff --git a/embodichain/learning/rl/evaluation.py b/embodichain/learning/rl/evaluation.py index 974a2ab7c..9e5277ea9 100644 --- a/embodichain/learning/rl/evaluation.py +++ b/embodichain/learning/rl/evaluation.py @@ -87,8 +87,20 @@ def evaluate_episodes( device: torch.device | str, seed: int | None = None, on_step: Callable[[dict[str, Any]], None] | None = None, + observation_transform: Callable[[torch.Tensor], torch.Tensor] | None = None, ) -> dict[str, float]: - """Evaluate exactly ``num_episodes`` completed asynchronous episodes.""" + """Evaluate exactly ``num_episodes`` completed asynchronous episodes. + + Args: + policy: Policy evaluated with deterministic actions. + env: Vector environment with automatic row reset. + num_episodes: Exact number of completed episodes to collect. + device: Policy device. + seed: Optional environment reset seed. + on_step: Optional callback for raw step info. + observation_transform: Optional frozen preprocessing transform, such as + a running observation normalizer. + """ if num_episodes <= 0: raise ValueError("num_episodes must be positive.") device = torch.device(device) @@ -107,6 +119,8 @@ def evaluate_episodes( observation, _ = env.reset(seed=seed) while len(returns) < num_episodes: flat_observation = _flat_observation(observation, device) + if observation_transform is not None: + flat_observation = observation_transform(flat_observation) policy_input = TensorDict( {"obs": flat_observation}, batch_size=[num_envs], diff --git a/embodichain/learning/rl/gradients.py b/embodichain/learning/rl/gradients.py new file mode 100644 index 000000000..99828fc5d --- /dev/null +++ b/embodichain/learning/rl/gradients.py @@ -0,0 +1,127 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Gradient stabilization primitives for differentiable rollouts.""" + +from __future__ import annotations + +import torch + +__all__ = ["BatchedGradientNormStats", "clip_batched_gradient_norm"] + + +class BatchedGradientNormStats: + """Accumulate row-wise adjoint norm and clipping statistics on-device. + + Args: + device: Device on which hook-side counters are accumulated. + """ + + def __init__(self, device: torch.device | str) -> None: + self.device = torch.device(device) + self.norm_sum = torch.zeros((), device=self.device) + self.norm_max = torch.zeros((), device=self.device) + self.finite_rows = torch.zeros((), device=self.device) + self.rows = torch.zeros((), device=self.device) + self.clipped_rows = torch.zeros((), device=self.device) + self.nonfinite_rows = torch.zeros((), device=self.device) + + def metrics(self) -> dict[str, float]: + """Return statistics after backward has invoked registered hooks. + + Returns: + Mean and maximum pre-clip norm plus clipped/non-finite fractions. + """ + rows = float(self.rows) + finite_rows = float(self.finite_rows) + return { + "action_adjoint_preclip_mean_norm": ( + float(self.norm_sum) / finite_rows if finite_rows > 0.0 else 0.0 + ), + "action_adjoint_preclip_max_norm": float(self.norm_max), + "action_adjoint_clipped_fraction": ( + float(self.clipped_rows) / rows if rows > 0.0 else 0.0 + ), + "action_adjoint_nonfinite_fraction": ( + float(self.nonfinite_rows) / rows if rows > 0.0 else 0.0 + ), + } + + +def clip_batched_gradient_norm( + gradient: torch.Tensor, + max_norm: float, + stats: BatchedGradientNormStats | None = None, +) -> torch.Tensor: + """Clip each batch row without shortening the differentiable time horizon. + + Norms are computed with max-absolute-value scaling to avoid overflow in + float32. A non-finite row is replaced with zeros while finite rows remain + independent from one another. + + Args: + gradient: Tensor whose first dimension identifies independent rows. + max_norm: Maximum L2 norm per row. Zero disables clipping. + stats: Optional on-device accumulator populated before clipping. + + Returns: + A finite tensor with the same shape, dtype, and device as ``gradient``. + + Raises: + ValueError: If ``max_norm`` is negative or ``gradient`` is not batched. + """ + if max_norm < 0.0: + raise ValueError("max_norm cannot be negative.") + if max_norm == 0.0: + return gradient + if gradient.ndim < 2: + raise ValueError("gradient must have a leading batch dimension.") + + flat = gradient.flatten(start_dim=1) + finite_rows = torch.isfinite(flat).all(dim=1, keepdim=True) + finite_values = torch.where(finite_rows, flat, torch.zeros_like(flat)) + max_abs = finite_values.abs().amax(dim=1, keepdim=True) + safe_max_abs = max_abs.clamp_min(1.0e-12) + scaled_norm = (finite_values / safe_max_abs).norm(dim=1, keepdim=True) + raw_norm = max_abs * scaled_norm + + if stats is not None: + with torch.no_grad(): + detached_norm = raw_norm.detach().flatten() + finite = finite_rows.detach().flatten() & torch.isfinite(detached_norm) + finite_norm = torch.where( + finite, + detached_norm, + torch.zeros_like(detached_norm), + ) + stats.norm_sum.add_(finite_norm.sum()) + stats.finite_rows.add_(finite.sum()) + stats.rows.add_(detached_norm.numel()) + stats.clipped_rows.add_((finite & (detached_norm > max_norm)).sum()) + stats.nonfinite_rows.add_((~finite).sum()) + stats.norm_max.copy_(torch.maximum(stats.norm_max, finite_norm.max())) + + scale = ((float(max_norm) / safe_max_abs) / scaled_norm.clamp_min(1.0)).clamp( + max=1.0 + ) + scale = torch.where(finite_rows, scale, torch.zeros_like(scale)) + broadcast_shape = (-1,) + (1,) * (gradient.ndim - 1) + safe_gradient = torch.where( + finite_rows.view(broadcast_shape), + gradient, + torch.zeros_like(gradient), + ) + return safe_gradient * scale.view(broadcast_shape) diff --git a/embodichain/learning/rl/models/__init__.py b/embodichain/learning/rl/models/__init__.py index fa6d46ae6..f7eae3d5c 100644 --- a/embodichain/learning/rl/models/__init__.py +++ b/embodichain/learning/rl/models/__init__.py @@ -28,6 +28,14 @@ from .actor_only import ActorOnly from .policy import Policy from .mlp import MLP +from .waypoint_transformer import ( + WaypointTransformerActor, + WaypointTransformerCritic, + WaypointTransformerEncoder, + parse_waypoint_observation, + waypoint_observation_dim, + waypoint_observation_normalize_mask, +) # In-module policy registry _POLICY_REGISTRY: Dict[str, Type[Policy]] = {} @@ -152,6 +160,62 @@ def build_mlp_from_cfg(module_cfg: Dict, in_dim: int, out_dim: int) -> MLP: return model +def build_model_from_cfg( + module_cfg: Dict, + in_dim: int, + out_dim: int, + *, + role: str, +) -> torch.nn.Module: + """Build an actor or critic module from a JSON-like config. + + Args: + module_cfg: Module configuration with ``type`` and ``network_cfg``. + in_dim: Flat observation dimension. + out_dim: Requested output dimension. + role: ``"actor"`` or ``"critic"``. + + Returns: + Constructed PyTorch module. + + Raises: + ValueError: If the module type or role is unsupported. + """ + module_type = module_cfg.get("type", "").lower() + if module_type == "mlp": + return build_mlp_from_cfg(module_cfg, in_dim, out_dim) + if module_type != "waypoint_transformer": + raise ValueError( + "Supported actor/critic module types are 'mlp' and " + f"'waypoint_transformer', got {module_type!r}." + ) + if role not in {"actor", "critic"}: + raise ValueError("role must be 'actor' or 'critic'.") + + network_cfg = dict(module_cfg.get("network_cfg", {})) + kwargs = { + "observation_dim": in_dim, + "num_waypoints": int(network_cfg["num_waypoints"]), + "joint_dim": int(network_cfg.get("joint_dim", 7)), + "use_relative_observations": bool( + network_cfg.get("use_relative_observations", True) + ), + "hidden_dim": int(network_cfg.get("hidden_dim", 256)), + "num_attention_heads": int(network_cfg.get("num_attention_heads", 4)), + "num_layers": int(network_cfg.get("num_layers", 2)), + "feedforward_dim": ( + int(network_cfg["feedforward_dim"]) + if int(network_cfg.get("feedforward_dim", 0)) > 0 + else None + ), + } + if role == "actor": + return WaypointTransformerActor(action_dim=out_dim, **kwargs) + if out_dim != 1: + raise ValueError("A waypoint Transformer critic must have out_dim=1.") + return WaypointTransformerCritic(**kwargs) + + # default registrations register_policy("actor_critic", ActorCritic) register_policy("actor_only", ActorOnly) @@ -162,8 +226,15 @@ def build_mlp_from_cfg(module_cfg: Dict, in_dim: int, out_dim: int) -> MLP: "register_policy", "get_registered_policy_names", "build_policy", + "build_model_from_cfg", "build_mlp_from_cfg", "get_policy_class", "Policy", "MLP", + "WaypointTransformerActor", + "WaypointTransformerCritic", + "WaypointTransformerEncoder", + "parse_waypoint_observation", + "waypoint_observation_dim", + "waypoint_observation_normalize_mask", ] diff --git a/embodichain/learning/rl/models/waypoint_transformer.py b/embodichain/learning/rl/models/waypoint_transformer.py new file mode 100644 index 000000000..806456eb9 --- /dev/null +++ b/embodichain/learning/rl/models/waypoint_transformer.py @@ -0,0 +1,538 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Transformer models for ordered Cartesian and joint waypoint constraints.""" + +from __future__ import annotations + +import math + +import torch +import torch.nn as nn + +__all__ = [ + "WaypointTransformerActor", + "WaypointTransformerCritic", + "WaypointTransformerEncoder", + "parse_waypoint_observation", + "waypoint_observation_dim", + "waypoint_observation_normalize_mask", +] + +_TOKEN_ACTION = 0 +_TOKEN_STATE = 1 +_TOKEN_ACTIVE_GOAL = 2 +_TOKEN_WAYPOINT = 3 +_NUM_TOKEN_TYPES = 4 +_NUM_WAYPOINT_TYPES = 3 + + +def _layer_init( + layer: nn.Linear, + std: float = math.sqrt(2.0), + bias_const: float = 0.0, +) -> nn.Linear: + torch.nn.init.orthogonal_(layer.weight, std) + torch.nn.init.constant_(layer.bias, bias_const) + return layer + + +def waypoint_observation_dim( + num_waypoints: int, + use_relative_observations: bool, + *, + joint_dim: int = 7, +) -> int: + """Return the flat dimension of the unified waypoint observation layout. + + Args: + num_waypoints: Maximum ordered waypoint count. + use_relative_observations: Whether relative pose/joint fields are present. + joint_dim: Controlled joint dimension. + + Returns: + Required flat observation dimension. + + Raises: + ValueError: If a dimension is not positive. + """ + if num_waypoints <= 0: + raise ValueError("num_waypoints must be positive.") + if joint_dim <= 0: + raise ValueError("joint_dim must be positive.") + num_waypoints = int(num_waypoints) + joint_dim = int(joint_dim) + dimension = ( + joint_dim + + 7 + + num_waypoints * (3 + 4 + joint_dim) + + 5 * num_waypoints + + joint_dim + + num_waypoints + ) + if use_relative_observations: + dimension += 7 + num_waypoints * (3 + 4 + joint_dim) + return dimension + + +def waypoint_observation_normalize_mask( + num_waypoints: int, + use_relative_observations: bool, + *, + joint_dim: int = 7, + device: torch.device | str | None = None, +) -> torch.Tensor: + """Build a mask that preserves waypoint semantics during normalization. + + Args: + num_waypoints: Maximum ordered waypoint count. + use_relative_observations: Whether relative pose/joint fields are present. + joint_dim: Controlled joint dimension. + device: Optional output device. + + Returns: + Boolean mask where ``True`` selects continuous normalized fields. + """ + num_waypoints = int(num_waypoints) + joint_dim = int(joint_dim) + dimension = waypoint_observation_dim( + num_waypoints, + use_relative_observations, + joint_dim=joint_dim, + ) + mask = torch.ones(dimension, dtype=torch.bool, device=device) + cursor = joint_dim + 7 + num_waypoints * (3 + 4 + joint_dim) + mask[cursor : cursor + 5 * num_waypoints] = False + cursor += 5 * num_waypoints + joint_dim + if use_relative_observations: + cursor += 7 + num_waypoints * (3 + 4) + mask[cursor : cursor + num_waypoints * joint_dim] = False + cursor += num_waypoints * joint_dim + mask[cursor : cursor + num_waypoints] = False + return mask + + +def parse_waypoint_observation( + observation: torch.Tensor, + num_waypoints: int, + use_relative_observations: bool, + *, + joint_dim: int = 7, +) -> dict[str, torch.Tensor | None]: + """Slice a flat unified waypoint observation into semantic fields. + + The layout is ``joint, eef, waypoint pose, waypoint joint, active, valid, + position/rotation/joint masks, last action, optional relative fields, type``. + + Args: + observation: Flat tensor shaped ``[batch, features]``. + num_waypoints: Maximum ordered waypoint count. + use_relative_observations: Whether relative fields are present. + joint_dim: Controlled joint dimension. + + Returns: + Mapping from semantic field names to tensor views. + + Raises: + ValueError: If the observation rank or feature dimension is invalid. + """ + if observation.ndim != 2: + raise ValueError("observation must have shape [batch, features].") + expected_dim = waypoint_observation_dim( + num_waypoints, + use_relative_observations, + joint_dim=joint_dim, + ) + if observation.shape[1] != expected_dim: + raise ValueError( + f"Expected waypoint observation dimension {expected_dim}, " + f"got {observation.shape[1]}." + ) + + n = int(num_waypoints) + d = int(joint_dim) + cursor = 0 + joint = observation[:, cursor : cursor + d] + cursor += d + end_effector = observation[:, cursor : cursor + 7] + cursor += 7 + waypoint_position = observation[:, cursor : cursor + 3 * n].reshape(-1, n, 3) + cursor += 3 * n + waypoint_quaternion = observation[:, cursor : cursor + 4 * n].reshape(-1, n, 4) + cursor += 4 * n + waypoint_joint = observation[:, cursor : cursor + d * n].reshape(-1, n, d) + cursor += d * n + active_onehot = observation[:, cursor : cursor + n] + cursor += n + valid_mask = observation[:, cursor : cursor + n] + cursor += n + position_mask = observation[:, cursor : cursor + n] + cursor += n + rotation_mask = observation[:, cursor : cursor + n] + cursor += n + joint_mask = observation[:, cursor : cursor + n] + cursor += n + last_action = observation[:, cursor : cursor + d] + cursor += d + + active_relative_pose = None + waypoint_relative_position = None + waypoint_relative_quaternion = None + waypoint_joint_error = None + if use_relative_observations: + active_relative_pose = observation[:, cursor : cursor + 7] + cursor += 7 + waypoint_relative_position = observation[:, cursor : cursor + 3 * n].reshape( + -1, n, 3 + ) + cursor += 3 * n + waypoint_relative_quaternion = observation[:, cursor : cursor + 4 * n].reshape( + -1, n, 4 + ) + cursor += 4 * n + waypoint_joint_error = observation[:, cursor : cursor + d * n].reshape(-1, n, d) + cursor += d * n + waypoint_type = observation[:, cursor : cursor + n] + + active_id = active_onehot.argmax(dim=-1).long() + valid_count = valid_mask.sum(dim=-1).long() + active_id = torch.clamp(torch.minimum(active_id, valid_count - 1), min=0) + return { + "joint": joint, + "end_effector": end_effector, + "last_action": last_action, + "active_relative_pose": active_relative_pose, + "waypoint_relative_position": waypoint_relative_position, + "waypoint_relative_quaternion": waypoint_relative_quaternion, + "waypoint_joint_error": waypoint_joint_error, + "active_onehot": active_onehot, + "valid_mask": valid_mask, + "position_mask": position_mask, + "rotation_mask": rotation_mask, + "joint_mask": joint_mask, + "active_id": active_id, + "waypoint_position": waypoint_position, + "waypoint_quaternion": waypoint_quaternion, + "waypoint_joint": waypoint_joint, + "waypoint_type": waypoint_type, + } + + +def _token_head( + hidden_dim: int, + output_dim: int, + std: float, + *, + input_dim: int | None = None, +) -> nn.Sequential: + input_dim = hidden_dim if input_dim is None else int(input_dim) + return nn.Sequential( + nn.LayerNorm(input_dim), + _layer_init(nn.Linear(input_dim, hidden_dim)), + nn.Tanh(), + _layer_init(nn.Linear(hidden_dim, output_dim), std=std), + ) + + +class WaypointTransformerEncoder(nn.Module): + """Encode state, active-goal, and all waypoint tokens bidirectionally. + + Args: + observation_dim: Flat unified observation dimension. + num_waypoints: Maximum ordered waypoint count. + joint_dim: Controlled joint dimension. + use_relative_observations: Whether relative fields are present. + hidden_dim: Transformer embedding dimension. + num_attention_heads: Attention-head count. + num_layers: Encoder-layer count. + feedforward_dim: Optional feed-forward dimension; defaults to four + times ``hidden_dim``. + """ + + def __init__( + self, + observation_dim: int, + num_waypoints: int, + *, + joint_dim: int = 7, + use_relative_observations: bool = True, + hidden_dim: int = 256, + num_attention_heads: int = 4, + num_layers: int = 2, + feedforward_dim: int | None = None, + ) -> None: + super().__init__() + self.num_waypoints = int(num_waypoints) + self.joint_dim = int(joint_dim) + self.use_relative_observations = bool(use_relative_observations) + expected_dim = waypoint_observation_dim( + self.num_waypoints, + self.use_relative_observations, + joint_dim=self.joint_dim, + ) + if int(observation_dim) != expected_dim: + raise ValueError( + f"WaypointTransformerEncoder expected observation_dim " + f"{expected_dim}, got {observation_dim}." + ) + if hidden_dim % num_attention_heads != 0: + raise ValueError("hidden_dim must be divisible by num_attention_heads.") + + self.state_dim = self.joint_dim + 7 + self.joint_dim + self.active_goal_dim = 7 + self.joint_dim + 3 + 1 + if self.use_relative_observations: + self.active_goal_dim += 7 + self.joint_dim + self.waypoint_token_dim = 3 + 4 + self.joint_dim + 3 + 1 + 1 + if self.use_relative_observations: + self.waypoint_token_dim += 7 + self.joint_dim + + self.action_token = nn.Parameter(torch.empty(1, 1, hidden_dim)) + nn.init.normal_(self.action_token, std=0.02) + self.token_type_embedding = nn.Embedding(_NUM_TOKEN_TYPES, hidden_dim) + nn.init.normal_(self.token_type_embedding.weight, std=0.02) + self.state_proj = _layer_init(nn.Linear(self.state_dim, hidden_dim)) + self.active_goal_proj = _layer_init(nn.Linear(self.active_goal_dim, hidden_dim)) + self.waypoint_proj = _layer_init(nn.Linear(self.waypoint_token_dim, hidden_dim)) + self.waypoint_index_embedding = nn.Parameter( + torch.empty(1, self.num_waypoints, hidden_dim) + ) + nn.init.normal_(self.waypoint_index_embedding, std=0.02) + self.waypoint_modality_embedding = nn.Embedding( + _NUM_WAYPOINT_TYPES, + hidden_dim, + ) + nn.init.normal_(self.waypoint_modality_embedding.weight, std=0.02) + + layer = nn.TransformerEncoderLayer( + d_model=hidden_dim, + nhead=num_attention_heads, + dim_feedforward=feedforward_dim or hidden_dim * 4, + dropout=0.0, + activation="gelu", + batch_first=True, + norm_first=False, + ) + self.encoder = nn.TransformerEncoder(layer, num_layers=num_layers) + + def _type_embedding( + self, + token_type: int, + batch_size: int, + device: torch.device, + ) -> torch.Tensor: + indices = torch.full( + (batch_size, 1), + token_type, + dtype=torch.long, + device=device, + ) + return self.token_type_embedding(indices) + + def encode_tokens(self, observation: torch.Tensor) -> torch.Tensor: + """Encode ``[ACTION, STATE, ACTIVE_GOAL, WP_1..WP_K]`` tokens. + + Args: + observation: Unified flat waypoint observation batch. + + Returns: + Encoded token tensor shaped ``[batch, 3 + K, hidden_dim]``. + """ + fields = parse_waypoint_observation( + observation, + self.num_waypoints, + self.use_relative_observations, + joint_dim=self.joint_dim, + ) + joint = fields["joint"] + assert isinstance(joint, torch.Tensor) + batch_size = observation.shape[0] + device = observation.device + action_tokens = self.action_token.expand( + batch_size, -1, -1 + ) + self._type_embedding( + _TOKEN_ACTION, + batch_size, + device, + ) + state_features = torch.cat( + [joint, fields["end_effector"], fields["last_action"]], + dim=-1, + ) + state_tokens = self.state_proj(state_features).unsqueeze( + 1 + ) + self._type_embedding(_TOKEN_STATE, batch_size, device) + + batch_indices = torch.arange(batch_size, device=device) + active_id = fields["active_id"] + assert isinstance(active_id, torch.Tensor) + active_parts = [ + fields["waypoint_position"][batch_indices, active_id], + fields["waypoint_quaternion"][batch_indices, active_id], + fields["waypoint_joint"][batch_indices, active_id], + fields["position_mask"][batch_indices, active_id].unsqueeze(-1), + fields["rotation_mask"][batch_indices, active_id].unsqueeze(-1), + fields["joint_mask"][batch_indices, active_id].unsqueeze(-1), + ] + if self.use_relative_observations: + active_parts.extend( + [ + fields["active_relative_pose"], + fields["waypoint_joint_error"][batch_indices, active_id], + ] + ) + progress = active_id.float().unsqueeze(-1) / max(self.num_waypoints - 1, 1) + active_parts.append(progress) + active_goal = torch.cat(active_parts, dim=-1) + active_goal_tokens = self.active_goal_proj(active_goal).unsqueeze( + 1 + ) + self._type_embedding(_TOKEN_ACTIVE_GOAL, batch_size, device) + + waypoint_features = [ + fields["waypoint_position"], + fields["waypoint_quaternion"], + fields["waypoint_joint"], + ] + if self.use_relative_observations: + waypoint_features.extend( + [ + fields["waypoint_relative_position"], + fields["waypoint_relative_quaternion"], + fields["waypoint_joint_error"], + ] + ) + waypoint_features.extend( + [ + fields["position_mask"].unsqueeze(-1), + fields["rotation_mask"].unsqueeze(-1), + fields["joint_mask"].unsqueeze(-1), + fields["active_onehot"].unsqueeze(-1), + fields["valid_mask"].unsqueeze(-1), + ] + ) + waypoint_token_types = self.token_type_embedding( + torch.full( + (batch_size, self.num_waypoints), + _TOKEN_WAYPOINT, + dtype=torch.long, + device=device, + ) + ) + waypoint_type = ( + fields["waypoint_type"] + .long() + .clamp( + 0, + _NUM_WAYPOINT_TYPES - 1, + ) + ) + waypoint_tokens = ( + self.waypoint_proj(torch.cat(waypoint_features, dim=-1)) + + waypoint_token_types + + self.waypoint_index_embedding + + self.waypoint_modality_embedding(waypoint_type) + ) + tokens = torch.cat( + [action_tokens, state_tokens, active_goal_tokens, waypoint_tokens], + dim=1, + ) + context_padding = torch.zeros( + batch_size, + 3, + dtype=torch.bool, + device=device, + ) + waypoint_padding = fields["valid_mask"] < 0.5 + padding_mask = torch.cat([context_padding, waypoint_padding], dim=1) + return self.encoder(tokens, src_key_padding_mask=padding_mask) + + +class WaypointTransformerActor(WaypointTransformerEncoder): + """Predict one joint action from fused action and active-goal tokens. + + Args: + observation_dim: Flat unified observation dimension. + action_dim: Output joint-action dimension. + num_waypoints: Maximum ordered waypoint count. + **kwargs: Encoder dimensions accepted by :class:`WaypointTransformerEncoder`. + """ + + def __init__( + self, + observation_dim: int, + action_dim: int, + num_waypoints: int, + **kwargs: object, + ) -> None: + super().__init__( + observation_dim=observation_dim, + num_waypoints=num_waypoints, + **kwargs, + ) + hidden_dim = int(kwargs.get("hidden_dim", 256)) + self.actor_head = _token_head( + hidden_dim, + action_dim, + std=0.01, + input_dim=2 * hidden_dim, + ) + + def forward(self, observation: torch.Tensor) -> torch.Tensor: + """Return deterministic mean actions. + + Args: + observation: Unified flat waypoint observation batch. + + Returns: + Mean action tensor shaped ``[batch, action_dim]``. + """ + encoded = self.encode_tokens(observation) + readout = torch.cat([encoded[:, 0], encoded[:, 2]], dim=-1) + return self.actor_head(readout) + + +class WaypointTransformerCritic(WaypointTransformerEncoder): + """Predict values from the full-context action token. + + Args: + observation_dim: Flat unified observation dimension. + num_waypoints: Maximum ordered waypoint count. + **kwargs: Encoder dimensions accepted by :class:`WaypointTransformerEncoder`. + """ + + def __init__( + self, + observation_dim: int, + num_waypoints: int, + **kwargs: object, + ) -> None: + super().__init__( + observation_dim=observation_dim, + num_waypoints=num_waypoints, + **kwargs, + ) + hidden_dim = int(kwargs.get("hidden_dim", 256)) + self.value_head = _token_head(hidden_dim, 1, std=1.0) + + def forward(self, observation: torch.Tensor) -> torch.Tensor: + """Return one value per observation. + + Args: + observation: Unified flat waypoint observation batch. + + Returns: + Value tensor shaped ``[batch, 1]``. + """ + return self.value_head(self.encode_tokens(observation)[:, 0]) diff --git a/embodichain/learning/rl/normalization.py b/embodichain/learning/rl/normalization.py new file mode 100644 index 000000000..248e798b1 --- /dev/null +++ b/embodichain/learning/rl/normalization.py @@ -0,0 +1,180 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Running statistics for learning-environment observations.""" + +from __future__ import annotations + +from collections.abc import Mapping +from typing import Any + +import torch + +__all__ = ["RunningObservationNormalizer"] + + +class RunningObservationNormalizer: + """Normalize continuous observation fields with running Welford statistics. + + The optional mask leaves semantic fields such as one-hot encodings and + validity bits unchanged while still tracking a single flat observation. + Statistics are updated only when :meth:`update` is called; :meth:`normalize` + is side-effect free so a complete differentiable rollout uses one frozen + normalization transform. + + Args: + observation_dim: Flat observation dimension. + device: Device used for the running statistics. + normalize_mask: Boolean mask where ``True`` selects normalized fields. + initial_count: Positive pseudo-count used to stabilize the first update. + """ + + def __init__( + self, + observation_dim: int, + device: torch.device | str, + normalize_mask: torch.Tensor | None = None, + *, + initial_count: float = 1.0e-4, + ) -> None: + if observation_dim <= 0: + raise ValueError("observation_dim must be positive.") + if initial_count <= 0.0: + raise ValueError("initial_count must be positive.") + + self.observation_dim = int(observation_dim) + self.device = torch.device(device) + self.mean = torch.zeros(self.observation_dim, device=self.device) + self.var = torch.ones(self.observation_dim, device=self.device) + self.count = float(initial_count) + if normalize_mask is None: + normalize_mask = torch.ones( + self.observation_dim, + dtype=torch.bool, + device=self.device, + ) + normalize_mask = torch.as_tensor( + normalize_mask, + dtype=torch.bool, + device=self.device, + ) + if normalize_mask.shape != (self.observation_dim,): + raise ValueError( + "normalize_mask must have shape " + f"({self.observation_dim},), got {tuple(normalize_mask.shape)}." + ) + self.normalize_mask = normalize_mask.clone() + + @torch.no_grad() + def update(self, observations: torch.Tensor) -> None: + """Merge one observation batch into the running statistics. + + Args: + observations: Finite tensor shaped ``[batch, observation_dim]``. + + Raises: + ValueError: If the shape is incompatible or values are non-finite. + """ + observations = torch.as_tensor(observations, device=self.device) + if observations.ndim != 2 or observations.shape[1] != self.observation_dim: + raise ValueError( + "observations must have shape [batch, observation_dim], got " + f"{tuple(observations.shape)}." + ) + if observations.shape[0] == 0: + return + if not bool(torch.isfinite(observations).all()): + raise ValueError("observations must contain only finite values.") + + batch_mean = observations.mean(dim=0) + batch_var = observations.var(dim=0, unbiased=False) + batch_count = int(observations.shape[0]) + delta = batch_mean - self.mean + total_count = self.count + batch_count + self.mean.add_(delta * batch_count / total_count) + merged_m2 = ( + self.var * self.count + + batch_var * batch_count + + delta.square() * self.count * batch_count / total_count + ) + self.var.copy_(merged_m2 / total_count) + self.count = float(total_count) + + def normalize(self, observations: torch.Tensor) -> torch.Tensor: + """Apply the frozen running transform without detaching observations. + + Args: + observations: Tensor ending in the configured observation dimension. + + Returns: + Tensor with continuous fields normalized and semantic fields intact. + + Raises: + ValueError: If the trailing observation dimension is incompatible. + """ + if observations.shape[-1] != self.observation_dim: + raise ValueError( + "observations must end with observation_dim " + f"{self.observation_dim}, got {tuple(observations.shape)}." + ) + normalized = (observations - self.mean) / (self.var.sqrt() + 1.0e-8) + return torch.where(self.normalize_mask, normalized, observations) + + def state_dict(self) -> dict[str, Any]: + """Return a device-independent checkpoint payload. + + Returns: + Mapping containing mean, variance, count, and normalization mask. + """ + return { + "mean": self.mean.detach().cpu(), + "var": self.var.detach().cpu(), + "count": self.count, + "normalize_mask": self.normalize_mask.detach().cpu(), + } + + @torch.no_grad() + def load_state_dict(self, state_dict: Mapping[str, Any]) -> None: + """Restore statistics while validating their observation layout. + + Args: + state_dict: Payload produced by :meth:`state_dict`. + + Raises: + ValueError: If dimensions or the pseudo-count are invalid. + """ + mean = torch.as_tensor(state_dict["mean"], device=self.device) + var = torch.as_tensor(state_dict["var"], device=self.device) + mask = torch.as_tensor( + state_dict.get("normalize_mask", self.normalize_mask), + dtype=torch.bool, + device=self.device, + ) + expected_shape = (self.observation_dim,) + if mean.shape != expected_shape or var.shape != expected_shape: + raise ValueError( + "Normalizer checkpoint shape does not match observation_dim " + f"{self.observation_dim}." + ) + if mask.shape != expected_shape: + raise ValueError("Normalizer checkpoint mask has an incompatible shape.") + count = float(state_dict["count"]) + if count <= 0.0: + raise ValueError("Normalizer checkpoint count must be positive.") + self.mean.copy_(mean) + self.var.copy_(var) + self.normalize_mask.copy_(mask) + self.count = count diff --git a/embodichain/learning/rl/train.py b/embodichain/learning/rl/train.py index f22a4cd80..2ce59f3fa 100644 --- a/embodichain/learning/rl/train.py +++ b/embodichain/learning/rl/train.py @@ -28,8 +28,11 @@ from torch.utils.tensorboard import SummaryWriter from copy import deepcopy -from embodichain.learning.rl.models import build_policy, get_registered_policy_names -from embodichain.learning.rl.models import build_mlp_from_cfg +from embodichain.learning.rl.models import ( + build_model_from_cfg, + build_policy, + get_registered_policy_names, +) from embodichain.learning.rl.algo import ( RolloutKind, build_algo, @@ -125,12 +128,14 @@ def _build_learning_policy( actor_cfg = policy_block.get("actor") critic_cfg = policy_block.get("critic") actor = ( - build_mlp_from_cfg(actor_cfg, obs_dim, action_dim) + build_model_from_cfg(actor_cfg, obs_dim, action_dim, role="actor") if actor_cfg is not None else None ) critic = ( - build_mlp_from_cfg(critic_cfg, obs_dim, 1) if critic_cfg is not None else None + build_model_from_cfg(critic_cfg, obs_dim, 1, role="critic") + if critic_cfg is not None + else None ) policy = build_policy( policy_block, @@ -241,9 +246,23 @@ def _train_learning_env( diff_cfg = DifferentiableTrainerCfg( segment_length=segment_length, update_horizon=update_horizon, + rollout_mode=str(trainer_cfg.get("rollout_mode", "segmented")), + gradient_accumulation_steps=int( + trainer_cfg.get("gradient_accumulation_steps", 1) + ), deterministic_actions=bool( trainer_cfg.get("deterministic_actions", False) ), + clip_actions_to_space=bool( + trainer_cfg.get("clip_actions_to_space", False) + ), + action_adjoint_max_norm=float( + trainer_cfg.get("action_adjoint_max_norm", 0.0) + ), + normalize_observations=bool( + trainer_cfg.get("normalize_observations", False) + ), + rollout_seed=seed, checkpoint_dir=str(checkpoint_dir), experiment_name=exp_name, save_frequency_updates=int( @@ -264,7 +283,12 @@ def _train_learning_env( writer=writer, eval_env=eval_env, ) - default_steps = iterations * update_horizon * num_envs + default_steps = ( + iterations + * update_horizon + * num_envs + * diff_cfg.gradient_accumulation_steps + ) else: buffer_size = int( trainer_cfg.get("buffer_size", trainer_cfg.get("rollout_steps", 256)) @@ -288,8 +312,15 @@ def _train_learning_env( best_eval_mode=trainer_cfg.get("best_eval_mode", "max"), ) default_steps = iterations * buffer_size * num_envs - total_timesteps = int(trainer_cfg.get("total_timesteps", default_steps)) - trainer.train(total_timesteps) + if ( + trainer_class is DifferentiableTrainer + and diff_cfg.rollout_mode == "complete" + and "total_timesteps" not in trainer_cfg + ): + trainer.train(total_updates=iterations) + else: + total_timesteps = int(trainer_cfg.get("total_timesteps", default_steps)) + trainer.train(total_timesteps) trainer.save_checkpoint() return trainer.get_summary() finally: diff --git a/tests/learning/test_apg.py b/tests/learning/test_apg.py index 32a7348fc..ea5e6ca14 100644 --- a/tests/learning/test_apg.py +++ b/tests/learning/test_apg.py @@ -30,9 +30,14 @@ APG, APGCfg, build_algo, + complete_discounted_return, get_registered_algo_names, segmented_discounted_return, ) +from embodichain.learning.rl.gradients import ( + BatchedGradientNormStats, + clip_batched_gradient_norm, +) from embodichain.learning.rl.collector import ( DifferentiableCollector, DifferentiableRollout, @@ -149,6 +154,110 @@ def test_segmented_discounted_return_restarts_discount_after_done() -> None: assert torch.equal(returns, torch.tensor([5.0])) +def test_complete_discounted_return_stops_after_first_done() -> None: + observation = torch.zeros((1, 1)) + transitions = [] + for reward, done in ((1.0, False), (2.0, True), (100.0, False)): + transitions.append( + DifferentiableTransition( + observation=observation, + policy_output=TensorDict( + {"action": torch.zeros((1, 1))}, + batch_size=[1], + ), + reward=torch.tensor([reward]), + terminated=torch.tensor([done]), + truncated=torch.tensor([False]), + next_observation=observation, + info={}, + ) + ) + + returns = complete_discounted_return( + DifferentiableRollout(observation, tuple(transitions)), + gamma=0.5, + ) + + assert torch.equal(returns, torch.tensor([2.0])) + + +def test_action_adjoint_clip_is_per_environment_and_overflow_safe() -> None: + gradient = torch.tensor( + [[3.0, 4.0], [0.3, 0.4], [2.0e30, 2.0e30], [float("inf"), 0.0]] + ) + stats = BatchedGradientNormStats("cpu") + + clipped = clip_batched_gradient_norm(gradient, 1.0, stats) + + torch.testing.assert_close(clipped[0], torch.tensor([0.6, 0.8])) + torch.testing.assert_close(clipped[1], gradient[1]) + torch.testing.assert_close( + clipped[2], + torch.full((2,), 2.0**-0.5), + rtol=1.0e-5, + atol=1.0e-6, + ) + torch.testing.assert_close(clipped[3], torch.zeros(2)) + assert torch.isfinite(clipped).all() + assert stats.rows == 4 + assert stats.finite_rows == 3 + assert stats.clipped_rows == 2 + assert stats.nonfinite_rows == 1 + + +def test_complete_rollout_objective_scale_equalizes_sequence_lengths() -> None: + policy = _make_policy() + algorithm = APG(APGCfg(device="cpu", max_grad_norm=100.0), policy) + observation = torch.zeros((4, 1)) + waypoint_counts = torch.tensor([1.0, 2.0, 4.0, 8.0]) + transition = DifferentiableTransition( + observation=observation, + policy_output=TensorDict( + {"action": policy.actor.weight.sum() * torch.ones((4, 1))}, + batch_size=[4], + ), + reward=30.0 * waypoint_counts + policy.actor.weight.sum() * 0.0, + terminated=torch.zeros(4, dtype=torch.bool), + truncated=torch.zeros(4, dtype=torch.bool), + next_observation=observation, + info={}, + ) + rollout = DifferentiableRollout(observation, (transition,)) + + algorithm.begin_update() + algorithm.accumulate_complete_rollout( + rollout, + objective_scale=waypoint_counts.reciprocal(), + ) + metrics = algorithm.finish_update() + + assert metrics["objective"] == pytest.approx(30.0) + + +def test_apg_skips_gradient_above_preclip_safety_limit() -> None: + policy = _make_policy() + algorithm = APG( + APGCfg( + device="cpu", + max_grad_norm=1.0, + max_grad_norm_before_clip=0.01, + ), + policy, + ) + initial_weight = policy.actor.weight.detach().clone() + rollout = DifferentiableCollector( + _LinearDifferentiableEnv(), + policy, + torch.device("cpu"), + ).collect(num_steps=2, deterministic=True) + + metrics = algorithm.update(rollout) + + assert metrics["skipped_update"] == 1.0 + assert metrics["skipped_excessive_gradient"] == 1.0 + assert torch.equal(policy.actor.weight, initial_weight) + + def test_apg_entropy_uses_reward_discount_and_done_reset_semantics() -> None: policy = _make_policy() observation = torch.zeros((1, 1)) diff --git a/tests/learning/test_differentiable_trainer.py b/tests/learning/test_differentiable_trainer.py index a36023eed..ea9e5d95e 100644 --- a/tests/learning/test_differentiable_trainer.py +++ b/tests/learning/test_differentiable_trainer.py @@ -27,8 +27,10 @@ from gymnasium.spaces import Box from embodichain.learning.rl import ( + DifferentiableRolloutSpec, DifferentiableTrainer, DifferentiableTrainerCfg, + stratified_rollout_value, ) from embodichain.learning.rl.algo import APG, APGCfg from embodichain.learning.rl.models import ActorOnly @@ -74,6 +76,35 @@ def close(self) -> None: return None +class _ScheduledCompleteRolloutEnv(_QuadraticActionEnv): + def __init__(self, num_envs: int = 2) -> None: + super().__init__(num_envs) + self.prepared_indices: list[int] = [] + self.reset_calls = 0 + self.current_waypoint_count = 1 + + def prepare_differentiable_rollout( + self, + rollout_index: int, + ) -> DifferentiableRolloutSpec: + self.prepared_indices.append(rollout_index) + self.current_waypoint_count = stratified_rollout_value(rollout_index, 1, 3) + return DifferentiableRolloutSpec( + num_steps=2 * self.current_waypoint_count, + objective_scale=1.0 / self.current_waypoint_count, + metadata={"waypoint_count": float(self.current_waypoint_count)}, + ) + + def reset( + self, + *, + seed: int | None = None, + options: Mapping[str, Any] | None = None, + ) -> tuple[torch.Tensor, dict[str, Any]]: + self.reset_calls += 1 + return super().reset(seed=seed, options=options) + + def _make_components( ent_coef: float = 0.0, ) -> tuple[_QuadraticActionEnv, ActorOnly, APG]: @@ -116,6 +147,76 @@ def test_trainer_updates_policy_and_detaches_each_segment() -> None: assert env.detach_calls == 2 +def test_complete_rollout_mode_resets_each_scheduled_microbatch() -> None: + env = _ScheduledCompleteRolloutEnv() + actor = nn.Linear(1, 1, bias=False) + nn.init.constant_(actor.weight, 0.5) + policy = ActorOnly(1, 1, env.device, actor=actor) + algorithm = APG( + APGCfg( + device="cpu", + optimizer=OptimizerCfg(learning_rate=0.05), + max_grad_norm=10.0, + ), + policy, + ) + trainer = DifferentiableTrainer( + DifferentiableTrainerCfg( + rollout_mode="complete", + update_horizon=6, + gradient_accumulation_steps=3, + deterministic_actions=True, + clip_actions_to_space=True, + rollout_seed=17, + ), + env, + policy, + algorithm, + ) + + summary = trainer.train(total_timesteps=24) + + assert summary["num_updates"] == 1 + assert summary["global_step"] == 24 + assert env.prepared_indices == [0, 1, 2] + assert env.reset_calls == 3 + assert env.detach_calls == 3 + assert summary["last_train_metrics"][ + "train/rollout_waypoint_count_mean" + ] == pytest.approx(2.0) + + +def test_complete_rollout_mode_honors_exact_optimizer_update_budget() -> None: + env = _ScheduledCompleteRolloutEnv() + actor = nn.Linear(1, 1, bias=False) + policy = ActorOnly(1, 1, env.device, actor=actor) + algorithm = APG(APGCfg(device="cpu", max_grad_norm=10.0), policy) + trainer = DifferentiableTrainer( + DifferentiableTrainerCfg( + rollout_mode="complete", + update_horizon=6, + deterministic_actions=True, + ), + env, + policy, + algorithm, + ) + + summary = trainer.train(total_updates=2) + + assert summary["num_updates"] == 2 + assert env.prepared_indices == [0, 1] + assert summary["global_step"] == 12 + + +def test_stratified_rollout_value_balances_and_rotates_cycles() -> None: + first = [stratified_rollout_value(index, 1, 3) for index in range(3)] + second = [stratified_rollout_value(index, 1, 3) for index in range(3, 6)] + + assert first == [1, 2, 3] + assert second == [2, 3, 1] + + def test_update_horizon_keeps_optimizer_budget_fixed_across_segment_lengths() -> None: short_env, short_policy, short_algorithm = _make_components(ent_coef=0.2) short_trainer = DifferentiableTrainer( diff --git a/tests/learning/test_observation_normalization.py b/tests/learning/test_observation_normalization.py new file mode 100644 index 000000000..d871ca877 --- /dev/null +++ b/tests/learning/test_observation_normalization.py @@ -0,0 +1,76 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Tests for semantic-mask-aware running observation normalization.""" + +from __future__ import annotations + +import pytest +import torch + +from embodichain.learning.rl.normalization import RunningObservationNormalizer + + +def test_running_normalizer_matches_combined_population_statistics() -> None: + normalizer = RunningObservationNormalizer(3, "cpu") + first = torch.tensor([[1.0, 2.0, 3.0], [3.0, 4.0, 5.0]]) + second = torch.tensor([[5.0, 6.0, 7.0]]) + + normalizer.update(first) + normalizer.update(second) + + combined = torch.cat([first, second]) + # The small pseudo-count has zero mean, unit variance, matching the NMG + # training reference rather than an exact batch-only moment. + expected_count = 3.0001 + expected_mean = combined.sum(dim=0) / expected_count + expected_m2 = ( + combined.square().sum(dim=0) + 1.0e-4 - expected_count * expected_mean.square() + ) + torch.testing.assert_close(normalizer.mean, expected_mean) + torch.testing.assert_close(normalizer.var, expected_m2 / expected_count) + assert normalizer.count == pytest.approx(expected_count) + + +def test_running_normalizer_preserves_semantic_fields_and_gradients() -> None: + normalizer = RunningObservationNormalizer( + 3, + "cpu", + normalize_mask=torch.tensor([True, False, True]), + ) + normalizer.update(torch.tensor([[1.0, 0.0, 3.0], [3.0, 1.0, 7.0]])) + observation = torch.tensor([[2.0, 1.0, 5.0]], requires_grad=True) + + normalized = normalizer.normalize(observation) + normalized.sum().backward() + + assert normalized[0, 1] == 1.0 + assert observation.grad is not None + assert torch.isfinite(observation.grad).all() + assert observation.grad[0, 1] == 1.0 + + +def test_running_normalizer_checkpoint_round_trip() -> None: + source = RunningObservationNormalizer(2, "cpu", torch.tensor([True, False])) + source.update(torch.tensor([[2.0, 1.0], [4.0, 0.0]])) + restored = RunningObservationNormalizer(2, "cpu") + + restored.load_state_dict(source.state_dict()) + + torch.testing.assert_close(restored.mean, source.mean) + torch.testing.assert_close(restored.var, source.var) + assert restored.count == source.count + assert torch.equal(restored.normalize_mask, source.normalize_mask) diff --git a/tests/learning/test_waypoint_transformer.py b/tests/learning/test_waypoint_transformer.py new file mode 100644 index 000000000..fe84dac18 --- /dev/null +++ b/tests/learning/test_waypoint_transformer.py @@ -0,0 +1,145 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Tests for unified Cartesian/joint waypoint Transformer models.""" + +from __future__ import annotations + +import torch + +from embodichain.learning.rl.models import ( + WaypointTransformerActor, + WaypointTransformerCritic, + build_model_from_cfg, + parse_waypoint_observation, + waypoint_observation_dim, + waypoint_observation_normalize_mask, +) + +_NUM_WAYPOINTS = 4 +_OBSERVATION_DIM = waypoint_observation_dim(_NUM_WAYPOINTS, True) + + +def _valid_observation(batch_size: int = 3) -> torch.Tensor: + observation = torch.zeros(batch_size, _OBSERVATION_DIM) + cursor = 7 + 7 + _NUM_WAYPOINTS * (3 + 4 + 7) + observation[:, cursor] = 1.0 + cursor += _NUM_WAYPOINTS + observation[:, cursor : cursor + _NUM_WAYPOINTS] = 1.0 + cursor += _NUM_WAYPOINTS + observation[:, cursor : cursor + _NUM_WAYPOINTS] = 1.0 + cursor += _NUM_WAYPOINTS + observation[:, cursor : cursor + _NUM_WAYPOINTS] = 1.0 + return observation + + +def _actor() -> WaypointTransformerActor: + return WaypointTransformerActor( + observation_dim=_OBSERVATION_DIM, + action_dim=7, + num_waypoints=_NUM_WAYPOINTS, + hidden_dim=64, + num_attention_heads=4, + num_layers=1, + ) + + +def test_waypoint_actor_and_critic_have_expected_shapes_and_gradients() -> None: + torch.manual_seed(7) + actor = _actor().train() + critic = WaypointTransformerCritic( + observation_dim=_OBSERVATION_DIM, + num_waypoints=_NUM_WAYPOINTS, + hidden_dim=64, + num_attention_heads=4, + num_layers=1, + ).train() + observation = _valid_observation() + + action = actor(observation) + value = critic(observation) + (action.square().mean() + value.square().mean()).backward() + + assert action.shape == (3, 7) + assert value.shape == (3, 1) + assert torch.isfinite(action).all() + assert torch.isfinite(value).all() + assert actor.action_token.grad is not None + assert critic.action_token.grad is not None + + +def test_waypoint_actor_uses_future_valid_waypoints() -> None: + torch.manual_seed(11) + actor = _actor().eval() + base = _valid_observation(batch_size=1) + changed = base.clone() + # WP3 position is a future token while WP0 remains active. + position_start = 7 + 7 + changed[:, position_start + 3 * 3 : position_start + 3 * 4] = 5.0 + + with torch.no_grad(): + base_action = actor(base) + changed_action = actor(changed) + + assert not torch.allclose(base_action, changed_action) + + +def test_waypoint_parser_clamps_active_id_to_last_valid_slot() -> None: + observation = _valid_observation(batch_size=1) + cursor = 7 + 7 + _NUM_WAYPOINTS * (3 + 4 + 7) + observation[:, cursor : cursor + _NUM_WAYPOINTS] = 0.0 + observation[:, cursor + 3] = 1.0 + cursor += _NUM_WAYPOINTS + observation[:, cursor : cursor + _NUM_WAYPOINTS] = 0.0 + observation[:, cursor : cursor + 2] = 1.0 + + fields = parse_waypoint_observation(observation, _NUM_WAYPOINTS, True) + + assert torch.equal(fields["active_id"], torch.tensor([1])) + + +def test_waypoint_normalization_mask_excludes_all_semantic_fields() -> None: + mask = waypoint_observation_normalize_mask(_NUM_WAYPOINTS, True) + + assert mask.shape == (_OBSERVATION_DIM,) + assert int((~mask).sum()) == 6 * _NUM_WAYPOINTS + 7 * _NUM_WAYPOINTS + + +def test_waypoint_transformer_builds_from_learning_config() -> None: + module_config = { + "type": "waypoint_transformer", + "network_cfg": { + "num_waypoints": _NUM_WAYPOINTS, + "hidden_dim": 64, + "num_attention_heads": 4, + "num_layers": 1, + }, + } + actor = build_model_from_cfg( + module_config, + _OBSERVATION_DIM, + 7, + role="actor", + ) + critic = build_model_from_cfg( + module_config, + _OBSERVATION_DIM, + 1, + role="critic", + ) + + assert isinstance(actor, WaypointTransformerActor) + assert isinstance(critic, WaypointTransformerCritic) From 24435ee957a8fd2daa101bdb4726897562498aca Mon Sep 17 00:00:00 2001 From: yuecideng Date: Mon, 31 Aug 2026 13:34:27 +0800 Subject: [PATCH 2/3] feat(learning): support deterministic lightweight training --- agent_context/topics/rl-learning/rl-learning.md | 3 +++ embodichain/learning/rl/train.py | 9 +++++---- 2 files changed, 8 insertions(+), 4 deletions(-) diff --git a/agent_context/topics/rl-learning/rl-learning.md b/agent_context/topics/rl-learning/rl-learning.md index b57daff79..79454f109 100644 --- a/agent_context/topics/rl-learning/rl-learning.md +++ b/agent_context/topics/rl-learning/rl-learning.md @@ -90,6 +90,9 @@ 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 diff --git a/embodichain/learning/rl/train.py b/embodichain/learning/rl/train.py index 2ce59f3fa..0dba17c58 100644 --- a/embodichain/learning/rl/train.py +++ b/embodichain/learning/rl/train.py @@ -46,7 +46,7 @@ from embodichain.learning.rl.routing import get_trainer_class from embodichain.learning.rl.utils import dict_to_tensordict, flatten_dict_observation from embodichain.learning.rl.utils.trainer import Trainer -from embodichain.utils import logger +from embodichain.utils import logger, set_seed from embodichain.lab.gym.utils.registration import ( build_env, discover_task_packages, @@ -183,9 +183,10 @@ def _train_learning_env( raise ValueError("CUDA was requested but is not available.") if device.type == "cuda": torch.cuda.set_device(device) - torch.cuda.manual_seed_all(seed) - np.random.seed(seed) - torch.manual_seed(seed) + set_seed( + seed, + deterministic=bool(trainer_cfg.get("torch_deterministic", False)), + ) env_block = trainer_cfg["learning_env"] if isinstance(env_block, str): From ce53bce60dab10eb75ef51bddc9a0b960379fb14 Mon Sep 17 00:00:00 2001 From: yuecideng Date: Mon, 31 Aug 2026 14:53:20 +0800 Subject: [PATCH 3/3] refactor(learning): split waypoint model from foundation --- agent_context/MAP.yaml | 1 - .../topics/rl-learning/rl-learning.md | 11 +- .../embodichain.learning.rl.models.rst | 13 +- docs/source/api_reference/public_api.rst | 17 - embodichain/learning/rl/models/__init__.py | 71 --- .../rl/models/waypoint_transformer.py | 538 ------------------ embodichain/learning/rl/train.py | 13 +- tests/learning/test_waypoint_transformer.py | 145 ----- 8 files changed, 8 insertions(+), 801 deletions(-) delete mode 100644 embodichain/learning/rl/models/waypoint_transformer.py delete mode 100644 tests/learning/test_waypoint_transformer.py diff --git a/agent_context/MAP.yaml b/agent_context/MAP.yaml index 89d7eb26a..07170e38f 100644 --- a/agent_context/MAP.yaml +++ b/agent_context/MAP.yaml @@ -497,7 +497,6 @@ topics: - gradient accumulation - action adjoint - observation normalization - - waypoint transformer - SyncCollector - learning_env - gym_config diff --git a/agent_context/topics/rl-learning/rl-learning.md b/agent_context/topics/rl-learning/rl-learning.md index 79454f109..5f20af008 100644 --- a/agent_context/topics/rl-learning/rl-learning.md +++ b/agent_context/topics/rl-learning/rl-learning.md @@ -39,7 +39,7 @@ CLI --config → trainer / policy / algorithm blocks → choose trainer.learning_env or trainer.gym_config → build environment - → build policy and optional MLP or waypoint-Transformer modules + → build policy and optional MLP modules → build algorithm config from the registry → route by algorithm.rollout_kind → train, evaluate, log, and checkpoint @@ -149,7 +149,7 @@ supports a pre-clip policy-gradient safety limit that skips unsafe updates. | PPO, GRPO, APG implementations | `algo/ppo.py`, `algo/grpo.py`, `algo/apg.py` | | Standard rollout storage and views | `buffer/` | | Standard and differentiable collection | `collector/` | -| Policy interface, actor-critic, actor-only, MLP/waypoint Transformer builders | `models/` | +| 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` | @@ -220,13 +220,6 @@ example is an experimental gradient reference, not a general simulator task. 3. Ensure its outputs satisfy every intended algorithm. 4. Provide graph-preserving sampling if used with differentiable rollouts. -The built-in `waypoint_transformer` module consumes the unified ordered -constraint layout: joint state, end-effector pose, absolute pose/joint targets, -active/valid and modality masks, last action, optional relative pose/joint -errors, and waypoint type. Its token sequence is -`[ACTION, STATE, ACTIVE_GOAL, WP_1..WP_K]`; attention is bidirectional so the -action can depend on future valid waypoints. - ### Add a Lightweight Environment 1. Implement `LearningVecEnv`, or `DifferentiableVecEnv` for APG. diff --git a/docs/source/api_reference/embodichain/embodichain.learning.rl.models.rst b/docs/source/api_reference/embodichain/embodichain.learning.rl.models.rst index 3eb79fefe..a4e382278 100644 --- a/docs/source/api_reference/embodichain/embodichain.learning.rl.models.rst +++ b/docs/source/api_reference/embodichain/embodichain.learning.rl.models.rst @@ -8,10 +8,8 @@ Overview Policy-network registration and model construction APIs for RL agents. Policies implement the :class:`Policy` ABC; the built-in actor-critic variants are -:class:`ActorCritic` and :class:`ActorOnly`. Backbones include :class:`MLP` and -the full-context :class:`WaypointTransformerActor` / -:class:`WaypointTransformerCritic` pair for ordered mixed-modality constraints. -:func:`build_policy` constructs a policy from a config block, with +:class:`ActorCritic` and :class:`ActorOnly`, both built on the :class:`MLP` +backbone. :func:`build_policy` constructs a policy from a config block, with :class:`~embodichain.learning.rl.utils.config.AlgorithmCfg`-style registration through :func:`register_policy` / :func:`get_policy_class`. @@ -23,23 +21,16 @@ through :func:`register_policy` / :func:`get_policy_class`. ActorCritic ActorOnly MLP - WaypointTransformerEncoder - WaypointTransformerActor - WaypointTransformerCritic .. rubric:: Functions .. autosummary:: - build_model_from_cfg build_mlp_from_cfg build_policy get_policy_class get_registered_policy_names register_policy - parse_waypoint_observation - waypoint_observation_dim - waypoint_observation_normalize_mask .. automodule:: embodichain.learning.rl.models :members: diff --git a/docs/source/api_reference/public_api.rst b/docs/source/api_reference/public_api.rst index 418c20f74..763816e62 100644 --- a/docs/source/api_reference/public_api.rst +++ b/docs/source/api_reference/public_api.rst @@ -1888,23 +1888,6 @@ embodichain.learning.rl.models.policy Policy -embodichain.learning.rl.models.waypoint_transformer ---------------------------------------------------- - -Unified tokenization and full-context models for ordered Cartesian, joint, and -mixed-modality waypoint policies. - -.. currentmodule:: embodichain.learning.rl.models.waypoint_transformer - -.. autosummary:: - - WaypointTransformerActor - WaypointTransformerCritic - WaypointTransformerEncoder - parse_waypoint_observation - waypoint_observation_dim - waypoint_observation_normalize_mask - embodichain.learning.rl.normalization ------------------------------------- diff --git a/embodichain/learning/rl/models/__init__.py b/embodichain/learning/rl/models/__init__.py index f7eae3d5c..fa6d46ae6 100644 --- a/embodichain/learning/rl/models/__init__.py +++ b/embodichain/learning/rl/models/__init__.py @@ -28,14 +28,6 @@ from .actor_only import ActorOnly from .policy import Policy from .mlp import MLP -from .waypoint_transformer import ( - WaypointTransformerActor, - WaypointTransformerCritic, - WaypointTransformerEncoder, - parse_waypoint_observation, - waypoint_observation_dim, - waypoint_observation_normalize_mask, -) # In-module policy registry _POLICY_REGISTRY: Dict[str, Type[Policy]] = {} @@ -160,62 +152,6 @@ def build_mlp_from_cfg(module_cfg: Dict, in_dim: int, out_dim: int) -> MLP: return model -def build_model_from_cfg( - module_cfg: Dict, - in_dim: int, - out_dim: int, - *, - role: str, -) -> torch.nn.Module: - """Build an actor or critic module from a JSON-like config. - - Args: - module_cfg: Module configuration with ``type`` and ``network_cfg``. - in_dim: Flat observation dimension. - out_dim: Requested output dimension. - role: ``"actor"`` or ``"critic"``. - - Returns: - Constructed PyTorch module. - - Raises: - ValueError: If the module type or role is unsupported. - """ - module_type = module_cfg.get("type", "").lower() - if module_type == "mlp": - return build_mlp_from_cfg(module_cfg, in_dim, out_dim) - if module_type != "waypoint_transformer": - raise ValueError( - "Supported actor/critic module types are 'mlp' and " - f"'waypoint_transformer', got {module_type!r}." - ) - if role not in {"actor", "critic"}: - raise ValueError("role must be 'actor' or 'critic'.") - - network_cfg = dict(module_cfg.get("network_cfg", {})) - kwargs = { - "observation_dim": in_dim, - "num_waypoints": int(network_cfg["num_waypoints"]), - "joint_dim": int(network_cfg.get("joint_dim", 7)), - "use_relative_observations": bool( - network_cfg.get("use_relative_observations", True) - ), - "hidden_dim": int(network_cfg.get("hidden_dim", 256)), - "num_attention_heads": int(network_cfg.get("num_attention_heads", 4)), - "num_layers": int(network_cfg.get("num_layers", 2)), - "feedforward_dim": ( - int(network_cfg["feedforward_dim"]) - if int(network_cfg.get("feedforward_dim", 0)) > 0 - else None - ), - } - if role == "actor": - return WaypointTransformerActor(action_dim=out_dim, **kwargs) - if out_dim != 1: - raise ValueError("A waypoint Transformer critic must have out_dim=1.") - return WaypointTransformerCritic(**kwargs) - - # default registrations register_policy("actor_critic", ActorCritic) register_policy("actor_only", ActorOnly) @@ -226,15 +162,8 @@ def build_model_from_cfg( "register_policy", "get_registered_policy_names", "build_policy", - "build_model_from_cfg", "build_mlp_from_cfg", "get_policy_class", "Policy", "MLP", - "WaypointTransformerActor", - "WaypointTransformerCritic", - "WaypointTransformerEncoder", - "parse_waypoint_observation", - "waypoint_observation_dim", - "waypoint_observation_normalize_mask", ] diff --git a/embodichain/learning/rl/models/waypoint_transformer.py b/embodichain/learning/rl/models/waypoint_transformer.py deleted file mode 100644 index 806456eb9..000000000 --- a/embodichain/learning/rl/models/waypoint_transformer.py +++ /dev/null @@ -1,538 +0,0 @@ -# ---------------------------------------------------------------------------- -# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ---------------------------------------------------------------------------- - -"""Transformer models for ordered Cartesian and joint waypoint constraints.""" - -from __future__ import annotations - -import math - -import torch -import torch.nn as nn - -__all__ = [ - "WaypointTransformerActor", - "WaypointTransformerCritic", - "WaypointTransformerEncoder", - "parse_waypoint_observation", - "waypoint_observation_dim", - "waypoint_observation_normalize_mask", -] - -_TOKEN_ACTION = 0 -_TOKEN_STATE = 1 -_TOKEN_ACTIVE_GOAL = 2 -_TOKEN_WAYPOINT = 3 -_NUM_TOKEN_TYPES = 4 -_NUM_WAYPOINT_TYPES = 3 - - -def _layer_init( - layer: nn.Linear, - std: float = math.sqrt(2.0), - bias_const: float = 0.0, -) -> nn.Linear: - torch.nn.init.orthogonal_(layer.weight, std) - torch.nn.init.constant_(layer.bias, bias_const) - return layer - - -def waypoint_observation_dim( - num_waypoints: int, - use_relative_observations: bool, - *, - joint_dim: int = 7, -) -> int: - """Return the flat dimension of the unified waypoint observation layout. - - Args: - num_waypoints: Maximum ordered waypoint count. - use_relative_observations: Whether relative pose/joint fields are present. - joint_dim: Controlled joint dimension. - - Returns: - Required flat observation dimension. - - Raises: - ValueError: If a dimension is not positive. - """ - if num_waypoints <= 0: - raise ValueError("num_waypoints must be positive.") - if joint_dim <= 0: - raise ValueError("joint_dim must be positive.") - num_waypoints = int(num_waypoints) - joint_dim = int(joint_dim) - dimension = ( - joint_dim - + 7 - + num_waypoints * (3 + 4 + joint_dim) - + 5 * num_waypoints - + joint_dim - + num_waypoints - ) - if use_relative_observations: - dimension += 7 + num_waypoints * (3 + 4 + joint_dim) - return dimension - - -def waypoint_observation_normalize_mask( - num_waypoints: int, - use_relative_observations: bool, - *, - joint_dim: int = 7, - device: torch.device | str | None = None, -) -> torch.Tensor: - """Build a mask that preserves waypoint semantics during normalization. - - Args: - num_waypoints: Maximum ordered waypoint count. - use_relative_observations: Whether relative pose/joint fields are present. - joint_dim: Controlled joint dimension. - device: Optional output device. - - Returns: - Boolean mask where ``True`` selects continuous normalized fields. - """ - num_waypoints = int(num_waypoints) - joint_dim = int(joint_dim) - dimension = waypoint_observation_dim( - num_waypoints, - use_relative_observations, - joint_dim=joint_dim, - ) - mask = torch.ones(dimension, dtype=torch.bool, device=device) - cursor = joint_dim + 7 + num_waypoints * (3 + 4 + joint_dim) - mask[cursor : cursor + 5 * num_waypoints] = False - cursor += 5 * num_waypoints + joint_dim - if use_relative_observations: - cursor += 7 + num_waypoints * (3 + 4) - mask[cursor : cursor + num_waypoints * joint_dim] = False - cursor += num_waypoints * joint_dim - mask[cursor : cursor + num_waypoints] = False - return mask - - -def parse_waypoint_observation( - observation: torch.Tensor, - num_waypoints: int, - use_relative_observations: bool, - *, - joint_dim: int = 7, -) -> dict[str, torch.Tensor | None]: - """Slice a flat unified waypoint observation into semantic fields. - - The layout is ``joint, eef, waypoint pose, waypoint joint, active, valid, - position/rotation/joint masks, last action, optional relative fields, type``. - - Args: - observation: Flat tensor shaped ``[batch, features]``. - num_waypoints: Maximum ordered waypoint count. - use_relative_observations: Whether relative fields are present. - joint_dim: Controlled joint dimension. - - Returns: - Mapping from semantic field names to tensor views. - - Raises: - ValueError: If the observation rank or feature dimension is invalid. - """ - if observation.ndim != 2: - raise ValueError("observation must have shape [batch, features].") - expected_dim = waypoint_observation_dim( - num_waypoints, - use_relative_observations, - joint_dim=joint_dim, - ) - if observation.shape[1] != expected_dim: - raise ValueError( - f"Expected waypoint observation dimension {expected_dim}, " - f"got {observation.shape[1]}." - ) - - n = int(num_waypoints) - d = int(joint_dim) - cursor = 0 - joint = observation[:, cursor : cursor + d] - cursor += d - end_effector = observation[:, cursor : cursor + 7] - cursor += 7 - waypoint_position = observation[:, cursor : cursor + 3 * n].reshape(-1, n, 3) - cursor += 3 * n - waypoint_quaternion = observation[:, cursor : cursor + 4 * n].reshape(-1, n, 4) - cursor += 4 * n - waypoint_joint = observation[:, cursor : cursor + d * n].reshape(-1, n, d) - cursor += d * n - active_onehot = observation[:, cursor : cursor + n] - cursor += n - valid_mask = observation[:, cursor : cursor + n] - cursor += n - position_mask = observation[:, cursor : cursor + n] - cursor += n - rotation_mask = observation[:, cursor : cursor + n] - cursor += n - joint_mask = observation[:, cursor : cursor + n] - cursor += n - last_action = observation[:, cursor : cursor + d] - cursor += d - - active_relative_pose = None - waypoint_relative_position = None - waypoint_relative_quaternion = None - waypoint_joint_error = None - if use_relative_observations: - active_relative_pose = observation[:, cursor : cursor + 7] - cursor += 7 - waypoint_relative_position = observation[:, cursor : cursor + 3 * n].reshape( - -1, n, 3 - ) - cursor += 3 * n - waypoint_relative_quaternion = observation[:, cursor : cursor + 4 * n].reshape( - -1, n, 4 - ) - cursor += 4 * n - waypoint_joint_error = observation[:, cursor : cursor + d * n].reshape(-1, n, d) - cursor += d * n - waypoint_type = observation[:, cursor : cursor + n] - - active_id = active_onehot.argmax(dim=-1).long() - valid_count = valid_mask.sum(dim=-1).long() - active_id = torch.clamp(torch.minimum(active_id, valid_count - 1), min=0) - return { - "joint": joint, - "end_effector": end_effector, - "last_action": last_action, - "active_relative_pose": active_relative_pose, - "waypoint_relative_position": waypoint_relative_position, - "waypoint_relative_quaternion": waypoint_relative_quaternion, - "waypoint_joint_error": waypoint_joint_error, - "active_onehot": active_onehot, - "valid_mask": valid_mask, - "position_mask": position_mask, - "rotation_mask": rotation_mask, - "joint_mask": joint_mask, - "active_id": active_id, - "waypoint_position": waypoint_position, - "waypoint_quaternion": waypoint_quaternion, - "waypoint_joint": waypoint_joint, - "waypoint_type": waypoint_type, - } - - -def _token_head( - hidden_dim: int, - output_dim: int, - std: float, - *, - input_dim: int | None = None, -) -> nn.Sequential: - input_dim = hidden_dim if input_dim is None else int(input_dim) - return nn.Sequential( - nn.LayerNorm(input_dim), - _layer_init(nn.Linear(input_dim, hidden_dim)), - nn.Tanh(), - _layer_init(nn.Linear(hidden_dim, output_dim), std=std), - ) - - -class WaypointTransformerEncoder(nn.Module): - """Encode state, active-goal, and all waypoint tokens bidirectionally. - - Args: - observation_dim: Flat unified observation dimension. - num_waypoints: Maximum ordered waypoint count. - joint_dim: Controlled joint dimension. - use_relative_observations: Whether relative fields are present. - hidden_dim: Transformer embedding dimension. - num_attention_heads: Attention-head count. - num_layers: Encoder-layer count. - feedforward_dim: Optional feed-forward dimension; defaults to four - times ``hidden_dim``. - """ - - def __init__( - self, - observation_dim: int, - num_waypoints: int, - *, - joint_dim: int = 7, - use_relative_observations: bool = True, - hidden_dim: int = 256, - num_attention_heads: int = 4, - num_layers: int = 2, - feedforward_dim: int | None = None, - ) -> None: - super().__init__() - self.num_waypoints = int(num_waypoints) - self.joint_dim = int(joint_dim) - self.use_relative_observations = bool(use_relative_observations) - expected_dim = waypoint_observation_dim( - self.num_waypoints, - self.use_relative_observations, - joint_dim=self.joint_dim, - ) - if int(observation_dim) != expected_dim: - raise ValueError( - f"WaypointTransformerEncoder expected observation_dim " - f"{expected_dim}, got {observation_dim}." - ) - if hidden_dim % num_attention_heads != 0: - raise ValueError("hidden_dim must be divisible by num_attention_heads.") - - self.state_dim = self.joint_dim + 7 + self.joint_dim - self.active_goal_dim = 7 + self.joint_dim + 3 + 1 - if self.use_relative_observations: - self.active_goal_dim += 7 + self.joint_dim - self.waypoint_token_dim = 3 + 4 + self.joint_dim + 3 + 1 + 1 - if self.use_relative_observations: - self.waypoint_token_dim += 7 + self.joint_dim - - self.action_token = nn.Parameter(torch.empty(1, 1, hidden_dim)) - nn.init.normal_(self.action_token, std=0.02) - self.token_type_embedding = nn.Embedding(_NUM_TOKEN_TYPES, hidden_dim) - nn.init.normal_(self.token_type_embedding.weight, std=0.02) - self.state_proj = _layer_init(nn.Linear(self.state_dim, hidden_dim)) - self.active_goal_proj = _layer_init(nn.Linear(self.active_goal_dim, hidden_dim)) - self.waypoint_proj = _layer_init(nn.Linear(self.waypoint_token_dim, hidden_dim)) - self.waypoint_index_embedding = nn.Parameter( - torch.empty(1, self.num_waypoints, hidden_dim) - ) - nn.init.normal_(self.waypoint_index_embedding, std=0.02) - self.waypoint_modality_embedding = nn.Embedding( - _NUM_WAYPOINT_TYPES, - hidden_dim, - ) - nn.init.normal_(self.waypoint_modality_embedding.weight, std=0.02) - - layer = nn.TransformerEncoderLayer( - d_model=hidden_dim, - nhead=num_attention_heads, - dim_feedforward=feedforward_dim or hidden_dim * 4, - dropout=0.0, - activation="gelu", - batch_first=True, - norm_first=False, - ) - self.encoder = nn.TransformerEncoder(layer, num_layers=num_layers) - - def _type_embedding( - self, - token_type: int, - batch_size: int, - device: torch.device, - ) -> torch.Tensor: - indices = torch.full( - (batch_size, 1), - token_type, - dtype=torch.long, - device=device, - ) - return self.token_type_embedding(indices) - - def encode_tokens(self, observation: torch.Tensor) -> torch.Tensor: - """Encode ``[ACTION, STATE, ACTIVE_GOAL, WP_1..WP_K]`` tokens. - - Args: - observation: Unified flat waypoint observation batch. - - Returns: - Encoded token tensor shaped ``[batch, 3 + K, hidden_dim]``. - """ - fields = parse_waypoint_observation( - observation, - self.num_waypoints, - self.use_relative_observations, - joint_dim=self.joint_dim, - ) - joint = fields["joint"] - assert isinstance(joint, torch.Tensor) - batch_size = observation.shape[0] - device = observation.device - action_tokens = self.action_token.expand( - batch_size, -1, -1 - ) + self._type_embedding( - _TOKEN_ACTION, - batch_size, - device, - ) - state_features = torch.cat( - [joint, fields["end_effector"], fields["last_action"]], - dim=-1, - ) - state_tokens = self.state_proj(state_features).unsqueeze( - 1 - ) + self._type_embedding(_TOKEN_STATE, batch_size, device) - - batch_indices = torch.arange(batch_size, device=device) - active_id = fields["active_id"] - assert isinstance(active_id, torch.Tensor) - active_parts = [ - fields["waypoint_position"][batch_indices, active_id], - fields["waypoint_quaternion"][batch_indices, active_id], - fields["waypoint_joint"][batch_indices, active_id], - fields["position_mask"][batch_indices, active_id].unsqueeze(-1), - fields["rotation_mask"][batch_indices, active_id].unsqueeze(-1), - fields["joint_mask"][batch_indices, active_id].unsqueeze(-1), - ] - if self.use_relative_observations: - active_parts.extend( - [ - fields["active_relative_pose"], - fields["waypoint_joint_error"][batch_indices, active_id], - ] - ) - progress = active_id.float().unsqueeze(-1) / max(self.num_waypoints - 1, 1) - active_parts.append(progress) - active_goal = torch.cat(active_parts, dim=-1) - active_goal_tokens = self.active_goal_proj(active_goal).unsqueeze( - 1 - ) + self._type_embedding(_TOKEN_ACTIVE_GOAL, batch_size, device) - - waypoint_features = [ - fields["waypoint_position"], - fields["waypoint_quaternion"], - fields["waypoint_joint"], - ] - if self.use_relative_observations: - waypoint_features.extend( - [ - fields["waypoint_relative_position"], - fields["waypoint_relative_quaternion"], - fields["waypoint_joint_error"], - ] - ) - waypoint_features.extend( - [ - fields["position_mask"].unsqueeze(-1), - fields["rotation_mask"].unsqueeze(-1), - fields["joint_mask"].unsqueeze(-1), - fields["active_onehot"].unsqueeze(-1), - fields["valid_mask"].unsqueeze(-1), - ] - ) - waypoint_token_types = self.token_type_embedding( - torch.full( - (batch_size, self.num_waypoints), - _TOKEN_WAYPOINT, - dtype=torch.long, - device=device, - ) - ) - waypoint_type = ( - fields["waypoint_type"] - .long() - .clamp( - 0, - _NUM_WAYPOINT_TYPES - 1, - ) - ) - waypoint_tokens = ( - self.waypoint_proj(torch.cat(waypoint_features, dim=-1)) - + waypoint_token_types - + self.waypoint_index_embedding - + self.waypoint_modality_embedding(waypoint_type) - ) - tokens = torch.cat( - [action_tokens, state_tokens, active_goal_tokens, waypoint_tokens], - dim=1, - ) - context_padding = torch.zeros( - batch_size, - 3, - dtype=torch.bool, - device=device, - ) - waypoint_padding = fields["valid_mask"] < 0.5 - padding_mask = torch.cat([context_padding, waypoint_padding], dim=1) - return self.encoder(tokens, src_key_padding_mask=padding_mask) - - -class WaypointTransformerActor(WaypointTransformerEncoder): - """Predict one joint action from fused action and active-goal tokens. - - Args: - observation_dim: Flat unified observation dimension. - action_dim: Output joint-action dimension. - num_waypoints: Maximum ordered waypoint count. - **kwargs: Encoder dimensions accepted by :class:`WaypointTransformerEncoder`. - """ - - def __init__( - self, - observation_dim: int, - action_dim: int, - num_waypoints: int, - **kwargs: object, - ) -> None: - super().__init__( - observation_dim=observation_dim, - num_waypoints=num_waypoints, - **kwargs, - ) - hidden_dim = int(kwargs.get("hidden_dim", 256)) - self.actor_head = _token_head( - hidden_dim, - action_dim, - std=0.01, - input_dim=2 * hidden_dim, - ) - - def forward(self, observation: torch.Tensor) -> torch.Tensor: - """Return deterministic mean actions. - - Args: - observation: Unified flat waypoint observation batch. - - Returns: - Mean action tensor shaped ``[batch, action_dim]``. - """ - encoded = self.encode_tokens(observation) - readout = torch.cat([encoded[:, 0], encoded[:, 2]], dim=-1) - return self.actor_head(readout) - - -class WaypointTransformerCritic(WaypointTransformerEncoder): - """Predict values from the full-context action token. - - Args: - observation_dim: Flat unified observation dimension. - num_waypoints: Maximum ordered waypoint count. - **kwargs: Encoder dimensions accepted by :class:`WaypointTransformerEncoder`. - """ - - def __init__( - self, - observation_dim: int, - num_waypoints: int, - **kwargs: object, - ) -> None: - super().__init__( - observation_dim=observation_dim, - num_waypoints=num_waypoints, - **kwargs, - ) - hidden_dim = int(kwargs.get("hidden_dim", 256)) - self.value_head = _token_head(hidden_dim, 1, std=1.0) - - def forward(self, observation: torch.Tensor) -> torch.Tensor: - """Return one value per observation. - - Args: - observation: Unified flat waypoint observation batch. - - Returns: - Value tensor shaped ``[batch, 1]``. - """ - return self.value_head(self.encode_tokens(observation)[:, 0]) diff --git a/embodichain/learning/rl/train.py b/embodichain/learning/rl/train.py index 0dba17c58..612df5691 100644 --- a/embodichain/learning/rl/train.py +++ b/embodichain/learning/rl/train.py @@ -28,11 +28,8 @@ from torch.utils.tensorboard import SummaryWriter from copy import deepcopy -from embodichain.learning.rl.models import ( - build_model_from_cfg, - build_policy, - get_registered_policy_names, -) +from embodichain.learning.rl.models import build_policy, get_registered_policy_names +from embodichain.learning.rl.models import build_mlp_from_cfg from embodichain.learning.rl.algo import ( RolloutKind, build_algo, @@ -128,14 +125,12 @@ def _build_learning_policy( actor_cfg = policy_block.get("actor") critic_cfg = policy_block.get("critic") actor = ( - build_model_from_cfg(actor_cfg, obs_dim, action_dim, role="actor") + build_mlp_from_cfg(actor_cfg, obs_dim, action_dim) if actor_cfg is not None else None ) critic = ( - build_model_from_cfg(critic_cfg, obs_dim, 1, role="critic") - if critic_cfg is not None - else None + build_mlp_from_cfg(critic_cfg, obs_dim, 1) if critic_cfg is not None else None ) policy = build_policy( policy_block, diff --git a/tests/learning/test_waypoint_transformer.py b/tests/learning/test_waypoint_transformer.py deleted file mode 100644 index fe84dac18..000000000 --- a/tests/learning/test_waypoint_transformer.py +++ /dev/null @@ -1,145 +0,0 @@ -# ---------------------------------------------------------------------------- -# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ---------------------------------------------------------------------------- - -"""Tests for unified Cartesian/joint waypoint Transformer models.""" - -from __future__ import annotations - -import torch - -from embodichain.learning.rl.models import ( - WaypointTransformerActor, - WaypointTransformerCritic, - build_model_from_cfg, - parse_waypoint_observation, - waypoint_observation_dim, - waypoint_observation_normalize_mask, -) - -_NUM_WAYPOINTS = 4 -_OBSERVATION_DIM = waypoint_observation_dim(_NUM_WAYPOINTS, True) - - -def _valid_observation(batch_size: int = 3) -> torch.Tensor: - observation = torch.zeros(batch_size, _OBSERVATION_DIM) - cursor = 7 + 7 + _NUM_WAYPOINTS * (3 + 4 + 7) - observation[:, cursor] = 1.0 - cursor += _NUM_WAYPOINTS - observation[:, cursor : cursor + _NUM_WAYPOINTS] = 1.0 - cursor += _NUM_WAYPOINTS - observation[:, cursor : cursor + _NUM_WAYPOINTS] = 1.0 - cursor += _NUM_WAYPOINTS - observation[:, cursor : cursor + _NUM_WAYPOINTS] = 1.0 - return observation - - -def _actor() -> WaypointTransformerActor: - return WaypointTransformerActor( - observation_dim=_OBSERVATION_DIM, - action_dim=7, - num_waypoints=_NUM_WAYPOINTS, - hidden_dim=64, - num_attention_heads=4, - num_layers=1, - ) - - -def test_waypoint_actor_and_critic_have_expected_shapes_and_gradients() -> None: - torch.manual_seed(7) - actor = _actor().train() - critic = WaypointTransformerCritic( - observation_dim=_OBSERVATION_DIM, - num_waypoints=_NUM_WAYPOINTS, - hidden_dim=64, - num_attention_heads=4, - num_layers=1, - ).train() - observation = _valid_observation() - - action = actor(observation) - value = critic(observation) - (action.square().mean() + value.square().mean()).backward() - - assert action.shape == (3, 7) - assert value.shape == (3, 1) - assert torch.isfinite(action).all() - assert torch.isfinite(value).all() - assert actor.action_token.grad is not None - assert critic.action_token.grad is not None - - -def test_waypoint_actor_uses_future_valid_waypoints() -> None: - torch.manual_seed(11) - actor = _actor().eval() - base = _valid_observation(batch_size=1) - changed = base.clone() - # WP3 position is a future token while WP0 remains active. - position_start = 7 + 7 - changed[:, position_start + 3 * 3 : position_start + 3 * 4] = 5.0 - - with torch.no_grad(): - base_action = actor(base) - changed_action = actor(changed) - - assert not torch.allclose(base_action, changed_action) - - -def test_waypoint_parser_clamps_active_id_to_last_valid_slot() -> None: - observation = _valid_observation(batch_size=1) - cursor = 7 + 7 + _NUM_WAYPOINTS * (3 + 4 + 7) - observation[:, cursor : cursor + _NUM_WAYPOINTS] = 0.0 - observation[:, cursor + 3] = 1.0 - cursor += _NUM_WAYPOINTS - observation[:, cursor : cursor + _NUM_WAYPOINTS] = 0.0 - observation[:, cursor : cursor + 2] = 1.0 - - fields = parse_waypoint_observation(observation, _NUM_WAYPOINTS, True) - - assert torch.equal(fields["active_id"], torch.tensor([1])) - - -def test_waypoint_normalization_mask_excludes_all_semantic_fields() -> None: - mask = waypoint_observation_normalize_mask(_NUM_WAYPOINTS, True) - - assert mask.shape == (_OBSERVATION_DIM,) - assert int((~mask).sum()) == 6 * _NUM_WAYPOINTS + 7 * _NUM_WAYPOINTS - - -def test_waypoint_transformer_builds_from_learning_config() -> None: - module_config = { - "type": "waypoint_transformer", - "network_cfg": { - "num_waypoints": _NUM_WAYPOINTS, - "hidden_dim": 64, - "num_attention_heads": 4, - "num_layers": 1, - }, - } - actor = build_model_from_cfg( - module_config, - _OBSERVATION_DIM, - 7, - role="actor", - ) - critic = build_model_from_cfg( - module_config, - _OBSERVATION_DIM, - 1, - role="critic", - ) - - assert isinstance(actor, WaypointTransformerActor) - assert isinstance(critic, WaypointTransformerCritic)