Add PointMass task and configurable optimizer for APG/PPO

.agents/skills/add-task-env/SKILL.md

@@ -18,14 +18,17 @@ Scaffold a new task environment following EmbodiChain's conventions and patterns
 
 Ask the user:
 
-- **Category**: `tableware`, `rl`, or `special` (maps to `embodichain/lab/gym/envs/tasks/<category>/`)
+- **Category**: `tableware`, `rl`, or `special` (maps to `embodichain_tasks/embodichain_tasks/<category>/`)
 - **Task name** (snake_case, e.g. `pick_place`)
 - **Gym ID** (e.g. `PickPlace-v1`)
 - **Task type**: RL task (needs reward functors) or expert demonstration task (needs `create_demo_action_list`)
 
 ### 2. Create the Task File
 
-Place at `embodichain/lab/gym/envs/tasks/<category>/<name>.py`.
+Place at `embodichain_tasks/embodichain_tasks/<category>/<name>.py`.
+Lightweight pure-PyTorch RL tasks (e.g. PointMass) also live under
+`embodichain_tasks/embodichain_tasks/rl/basic/` and register through
+`@register_learning_env` when they are not `EmbodiedEnv` subclasses.
 
 Template:
 
@@ -76,32 +79,38 @@ class <CamelCaseName>Env(EmbodiedEnv):
 
 ### 3. Update Exports
 
-Add to `embodichain/lab/gym/envs/tasks/__init__.py`:
+Task packages under `embodichain_tasks` are auto-imported via
+`import_packages()`. Prefer an explicit export in the category
+`__init__.py`:
 
 ```python
-from embodichain.lab.gym.envs.tasks.<category>.<name> import <CamelCaseName>Env
+from .<name> import <CamelCaseName>Env
+
+__all__ = [..., "<CamelCaseName>Env"]
 ```
 
-Add `"<CamelCaseName>Env"` to the `__all__` list.
+Optional compatibility re-export may also be added in
+`embodichain/lab/gym/envs/tasks/__init__.py`.
 
 ### 4. Create Test Stub
 
-Place at `tests/gym/envs/tasks/test_<name>.py`.
+Place at `tests/gym/envs/tasks/test_<name>.py` (or `tests/learning/` for
+lightweight learning environments).
 
 ### 5. Format
 
 ```bash
-black embodichain/lab/gym/envs/tasks/<category>/<name>.py
+black embodichain_tasks/embodichain_tasks/<category>/<name>.py
 black tests/gym/envs/tasks/test_<name>.py
 ```
 
 ## Checklist
 
 - [ ] File has Apache 2.0 header
 - [ ] Uses `from __future__ import annotations`
-- [ ] `@register_env` decorator with unique gym ID
+- [ ] `@register_env` or `@register_learning_env` with a unique ID
 - [ ] `__all__` defined in the task module
-- [ ] Default `cfg = EmbodiedEnvCfg()` in `__init__`
-- [ ] Import and `__all__` added to `tasks/__init__.py`
+- [ ] Default `cfg = EmbodiedEnvCfg()` in `__init__` for EmbodiedEnv tasks
+- [ ] Category `__init__.py` export updated
 - [ ] Test stub created
 - [ ] `black` run on both files

agent_context/MAP.yaml

@@ -341,13 +341,16 @@ topics:
       - topics/rl-learning/rl-learning.md
     source_of_truth:
       - embodichain/learning/rl/env.py
+      - embodichain/learning/rl/evaluation.py
+      - embodichain/learning/rl/routing.py
       - embodichain/learning/rl/differentiable_trainer.py
       - embodichain/learning/rl/experimental/newton/
       - embodichain/learning/rl/train.py
       - embodichain/learning/rl/algo/
       - embodichain/learning/rl/buffer/
       - embodichain/learning/rl/models/
       - embodichain/learning/rl/collector/
+      - embodichain_tasks/embodichain_tasks/rl/basic/point_mass.py
     related_topics:
       - env-framework
       - manager-functor

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

@@ -18,7 +18,9 @@ python -m embodichain.learning.rl.train --config <path-to-yaml-or-json> [--distr
 The RL subsystem implements on-policy reinforcement learning with a modular pipeline:
 
 1. **Config** — JSON/YAML file defines `trainer`, `policy`, and `algorithm` blocks.
-2. **Environment** — Built via `build_env()` from `embodichain.lab.gym.envs.tasks.rl`.
+2. **Environment** — Simulator tasks use `build_env()`; lightweight tensor
+   tasks use the `learning_env` registry. Official tasks live in
+   `embodichain_tasks`.
 3. **Policy** — Neural-network module (`Policy` ABC) producing actions from observations.
 4. **Collector** — Steps the env, writes transitions into a preallocated `TensorDict`.
 5. **Buffer** — `RolloutBuffer` owns the preallocated storage; marks it full after collection.
@@ -48,17 +50,25 @@ Standard PPO/GRPO rollout data flows as `TensorDict` objects (from the
   update. Segment losses backpropagate immediately, state is detached after
   each segment, and policy gradients accumulate until one update horizon is
   complete.
-- Checkpoints store policy, optimizer, and trainer counters.
+- Checkpoints store policy, optimizer, optional LR scheduler, trainer counters, and best-evaluation
+  state.
 - Training selects `policy.train()`, while evaluation temporarily selects
   `policy.eval()` and restores the previous mode afterward.
 
 The collector paths are intentionally separate:
 - PPO/GRPO use `SyncCollector` and `TensorDict`.
 - APG uses `DifferentiableCollector` and `DifferentiableRollout`.
 
-APG is exported but intentionally absent from `_ALGO_REGISTRY`; the standard
-`train.py` and `Trainer` cannot consume differentiable rollouts yet. Construct
-APG directly with `DifferentiableTrainer`.
+APG is registered with `RolloutKind.DIFFERENTIABLE`. The shared `train.py`
+entry point routes it to `DifferentiableTrainer`; PPO/GRPO declare
+`RolloutKind.STANDARD` and continue to use `Trainer`. Collectors and rollout
+types remain separate even though environment construction, policy building,
+evaluation, logging, and checkpoint conventions are shared.
+
+`PointMassRL` in
+`embodichain_tasks/embodichain_tasks/rl/basic/point_mass.py` is the canonical
+dual-path task. The same differentiable PyTorch dynamics run under
+`torch.no_grad()` for PPO and preserve the action-to-reward graph for APG.
 
 ### Newton Reference
 
@@ -85,10 +95,20 @@ gradients, TBPTT, APG updates, and held-out improvement.
 
 ```
 train_from_config()
-  ├─ build_env()                     → Gym env
+  ├─ build_env() | build_learning_env()
   ├─ build_policy(policy_block, ...) → Policy (ActorCritic | ActorOnly | custom)
-  ├─ build_algo(name, cfg, policy)   → Algorithm (PPO | GRPO)
-  └─ Trainer(policy, env, algorithm, ...)
+  ├─ build_algo(name, cfg, policy)   → Algorithm (PPO | GRPO | APG)
+  └─ route by RolloutKind
+       ├─ STANDARD       → Trainer + SyncCollector + RolloutBuffer
+       └─ DIFFERENTIABLE → DifferentiableTrainer + DifferentiableCollector
+```
+
+Both trainers call `evaluation.evaluate_episodes()` with an independent
+evaluation environment. It counts actual completed episodes under asynchronous
+auto-reset and logs terminal-only metrics as `eval/*`.
+
+```text
+Trainer(policy, env, algorithm, ...)
        ├─ RolloutBuffer  [buffer/standard_buffer.py]
        ├─ SyncCollector  [collector/sync_collector.py]
        └─ .train(total_timesteps)
@@ -105,7 +125,7 @@ train_from_config()
 **Source**: `embodichain/learning/rl/algo/ppo.py`
 
 - Config: `PPOCfg(AlgorithmCfg)` — `n_epochs=10`, `clip_coef=0.2`, `ent_coef=0.01`, `vf_coef=0.5`.
-- Inherits `AlgorithmCfg` defaults: `lr=3e-4`, `batch_size=64`, `gamma=0.99`, `gae_lambda=0.95`, `max_grad_norm=0.5`.
+- Inherits `AlgorithmCfg` defaults: nested `optimizer` (`adam`, `lr=3e-4`), optional `lr_scheduler`, `batch_size=64`, `gamma=0.99`, `gae_lambda=0.95`, `max_grad_norm=0.5`.
 - `update(rollout)` flow:
   1. `compute_gae(rollout, gamma, gae_lambda)` — writes `advantage` and `return` into the TensorDict.
   2. `transition_view(rollout, flatten=True)` — drops padded final slot, flattens to `[N*T]`.
@@ -127,11 +147,18 @@ train_from_config()
 **Source**: `embodichain/learning/rl/algo/__init__.py`
 
 ```python
-_ALGO_REGISTRY = {"ppo": (PPOCfg, PPO), "grpo": (GRPOCfg, GRPO)}
+_ALGO_REGISTRY = {
+    "apg": (APGCfg, APG),
+    "ppo": (PPOCfg, PPO),
+    "grpo": (GRPOCfg, GRPO),
+}
 build_algo(name, cfg_kwargs, policy, device, distributed=False)
 ```
 
-When `distributed=True`, wraps the policy in `DistributedDataParallel` before passing to the algorithm.
+`APG.rollout_kind` is `RolloutKind.DIFFERENTIABLE`; PPO/GRPO use
+`RolloutKind.STANDARD`. When `distributed=True`, wraps the policy in
+`DistributedDataParallel` before passing to the algorithm and rejects
+differentiable algorithms.
 
 ## Rollout Buffer
 
@@ -233,3 +260,6 @@ Distributed training:
 | `GRPO: group_size >= 2` | GRPO requires at least 2 environments per group for normalization. |
 | NaN losses | Check `log_std` bounds, gradient clipping, and reward scale. `max_grad_norm` defaults to 0.5. |
 | Stale observations after reset | `SyncCollector` resets obs via `_reset_env()` on init; set `reset_every_rollout=True` if episodes must fully reset between rollouts. |
+| `Learning environment 'X' is not registered` | Task package was not imported. Call `discover_task_packages()` / `execute_init_hooks()`, or import the task module that uses `@register_learning_env`. |
+| `Differentiable algorithms require trainer.learning_env` | APG cannot train simulator `gym_config` envs; use a registered `learning_env` such as `PointMassRL`. |
+| Eval success rate stuck near zero | Confirm `evaluate_episodes` reads terminal-only `success`/`metrics`, and that the eval env uses a held-out `eval_seed`. |

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

@@ -7,16 +7,20 @@ Overview
 --------
 
 Algorithm registry and algorithm-construction helpers for RL training. The
-on-policy algorithms :class:`PPO` and :class:`GRPO` both derive from
-:class:`BaseAlgorithm`; :func:`build_algo` looks up a registered algorithm by
-name and wires it to a policy, while :func:`compute_gae` provides generalized
-advantage estimation.
+on-policy algorithms :class:`PPO` and :class:`GRPO` and the differentiable
+algorithm :class:`APG` all derive from :class:`BaseAlgorithm`. Each algorithm
+declares a :class:`RolloutKind` that selects the compatible trainer.
+:func:`build_algo` looks up a registered algorithm by name and wires it to a
+policy, while :func:`compute_gae` provides generalized advantage estimation.
 
    .. rubric:: Classes
 
    .. autosummary::
 
       BaseAlgorithm
+      RolloutKind
+      APGCfg
+      APG
       PPOCfg
       PPO
       GRPOCfg
@@ -29,9 +33,9 @@ advantage estimation.
       build_algo
       get_registered_algo_names
       compute_gae
+      segmented_discounted_return
 
 .. automodule:: embodichain.learning.rl.algo
    :members:
    :undoc-members:
    :show-inheritance:
-

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

@@ -21,6 +21,19 @@ collection logic, policy/model builders, and training entry points.
       train
       utils
 
+   .. rubric:: Top-level APIs
+
+   .. autosummary::
+
+      DifferentiableTrainer
+      DifferentiableTrainerCfg
+      DifferentiableVecEnv
+      LearningVecEnv
+      build_learning_env
+      evaluate_episodes
+      get_trainer_class
+      register_learning_env
+
 Algorithms
 ----------
 
@@ -29,6 +42,38 @@ Algorithms
    :undoc-members:
    :show-inheritance:
 
+Environments
+------------
+
+.. automodule:: embodichain.learning.rl.env
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+Evaluation
+----------
+
+.. automodule:: embodichain.learning.rl.evaluation
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+Routing
+-------
+
+.. automodule:: embodichain.learning.rl.routing
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+Differentiable Trainer
+----------------------
+
+.. automodule:: embodichain.learning.rl.differentiable_trainer
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
 Rollout Buffer
 --------------
 

docs/source/overview/rl/train_script.md

@@ -7,8 +7,11 @@ This module provides the RL training entry script, responsible for parsing confi
 ### train.py
 - Main training script, supports command-line arguments (such as --config), automatically loads JSON or YAML config.
 - Initializes device, random seed, output directory, and logging (TensorBoard/WandB).
-- Loads environment config, supports multi-environment parallelism and evaluation environments.
-- Builds policy (e.g., actor-critic), algorithm (e.g., PPO), and Trainer.
+- Loads either a simulator `gym_config` or a lightweight registered
+  `learning_env`, with independent vectorized evaluation environments.
+- Builds the policy and algorithm, then routes standard rollouts
+  (PPO/GRPO) to `Trainer` and differentiable rollouts (APG) to
+  `DifferentiableTrainer`.
 - Supports event management (e.g., environment randomization, data logging, evaluation events).
 - Automatically saves model checkpoints and performs periodic evaluation.
 
@@ -35,8 +38,19 @@ This module provides the RL training entry script, responsible for parsing confi
 ## Usage Example
 ```bash
 embodichain train-rl --config embodichain_tasks/configs/agents/rl/basic/cart_pole/train_config.yaml
+embodichain train-rl --config embodichain_tasks/configs/agents/rl/basic/point_mass/train_apg.yaml
+embodichain train-rl --config embodichain_tasks/configs/agents/rl/basic/point_mass/train_ppo.yaml
 ```
 
+PointMass intentionally uses one differentiable PyTorch environment for both
+algorithms. The standard collector runs it under `torch.no_grad()`, whereas
+APG retains the dynamics graph and detaches it only at configured TBPTT
+segment boundaries.
+
+Evaluation always uses a separate environment and deterministic actions.
+Metrics are logged under `eval/*`; `num_eval_episodes` is the exact number of
+completed episodes rather than episodes per parallel environment.
+
 ## Extension and Customization
 - Supports custom event modules for flexible training flow extension.
 - Can integrate multi-task and multi-environment training.

docs/source/overview/rl/trainer.md

@@ -22,9 +22,23 @@ This module implements the main RL training loop, logging management, and event-
 - `train(total_timesteps)`: Main training loop, automatically collects data, updates policy, and logs.
 - `_collect_rollout()`: Collect one rollout through `SyncCollector`, supports custom callback statistics.
 - `_log_train(losses)`: Log training loss, reward, sampling speed, etc.
-- `_eval_once()`: Periodic evaluation, records evaluation metrics.
+- `_eval_once()`: Calls the shared evaluator and records exact completed
+  episode metrics.
 - `save_checkpoint()`: Save model parameters and training state.
 
+### DifferentiableTrainer
+- Coordinates APG rollouts without copying graph-connected transitions into
+  the standard preallocated buffer.
+- Accumulates pathwise gradients across `segment_length` TBPTT segments until
+  `update_horizon` is reached.
+- Uses the same episode statistics, deterministic evaluator, TensorBoard/W&B
+  namespaces, best-checkpoint selection, and global-step semantics as
+  `Trainer`.
+
+The trainers remain separate because their rollout ownership and autograd
+requirements differ. `train-rl` selects between them through each algorithm's
+`RolloutKind`.
+
 ## Event Management
 - Supports custom events (e.g., environment randomization, data logging) injected via EventManager.
 - Events can be executed by interval/step/trigger, enabling flexible extension.

docs/source/tutorial/rl.rst

@@ -37,12 +37,20 @@ The core components and their relationships:
 Configuration via JSON or YAML
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
-Training is configured via a JSON or YAML file that defines runtime settings, environment, policy, and algorithm parameters. EmbodiChain loads either format with ``load_config()``; the nested ``trainer.gym_config`` path supports the same extensions.
+Training is configured via a JSON or YAML file that defines runtime settings, environment, policy, and algorithm parameters. EmbodiChain loads either format with ``load_config()``. Two environment sources are supported:
+
+- ``trainer.gym_config``: simulator ``EmbodiedEnv`` tasks such as CartPole and PushCube
+- ``trainer.learning_env``: lightweight registered tensor environments such as PointMass
+
+``algorithm.name`` selects the update rule. Standard algorithms (PPO/GRPO) use
+``Trainer`` and ``SyncCollector``. Differentiable algorithms (APG) use
+``DifferentiableTrainer`` and ``DifferentiableCollector``. Both paths share
+``evaluate_episodes()`` for deterministic evaluation metrics under ``eval/*``.
 
 Example Configuration
 ---------------------   
 
-The configuration file (e.g., ``train_config.json`` or ``train_config.yaml``) is located in ``embodichain_tasks/configs/agents/rl/push_cube`` or ``embodichain_tasks/configs/agents/rl/basic/cart_pole``:
+The configuration file (e.g., ``train_config.json`` or ``train_config.yaml``) is located in ``embodichain_tasks/configs/agents/rl/push_cube`` or ``embodichain_tasks/configs/agents/rl/basic/cart_pole``. PointMass APG/PPO configs live under ``embodichain_tasks/configs/agents/rl/basic/point_mass``:
 
 .. dropdown:: Example: train_config.json
    :icon: code
@@ -58,6 +66,33 @@ The configuration file (e.g., ``train_config.json`` or ``train_config.yaml``) is
       :language: yaml
       :linenos:
 
+.. dropdown:: Example: train_apg.yaml (PointMass)
+   :icon: code
+
+   .. literalinclude:: ../../../embodichain_tasks/configs/agents/rl/basic/point_mass/train_apg.yaml
+      :language: yaml
+      :linenos:
+
+PointMass uses one differentiable PyTorch dynamics implementation for both
+APG and PPO. Launch either config with the same CLI:
+
+.. code-block:: bash
+
+   embodichain train-rl --config embodichain_tasks/configs/agents/rl/basic/point_mass/train_apg.yaml
+   embodichain train-rl --config embodichain_tasks/configs/agents/rl/basic/point_mass/train_ppo.yaml
+
+After training, plot deterministic 2D trajectories from a checkpoint:
+
+.. code-block:: bash
+
+   python scripts/tutorials/rl/visualize_point_mass.py \
+       --config embodichain_tasks/configs/agents/rl/basic/point_mass/train_apg.yaml \
+       --checkpoint outputs/<exp_name>/checkpoints/<exp_name>_best.pt \
+       --tag apg_best \
+       --num-episodes 4
+
+Plots are written under ``outputs/point_mass_viz/`` by default.
+
 Configuration Sections
 ---------------------