From 9d9a14f7d1f3035b62f382dc5a2b3108dd571b17 Mon Sep 17 00:00:00 2001
From: yangchen73 <122090643@link.cuhk.edu.cn>
Date: Sun, 2 Aug 2026 01:57:24 +0800
Subject: [PATCH 1/7] Add PointMass RL task with shared APG/PPO training path

---
 .agents/skills/add-task-env/SKILL.md          |  29 +-
 agent_context/MAP.yaml                        |   3 +
 .../topics/rl-learning/rl-learning.md         |  50 ++-
 .../embodichain.learning.rl.algo.rst          |  14 +-
 .../embodichain/embodichain.learning.rl.rst   |  45 +++
 docs/source/overview/rl/train_script.md       |  18 +-
 docs/source/overview/rl/trainer.md            |  16 +-
 docs/source/tutorial/rl.rst                   |  27 +-
 embodichain/learning/rl/__init__.py           |  17 +-
 embodichain/learning/rl/algo/__init__.py      |  15 +-
 embodichain/learning/rl/algo/apg.py           |   4 +-
 embodichain/learning/rl/algo/base.py          |  11 +-
 .../learning/rl/differentiable_trainer.py     | 135 +++++++-
 embodichain/learning/rl/env.py                | 101 ++++--
 embodichain/learning/rl/evaluation.py         | 142 +++++++++
 embodichain/learning/rl/models/__init__.py    |  16 +-
 embodichain/learning/rl/routing.py            |  39 +++
 embodichain/learning/rl/train.py              | 208 +++++++++++-
 embodichain/learning/rl/utils/helper.py       |  12 +-
 embodichain/learning/rl/utils/trainer.py      | 252 ++++++---------
 .../agents/rl/basic/point_mass/train_apg.yaml |  42 +++
 .../agents/rl/basic/point_mass/train_ppo.yaml |  52 +++
 .../embodichain_tasks/rl/basic/__init__.py    |  24 ++
 .../embodichain_tasks/rl/basic/point_mass.py  | 296 ++++++++++++++++++
 scripts/benchmark/rl/algorithms/apg.yaml      |  19 ++
 scripts/benchmark/rl/runtime.py               | 280 +++++++++++------
 scripts/benchmark/rl/suites/point_mass.yaml   |  21 ++
 scripts/benchmark/rl/tasks/point_mass.yaml    |  24 ++
 tests/learning/test_apg.py                    |   9 +-
 tests/learning/test_differentiable_env.py     |   3 +
 tests/learning/test_differentiable_trainer.py |   3 +
 tests/learning/test_evaluation.py             |  90 ++++++
 tests/learning/test_point_mass.py             | 226 +++++++++++++
 tests/learning/test_routing.py                |  48 +++
 34 files changed, 1966 insertions(+), 325 deletions(-)
 create mode 100644 embodichain/learning/rl/evaluation.py
 create mode 100644 embodichain/learning/rl/routing.py
 create mode 100644 embodichain_tasks/configs/agents/rl/basic/point_mass/train_apg.yaml
 create mode 100644 embodichain_tasks/configs/agents/rl/basic/point_mass/train_ppo.yaml
 create mode 100644 embodichain_tasks/embodichain_tasks/rl/basic/point_mass.py
 create mode 100644 scripts/benchmark/rl/algorithms/apg.yaml
 create mode 100644 scripts/benchmark/rl/suites/point_mass.yaml
 create mode 100644 scripts/benchmark/rl/tasks/point_mass.yaml
 create mode 100644 tests/learning/test_evaluation.py
 create mode 100644 tests/learning/test_point_mass.py
 create mode 100644 tests/learning/test_routing.py

diff --git a/.agents/skills/add-task-env/SKILL.md b/.agents/skills/add-task-env/SKILL.md
index b6092cfc9..d063fa460 100644
--- a/.agents/skills/add-task-env/SKILL.md
+++ b/.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,22 +79,28 @@ 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
 ```
 
@@ -99,9 +108,9 @@ black tests/gym/envs/tasks/test_<name>.py
 
 - [ ] 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
diff --git a/agent_context/MAP.yaml b/agent_context/MAP.yaml
index 8389f34f7..65e5aa521 100644
--- a/agent_context/MAP.yaml
+++ b/agent_context/MAP.yaml
@@ -341,6 +341,8 @@ 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
@@ -348,6 +350,7 @@ topics:
       - 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
diff --git a/agent_context/topics/rl-learning/rl-learning.md b/agent_context/topics/rl-learning/rl-learning.md
index f892c916c..65cc28f85 100644
--- a/agent_context/topics/rl-learning/rl-learning.md
+++ b/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,7 +50,8 @@ 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, trainer counters, and best-evaluation
+  state.
 - Training selects `policy.train()`, while evaluation temporarily selects
   `policy.eval()` and restores the previous mode afterward.
 
@@ -56,9 +59,16 @@ 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)
@@ -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`. |
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 1dbec5f66..0db256501 100644
--- a/docs/source/api_reference/embodichain/embodichain.learning.rl.algo.rst
+++ b/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:
-   
\ 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 675f004a6..bf1b5b01e 100644
--- a/docs/source/api_reference/embodichain/embodichain.learning.rl.rst
+++ b/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
 --------------
 
diff --git a/docs/source/overview/rl/train_script.md b/docs/source/overview/rl/train_script.md
index 930e407f9..f39b8b69d 100644
--- a/docs/source/overview/rl/train_script.md
+++ b/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.
diff --git a/docs/source/overview/rl/trainer.md b/docs/source/overview/rl/trainer.md
index 5ef4ee995..433d278e9 100644
--- a/docs/source/overview/rl/trainer.md
+++ b/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.
diff --git a/docs/source/tutorial/rl.rst b/docs/source/tutorial/rl.rst
index a14dac3b4..54ce4590f 100644
--- a/docs/source/tutorial/rl.rst
+++ b/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,21 @@ 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
+
 Configuration Sections
 ---------------------
 
diff --git a/embodichain/learning/rl/__init__.py b/embodichain/learning/rl/__init__.py
index 528ba7e91..9c97766b1 100644
--- a/embodichain/learning/rl/__init__.py
+++ b/embodichain/learning/rl/__init__.py
@@ -27,13 +27,28 @@
     DifferentiableTrainer,
     DifferentiableTrainerCfg,
 )
-from .env import DifferentiableObservation, DifferentiableVecEnv
+from .env import (
+    DifferentiableObservation,
+    DifferentiableVecEnv,
+    LearningVecEnv,
+    build_learning_env,
+    get_registered_learning_env_names,
+    register_learning_env,
+)
+from .evaluation import evaluate_episodes
+from .routing import get_trainer_class
 
 __all__ = [
     "DifferentiableObservation",
     "DifferentiableTrainer",
     "DifferentiableTrainerCfg",
     "DifferentiableVecEnv",
+    "LearningVecEnv",
+    "build_learning_env",
+    "evaluate_episodes",
+    "get_registered_learning_env_names",
+    "get_trainer_class",
+    "register_learning_env",
     "algo",
     "buffer",
     "models",
diff --git a/embodichain/learning/rl/algo/__init__.py b/embodichain/learning/rl/algo/__init__.py
index 305cf7944..5fb312ea7 100644
--- a/embodichain/learning/rl/algo/__init__.py
+++ b/embodichain/learning/rl/algo/__init__.py
@@ -24,12 +24,13 @@
 from torch.nn.parallel import DistributedDataParallel as DDP
 
 from .apg import APG, APGCfg, segmented_discounted_return
-from .base import BaseAlgorithm
+from .base import BaseAlgorithm, RolloutKind
 from .common import compute_gae
 from .grpo import GRPO, GRPOCfg
 from .ppo import PPO, PPOCfg
 
 _ALGO_REGISTRY: Dict[str, Tuple[Type[Any], Type[Any]]] = {
+    "apg": (APGCfg, APG),
     "ppo": (PPOCfg, PPO),
     "grpo": (GRPOCfg, GRPO),
 }
@@ -48,13 +49,6 @@ def build_algo(
     distributed: bool = False,
 ):
     key = name.lower()
-    if key == "apg":
-        raise ValueError(
-            "APG uses differentiable rollouts and is not supported by the "
-            "standard build_algo()/train.py path. Construct APG with "
-            "DifferentiableTrainer directly, or use the experimental Newton "
-            "training entry point."
-        )
     if key not in _ALGO_REGISTRY:
         raise ValueError(
             f"Algorithm '{name}' not found. Available: {get_registered_algo_names()}"
@@ -62,6 +56,10 @@ def build_algo(
     CfgCls, AlgoCls = _ALGO_REGISTRY[key]
     cfg = CfgCls(device=str(device), **cfg_kwargs)
     if distributed:
+        if AlgoCls.rollout_kind is RolloutKind.DIFFERENTIABLE:
+            raise ValueError(
+                "Differentiable algorithms do not support distributed training."
+            )
         if not (
             torch.distributed.is_available() and torch.distributed.is_initialized()
         ):
@@ -79,6 +77,7 @@ def build_algo(
 
 __all__ = [
     "BaseAlgorithm",
+    "RolloutKind",
     "APGCfg",
     "APG",
     "segmented_discounted_return",
diff --git a/embodichain/learning/rl/algo/apg.py b/embodichain/learning/rl/algo/apg.py
index 144f52266..76f1f33d3 100644
--- a/embodichain/learning/rl/algo/apg.py
+++ b/embodichain/learning/rl/algo/apg.py
@@ -26,7 +26,7 @@
 from embodichain.learning.rl.utils import AlgorithmCfg
 from embodichain.utils import configclass
 
-from .base import BaseAlgorithm
+from .base import BaseAlgorithm, RolloutKind
 
 __all__ = ["APG", "APGCfg", "segmented_discounted_return"]
 
@@ -66,6 +66,8 @@ class APGCfg(AlgorithmCfg):
 class APG(BaseAlgorithm[DifferentiableRollout]):
     """Optimize policy parameters through differentiable rollout rewards."""
 
+    rollout_kind = RolloutKind.DIFFERENTIABLE
+
     def __init__(self, cfg: APGCfg, policy: torch.nn.Module) -> None:
         self.cfg = cfg
         self.policy = policy
diff --git a/embodichain/learning/rl/algo/base.py b/embodichain/learning/rl/algo/base.py
index 8567668c7..57235fbdc 100644
--- a/embodichain/learning/rl/algo/base.py
+++ b/embodichain/learning/rl/algo/base.py
@@ -17,15 +17,23 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
+from enum import Enum
 from typing import Dict, Generic, TypeVar
 
 import torch
 
-__all__ = ["BaseAlgorithm"]
+__all__ = ["BaseAlgorithm", "RolloutKind"]
 
 RolloutT = TypeVar("RolloutT")
 
 
+class RolloutKind(str, Enum):
+    """Rollout semantics required by an algorithm."""
+
+    STANDARD = "standard"
+    DIFFERENTIABLE = "differentiable"
+
+
 class BaseAlgorithm(ABC, Generic[RolloutT]):
     """Base class for RL algorithms.
 
@@ -33,6 +41,7 @@ class BaseAlgorithm(ABC, Generic[RolloutT]):
     """
 
     device: torch.device
+    rollout_kind = RolloutKind.STANDARD
 
     @abstractmethod
     def update(self, rollout: RolloutT) -> Dict[str, float]:
diff --git a/embodichain/learning/rl/differentiable_trainer.py b/embodichain/learning/rl/differentiable_trainer.py
index b07faccf7..46695c930 100644
--- a/embodichain/learning/rl/differentiable_trainer.py
+++ b/embodichain/learning/rl/differentiable_trainer.py
@@ -19,14 +19,18 @@
 from __future__ import annotations
 
 import math
+import time
+from collections import deque
 from pathlib import Path
 from typing import TYPE_CHECKING, Any
 
 import torch
+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.evaluation import evaluate_episodes
 from embodichain.learning.rl.models import Policy
 from embodichain.utils import configclass
 
@@ -48,6 +52,12 @@ class DifferentiableTrainerCfg:
     checkpoint_dir: str = "outputs/checkpoints"
     experiment_name: str = "apg"
     save_frequency_updates: int = 0
+    eval_frequency_steps: int = 0
+    num_eval_episodes: int = 5
+    eval_seed: int | None = None
+    use_wandb: bool = False
+    best_eval_metric: str = "eval/avg_reward"
+    best_eval_mode: str = "max"
 
 
 class DifferentiableTrainer:
@@ -60,6 +70,7 @@ def __init__(
         policy: Policy,
         algorithm: APG,
         writer: SummaryWriter | None = None,
+        eval_env: DifferentiableVecEnv | None = None,
     ) -> None:
         if cfg.segment_length <= 0:
             raise ValueError("segment_length must be positive.")
@@ -72,6 +83,10 @@ def __init__(
             raise ValueError("update_horizon must be divisible by segment_length.")
         if cfg.save_frequency_updates < 0:
             raise ValueError("save_frequency_updates cannot be negative.")
+        if cfg.eval_frequency_steps < 0:
+            raise ValueError("eval_frequency_steps cannot be negative.")
+        if cfg.best_eval_mode not in {"min", "max"}:
+            raise ValueError("best_eval_mode must be 'min' or 'max'.")
         if algorithm.policy is not policy:
             raise ValueError("Trainer and APG must reference the same policy instance.")
         if torch.device(env.device) != algorithm.device:
@@ -83,6 +98,7 @@ def __init__(
         self.policy = policy
         self.algorithm = algorithm
         self.writer = writer
+        self.eval_env = eval_env
         self.collector = DifferentiableCollector(
             env=env,
             policy=policy,
@@ -91,7 +107,23 @@ def __init__(
         self.global_step = 0
         self.num_updates = 0
         self.train_history: list[dict[str, float]] = []
+        self.eval_history: list[dict[str, float]] = []
+        self.last_eval_metrics: dict[str, float] = {}
         self.latest_checkpoint_path: str | None = None
+        self.best_checkpoint_path: str | None = None
+        self.best_eval_value: float | None = None
+        self.start_time = time.time()
+        self.ret_window: deque[float] = deque(maxlen=100)
+        self.len_window: deque[float] = deque(maxlen=100)
+        self._episode_return = torch.zeros(
+            env.num_envs, dtype=torch.float32, device=algorithm.device
+        )
+        self._episode_length = torch.zeros(
+            env.num_envs, dtype=torch.long, device=algorithm.device
+        )
+        self._next_eval_step = (
+            cfg.eval_frequency_steps if cfg.eval_frequency_steps > 0 else None
+        )
 
     def train(self, total_timesteps: int) -> dict[str, Any]:
         """Train until at least ``total_timesteps`` vector transitions exist."""
@@ -115,6 +147,7 @@ def train(self, total_timesteps: int) -> dict[str, Any]:
                     rollout = self.collector.collect(
                         segment_steps,
                         deterministic=self.cfg.deterministic_actions,
+                        on_step_callback=self._on_step,
                     )
                     self.algorithm.accumulate_segment(rollout)
                     self.collector.detach_state()
@@ -126,14 +159,35 @@ def train(self, total_timesteps: int) -> dict[str, Any]:
 
             self.global_step += collected_steps * self.env.num_envs
             self.num_updates += 1
+            elapsed = max(time.time() - self.start_time, 1e-6)
             entry = {
                 "global_step": float(self.global_step),
                 "num_updates": float(self.num_updates),
+                "charts/SPS": float(self.global_step / elapsed),
+                "charts/episode_reward_avg_100": (
+                    float(sum(self.ret_window) / len(self.ret_window))
+                    if self.ret_window
+                    else float("nan")
+                ),
+                "charts/episode_length_avg_100": (
+                    float(sum(self.len_window) / len(self.len_window))
+                    if self.len_window
+                    else float("nan")
+                ),
                 **{f"train/{key}": value for key, value in metrics.items()},
             }
             self.train_history.append(entry)
             self._log(metrics)
 
+            if (
+                self._next_eval_step is not None
+                and self.eval_env is not None
+                and self.global_step >= self._next_eval_step
+            ):
+                self._evaluate()
+                while self._next_eval_step <= self.global_step:
+                    self._next_eval_step += self.cfg.eval_frequency_steps
+
             if (
                 self.cfg.save_frequency_updates > 0
                 and self.num_updates % self.cfg.save_frequency_updates == 0
@@ -158,6 +212,7 @@ def save_checkpoint(self, path: str | Path | None = None) -> str:
                 "num_updates": self.num_updates,
                 "policy": self.policy.state_dict(),
                 "optimizer": self.algorithm.optimizer.state_dict(),
+                "best_eval_value": self.best_eval_value,
             },
             checkpoint_path,
         )
@@ -184,6 +239,7 @@ def load_checkpoint(self, path: str | Path) -> None:
         self.algorithm.optimizer.load_state_dict(checkpoint["optimizer"])
         self.global_step = int(checkpoint["global_step"])
         self.num_updates = int(checkpoint["num_updates"])
+        self.best_eval_value = checkpoint.get("best_eval_value")
         self.latest_checkpoint_path = str(path)
 
     def get_summary(self) -> dict[str, Any]:
@@ -194,12 +250,83 @@ def get_summary(self) -> dict[str, Any]:
             "last_train_metrics": (
                 dict(self.train_history[-1]) if self.train_history else {}
             ),
+            "last_eval_metrics": dict(self.last_eval_metrics),
             "train_history": list(self.train_history),
+            "eval_history": list(self.eval_history),
             "latest_checkpoint_path": self.latest_checkpoint_path,
+            "best_checkpoint_path": self.best_checkpoint_path,
         }
 
-    def _log(self, metrics: dict[str, float]) -> None:
-        if self.writer is None:
+    def _on_step(self, transition: Any) -> None:
+        reward = transition.reward.detach()
+        done = transition.done.detach()
+        self._episode_return += reward
+        self._episode_length += 1
+        done_indices = torch.nonzero(done, as_tuple=False).squeeze(-1)
+        if done_indices.numel() == 0:
             return
-        for key, value in metrics.items():
-            self.writer.add_scalar(f"train/{key}", value, self.global_step)
+        self.ret_window.extend(
+            float(value) for value in self._episode_return[done_indices].cpu().tolist()
+        )
+        self.len_window.extend(
+            float(value) for value in self._episode_length[done_indices].cpu().tolist()
+        )
+        self._episode_return[done_indices] = 0.0
+        self._episode_length[done_indices] = 0
+
+    def _evaluate(self) -> dict[str, float]:
+        if self.eval_env is None:
+            return {}
+        metrics = evaluate_episodes(
+            policy=self.policy,
+            env=self.eval_env,
+            num_episodes=self.cfg.num_eval_episodes,
+            device=self.algorithm.device,
+            seed=self.cfg.eval_seed,
+        )
+        entry = {"global_step": float(self.global_step), **metrics}
+        self.eval_history.append(entry)
+        self.last_eval_metrics = entry
+        if self.writer is not None:
+            for key, value in metrics.items():
+                if math.isfinite(value):
+                    self.writer.add_scalar(key, value, self.global_step)
+        if self.cfg.use_wandb:
+            wandb.log(
+                {key: value for key, value in metrics.items() if math.isfinite(value)},
+                step=self.global_step,
+            )
+        candidate = metrics.get(self.cfg.best_eval_metric)
+        if candidate is not None and math.isfinite(candidate):
+            improved = self.best_eval_value is None or (
+                candidate > self.best_eval_value
+                if self.cfg.best_eval_mode == "max"
+                else candidate < self.best_eval_value
+            )
+            if improved:
+                self.best_eval_value = candidate
+                path = (
+                    Path(self.cfg.checkpoint_dir)
+                    / f"{self.cfg.experiment_name}_best.pt"
+                )
+                self.best_checkpoint_path = self.save_checkpoint(path)
+        return metrics
+
+    def _log(self, metrics: dict[str, float]) -> None:
+        elapsed = max(time.time() - self.start_time, 1e-6)
+        values = {
+            **{f"train/{key}": value for key, value in metrics.items()},
+            "charts/SPS": self.global_step / elapsed,
+        }
+        if self.ret_window:
+            values["charts/episode_reward_avg_100"] = sum(self.ret_window) / len(
+                self.ret_window
+            )
+            values["charts/episode_length_avg_100"] = sum(self.len_window) / len(
+                self.len_window
+            )
+        if self.writer is not None:
+            for key, value in values.items():
+                self.writer.add_scalar(key, value, self.global_step)
+        if self.cfg.use_wandb:
+            wandb.log(values, step=self.global_step)
diff --git a/embodichain/learning/rl/env.py b/embodichain/learning/rl/env.py
index c683a223a..5bf3716a4 100644
--- a/embodichain/learning/rl/env.py
+++ b/embodichain/learning/rl/env.py
@@ -14,39 +14,35 @@
 # limitations under the License.
 # ----------------------------------------------------------------------------
 
-"""Contracts for differentiable vector environments."""
+"""Contracts and registration helpers for lightweight learning environments."""
 
 from __future__ import annotations
 
+from collections.abc import Callable
 from typing import Any, Mapping, Protocol, TypeAlias, runtime_checkable
 
 import torch
 from gymnasium.spaces import Space
 from tensordict import TensorDict
 
-__all__ = ["DifferentiableObservation", "DifferentiableVecEnv"]
+__all__ = [
+    "DifferentiableObservation",
+    "DifferentiableVecEnv",
+    "LearningVecEnv",
+    "build_learning_env",
+    "get_registered_learning_env_names",
+    "register_learning_env",
+]
 
 DifferentiableObservation: TypeAlias = torch.Tensor | TensorDict
+LearningEnvFactory: TypeAlias = Callable[..., "LearningVecEnv"]
+
+_LEARNING_ENV_REGISTRY: dict[str, LearningEnvFactory] = {}
 
 
 @runtime_checkable
-class DifferentiableVecEnv(Protocol):
-    """Structural interface for batched differentiable environments.
-
-    Implementations must preserve the autograd path from ``action`` through
-    both the returned reward and the differentiable portion of the next
-    observation. Unlike a regular Gym environment, callers may retain this
-    graph across multiple calls to :meth:`step`.
-
-    ``detach_state`` defines an explicit truncated-backpropagation boundary.
-    It must detach differentiable internal state and return the corresponding
-    detached observation without resetting episode state, randomizing the
-    environment, or changing tensor values.
-
-    Implementations must auto-reset each environment that terminates or
-    truncates. For those environments, :meth:`step` returns the terminal reward
-    and done flags together with the initial observation of the next episode.
-    """
+class LearningVecEnv(Protocol):
+    """Structural interface shared by lightweight vector environments."""
 
     num_envs: int
     device: torch.device
@@ -69,9 +65,74 @@ def step(self, action: torch.Tensor) -> tuple[
         torch.Tensor,
         dict[str, Any],
     ]:
-        """Advance one step and auto-reset environments marked done."""
+        """Advance all environments by one step."""
         ...
 
+    def close(self) -> None:
+        """Release owned resources."""
+        ...
+
+
+@runtime_checkable
+class DifferentiableVecEnv(LearningVecEnv, Protocol):
+    """Batched env that preserves the autograd path through ``step``.
+
+    ``detach_state`` is the truncated-backpropagation boundary: detach
+    differentiable internal state and return the current observation without
+    resetting or resampling the episode. Finished rows must auto-reset inside
+    ``step``, returning the terminal reward/done with the next initial observation.
+    """
+
     def detach_state(self) -> DifferentiableObservation:
         """Detach internal state and return its current detached observation."""
         ...
+
+
+def register_learning_env(
+    name: str,
+    factory: LearningEnvFactory | None = None,
+    *,
+    override: bool = False,
+) -> Callable[[LearningEnvFactory], LearningEnvFactory] | LearningEnvFactory:
+    """Register a lightweight vector-environment factory.
+
+    The function supports both ``@register_learning_env("Name")`` and direct
+    ``register_learning_env("Name", Factory)`` use.
+    """
+
+    def decorator(env_factory: LearningEnvFactory) -> LearningEnvFactory:
+        key = name.lower()
+        if key in _LEARNING_ENV_REGISTRY and not override:
+            raise ValueError(f"Learning environment '{name}' is already registered.")
+        _LEARNING_ENV_REGISTRY[key] = env_factory
+        return env_factory
+
+    if factory is None:
+        return decorator
+    return decorator(factory)
+
+
+def get_registered_learning_env_names() -> list[str]:
+    """Return registered lightweight environment names."""
+    return sorted(_LEARNING_ENV_REGISTRY)
+
+
+def build_learning_env(
+    name: str,
+    *,
+    num_envs: int,
+    device: torch.device | str,
+    **cfg: Any,
+) -> LearningVecEnv:
+    """Build a registered lightweight vector environment."""
+    key = name.lower()
+    if key not in _LEARNING_ENV_REGISTRY:
+        available = ", ".join(get_registered_learning_env_names()) or "<none>"
+        raise ValueError(
+            f"Learning environment '{name}' is not registered. Available: {available}"
+        )
+    return _LEARNING_ENV_REGISTRY[key](
+        num_envs=num_envs,
+        device=torch.device(device),
+        **cfg,
+    )
diff --git a/embodichain/learning/rl/evaluation.py b/embodichain/learning/rl/evaluation.py
new file mode 100644
index 000000000..d64ebcfe3
--- /dev/null
+++ b/embodichain/learning/rl/evaluation.py
@@ -0,0 +1,142 @@
+# ----------------------------------------------------------------------------
+# 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.
+# ----------------------------------------------------------------------------
+
+"""Shared deterministic episode evaluation for all RL trainers."""
+
+from __future__ import annotations
+
+from collections.abc import Callable, Mapping
+from typing import Any
+
+import numpy as np
+import torch
+from tensordict import TensorDict
+
+from embodichain.learning.rl.utils import (
+    dict_to_tensordict,
+    flatten_dict_observation,
+)
+
+__all__ = ["evaluate_episodes"]
+
+
+def _flat_observation(observation: Any, device: torch.device) -> torch.Tensor:
+    tensor_dict = dict_to_tensordict(observation, device)
+    return flatten_dict_observation(tensor_dict)
+
+
+def _action_for_env(env: Any, action: torch.Tensor) -> Any:
+    action_manager = getattr(env, "action_manager", None)
+    if action_manager is None and hasattr(env, "get_wrapper_attr"):
+        try:
+            action_manager = env.get_wrapper_attr("action_manager")
+        except AttributeError:
+            action_manager = None
+    if action_manager is None:
+        return action
+    return action_manager.convert_policy_action_to_env_action(action)
+
+
+def _selected_values(value: Any, indices: torch.Tensor) -> list[float]:
+    if isinstance(value, torch.Tensor):
+        selected = value.detach().reshape(-1)[indices.to(value.device)]
+        return [float(item) for item in selected.cpu().tolist()]
+    array = np.asarray(value).reshape(-1)
+    return [float(array[index]) for index in indices.cpu().tolist()]
+
+
+@torch.no_grad()
+def evaluate_episodes(
+    *,
+    policy: torch.nn.Module,
+    env: Any,
+    num_episodes: int,
+    device: torch.device | str,
+    seed: int | None = None,
+    on_step: Callable[[dict[str, Any]], None] | None = None,
+) -> dict[str, float]:
+    """Evaluate exactly ``num_episodes`` completed asynchronous episodes."""
+    if num_episodes <= 0:
+        raise ValueError("num_episodes must be positive.")
+    device = torch.device(device)
+    previous_training = policy.training
+    policy.eval()
+
+    returns: list[float] = []
+    lengths: list[float] = []
+    successes: list[float] = []
+    metric_values: dict[str, list[float]] = {}
+    num_envs = int(env.num_envs)
+    current_return = torch.zeros(num_envs, dtype=torch.float32, device=device)
+    current_length = torch.zeros(num_envs, dtype=torch.long, device=device)
+
+    try:
+        observation, _ = env.reset(seed=seed)
+        while len(returns) < num_episodes:
+            flat_observation = _flat_observation(observation, device)
+            policy_input = TensorDict(
+                {"obs": flat_observation},
+                batch_size=[num_envs],
+                device=device,
+            )
+            policy_output = policy.get_action(policy_input, deterministic=True)
+            observation, reward, terminated, truncated, info = env.step(
+                _action_for_env(env, policy_output["action"])
+            )
+            reward = torch.as_tensor(reward, device=device).reshape(num_envs)
+            done = (
+                torch.as_tensor(terminated, device=device, dtype=torch.bool)
+                | torch.as_tensor(truncated, device=device, dtype=torch.bool)
+            ).reshape(num_envs)
+            current_return += reward.float()
+            current_length += 1
+
+            if on_step is not None:
+                on_step(info)
+
+            done_indices = torch.nonzero(done, as_tuple=False).squeeze(-1)
+            if done_indices.numel() == 0:
+                continue
+            remaining = num_episodes - len(returns)
+            selected = done_indices[:remaining]
+            returns.extend(_selected_values(current_return, selected))
+            lengths.extend(_selected_values(current_length, selected))
+
+            if isinstance(info, Mapping):
+                if "success" in info:
+                    successes.extend(_selected_values(info["success"], selected))
+                metrics = info.get("metrics", {})
+                if isinstance(metrics, Mapping):
+                    for name, value in metrics.items():
+                        metric_values.setdefault(str(name), []).extend(
+                            _selected_values(value, selected)
+                        )
+
+
+            current_return[done_indices] = 0.0
+            current_length[done_indices] = 0
+    finally:
+        policy.train(previous_training)
+
+    result = {
+        "eval/avg_reward": float(np.mean(returns)),
+        "eval/avg_length": float(np.mean(lengths)),
+        "eval/success_rate": (float(np.mean(successes)) if successes else float("nan")),
+    }
+    for name, values in metric_values.items():
+        if values:
+            result[f"eval/metrics/{name}"] = float(np.mean(values))
+    return result
diff --git a/embodichain/learning/rl/models/__init__.py b/embodichain/learning/rl/models/__init__.py
index 122960066..fa6d46ae6 100644
--- a/embodichain/learning/rl/models/__init__.py
+++ b/embodichain/learning/rl/models/__init__.py
@@ -137,9 +137,19 @@ def build_mlp_from_cfg(module_cfg: Dict, in_dim: int, out_dim: int) -> MLP:
     if module_cfg.get("type", "").lower() != "mlp":
         raise ValueError("Only 'mlp' type is supported for actor/critic in this setup.")
 
-    hidden_sizes = module_cfg["network_cfg"]["hidden_sizes"]
-    activation = module_cfg["network_cfg"]["activation"]
-    return MLP(in_dim, out_dim, hidden_sizes, activation)
+    network_cfg = module_cfg["network_cfg"]
+    model = MLP(
+        in_dim,
+        out_dim,
+        network_cfg["hidden_sizes"],
+        network_cfg.get("activation", "elu"),
+        last_activation=network_cfg.get("last_activation"),
+        use_layernorm=bool(network_cfg.get("use_layernorm", False)),
+        dropout_p=float(network_cfg.get("dropout_p", 0.0)),
+    )
+    if "orthogonal_init" in network_cfg:
+        model.init_orthogonal(network_cfg["orthogonal_init"])
+    return model
 
 
 # default registrations
diff --git a/embodichain/learning/rl/routing.py b/embodichain/learning/rl/routing.py
new file mode 100644
index 000000000..8d7504284
--- /dev/null
+++ b/embodichain/learning/rl/routing.py
@@ -0,0 +1,39 @@
+# ----------------------------------------------------------------------------
+# 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.
+# ----------------------------------------------------------------------------
+
+"""Trainer routing based on an algorithm's rollout semantics."""
+
+from __future__ import annotations
+
+from typing import Any, TypeAlias
+
+from embodichain.learning.rl.algo import RolloutKind
+from embodichain.learning.rl.differentiable_trainer import DifferentiableTrainer
+from embodichain.learning.rl.utils.trainer import Trainer
+
+__all__ = ["TrainerType", "get_trainer_class"]
+
+TrainerType: TypeAlias = type[Trainer] | type[DifferentiableTrainer]
+
+
+def get_trainer_class(algorithm: Any) -> TrainerType:
+    """Return the trainer compatible with ``algorithm``."""
+    rollout_kind = getattr(algorithm, "rollout_kind", RolloutKind.STANDARD)
+    if rollout_kind == RolloutKind.DIFFERENTIABLE:
+        return DifferentiableTrainer
+    if rollout_kind == RolloutKind.STANDARD:
+        return Trainer
+    raise ValueError(f"Unsupported rollout kind: {rollout_kind!r}.")
diff --git a/embodichain/learning/rl/train.py b/embodichain/learning/rl/train.py
index 33234483e..47677ca68 100644
--- a/embodichain/learning/rl/train.py
+++ b/embodichain/learning/rl/train.py
@@ -30,7 +30,17 @@
 
 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 build_algo, get_registered_algo_names
+from embodichain.learning.rl.algo import (
+    RolloutKind,
+    build_algo,
+    get_registered_algo_names,
+)
+from embodichain.learning.rl.differentiable_trainer import (
+    DifferentiableTrainer,
+    DifferentiableTrainerCfg,
+)
+from embodichain.learning.rl.env import build_learning_env
+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
@@ -76,6 +86,188 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
     return parser.parse_args(argv)
 
 
+def _build_learning_policy(
+    policy_block: dict,
+    env,
+    device: torch.device,
+):
+    obs_dim = int(env.single_observation_space.shape[-1])
+    action_dim = int(env.single_action_space.shape[-1])
+    policy_name = policy_block["name"].lower()
+    actor_cfg = policy_block.get("actor")
+    critic_cfg = policy_block.get("critic")
+    actor = (
+        build_mlp_from_cfg(actor_cfg, obs_dim, action_dim)
+        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
+    )
+    policy = build_policy(
+        policy_block,
+        env.single_observation_space,
+        env.single_action_space,
+        device,
+        actor=actor,
+        critic=critic,
+    )
+    if "initial_log_std" in policy_block and hasattr(policy, "log_std"):
+        with torch.no_grad():
+            policy.log_std.fill_(float(policy_block["initial_log_std"]))
+    return policy
+
+
+def _train_learning_env(
+    cfg_data: dict,
+    *,
+    distributed: bool | None,
+):
+    """Train a lightweight registered environment through the unified CLI."""
+    trainer_cfg = cfg_data["trainer"]
+    policy_block = cfg_data["policy"]
+    algorithm_block = cfg_data["algorithm"]
+    distributed = (
+        bool(trainer_cfg.get("distributed", False))
+        if distributed is None
+        else distributed
+    )
+    if distributed:
+        raise ValueError(
+            "Learning environments do not yet support distributed training."
+        )
+
+    discover_task_packages()
+    execute_init_hooks()
+    seed = int(trainer_cfg.get("seed", 1))
+    device = torch.device(trainer_cfg.get("device", "cpu"))
+    if device.type == "cuda" and not torch.cuda.is_available():
+        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)
+
+    env_block = trainer_cfg["learning_env"]
+    if isinstance(env_block, str):
+        env_name = env_block
+        env_cfg = {}
+    else:
+        env_name = env_block["name"]
+        env_cfg = dict(env_block.get("cfg", {}))
+    num_envs = int(trainer_cfg.get("num_envs", 64))
+    env = build_learning_env(
+        env_name,
+        num_envs=num_envs,
+        device=device,
+        **env_cfg,
+    )
+
+    enable_eval = bool(trainer_cfg.get("enable_eval", False))
+    eval_env = None
+    if enable_eval:
+        eval_env = build_learning_env(
+            env_name,
+            num_envs=int(trainer_cfg.get("num_eval_envs", 16)),
+            device=device,
+            **env_cfg,
+        )
+
+    policy = _build_learning_policy(policy_block, env, device)
+    algorithm = build_algo(
+        algorithm_block["name"],
+        dict(algorithm_block.get("cfg", {})),
+        policy,
+        device,
+    )
+    trainer_class = get_trainer_class(algorithm)
+
+    exp_name = trainer_cfg.get("exp_name", f"{env_name}_{algorithm_block['name']}")
+    run_stamp = time.strftime("%Y%m%d_%H%M%S")
+    run_base = Path("outputs") / f"{exp_name}_{run_stamp}"
+    log_dir = run_base / "logs" / exp_name
+    checkpoint_dir = run_base / "checkpoints"
+    checkpoint_dir.mkdir(parents=True, exist_ok=True)
+    writer = SummaryWriter(str(log_dir))
+    use_wandb = bool(trainer_cfg.get("use_wandb", False))
+    if use_wandb:
+        wandb.init(
+            project=trainer_cfg.get("wandb_project_name", "embodichain-point-mass"),
+            name=exp_name,
+            config=cfg_data,
+        )
+
+    eval_freq = int(trainer_cfg.get("eval_freq", 0)) if enable_eval else 0
+    eval_seed = int(trainer_cfg.get("eval_seed", seed + 10_000))
+    iterations = int(trainer_cfg.get("iterations", 250))
+    try:
+        if trainer_class is DifferentiableTrainer:
+            segment_length = int(trainer_cfg.get("segment_length", 16))
+            update_horizon = int(trainer_cfg.get("update_horizon", segment_length))
+            diff_cfg = DifferentiableTrainerCfg(
+                segment_length=segment_length,
+                update_horizon=update_horizon,
+                deterministic_actions=bool(
+                    trainer_cfg.get("deterministic_actions", False)
+                ),
+                checkpoint_dir=str(checkpoint_dir),
+                experiment_name=exp_name,
+                save_frequency_updates=int(
+                    trainer_cfg.get("save_frequency_updates", 0)
+                ),
+                eval_frequency_steps=eval_freq,
+                num_eval_episodes=int(trainer_cfg.get("num_eval_episodes", 5)),
+                eval_seed=eval_seed,
+                use_wandb=use_wandb,
+                best_eval_metric=trainer_cfg.get("best_eval_metric", "eval/avg_reward"),
+                best_eval_mode=trainer_cfg.get("best_eval_mode", "max"),
+            )
+            trainer = DifferentiableTrainer(
+                cfg=diff_cfg,
+                env=env,
+                policy=policy,
+                algorithm=algorithm,
+                writer=writer,
+                eval_env=eval_env,
+            )
+            default_steps = iterations * update_horizon * num_envs
+        else:
+            buffer_size = int(
+                trainer_cfg.get("buffer_size", trainer_cfg.get("rollout_steps", 256))
+            )
+            trainer = Trainer(
+                policy=policy,
+                env=env,
+                algorithm=algorithm,
+                buffer_size=buffer_size,
+                batch_size=int(algorithm.cfg.batch_size),
+                writer=writer,
+                eval_freq=eval_freq,
+                save_freq=int(trainer_cfg.get("save_freq", 0)),
+                checkpoint_dir=str(checkpoint_dir),
+                exp_name=exp_name,
+                use_wandb=use_wandb,
+                eval_env=eval_env,
+                num_eval_episodes=int(trainer_cfg.get("num_eval_episodes", 5)),
+                eval_seed=eval_seed,
+                best_eval_metric=trainer_cfg.get("best_eval_metric", "eval/avg_reward"),
+                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)
+        trainer.save_checkpoint()
+        return trainer.get_summary()
+    finally:
+        writer.close()
+        if use_wandb:
+            wandb.finish()
+        env.close()
+        if eval_env is not None:
+            eval_env.close()
+
+
 def train_from_config(config_path: str, distributed: bool | None = None):
     """Run training from a config file path.
 
@@ -87,14 +279,14 @@ def train_from_config(config_path: str, distributed: bool | None = None):
     cfg_data = load_config(config_path)
 
     trainer_cfg = cfg_data["trainer"]
+    if "learning_env" in trainer_cfg:
+        return _train_learning_env(cfg_data, distributed=distributed)
     policy_block = cfg_data["policy"]
     algo_block = cfg_data["algorithm"]
 
-    # Resolve distributed flag
     if distributed is None:
         distributed = bool(trainer_cfg.get("distributed", False))
 
-    # Distributed setup
     rank = 0
     world_size = 1
     local_rank = 0
@@ -120,7 +312,6 @@ def train_from_config(config_path: str, distributed: bool | None = None):
         rank = torch.distributed.get_rank()
         world_size = torch.distributed.get_world_size()
 
-    # Runtime
     exp_name = trainer_cfg.get("exp_name", "generic_exp")
     seed = int(trainer_cfg.get("seed", 1))
     device_str = trainer_cfg.get("device", "cpu")
@@ -140,7 +331,6 @@ def train_from_config(config_path: str, distributed: bool | None = None):
     num_envs = trainer_cfg.get("num_envs", None)
     wandb_project_name = trainer_cfg.get("wandb_project_name", "embodichain-generic")
 
-    # Device
     if not isinstance(device_str, str):
         raise ValueError(
             f"runtime.device must be a string such as 'cpu' or 'cuda:0'. Got: {device_str!r}"
@@ -313,6 +503,11 @@ def train_from_config(config_path: str, distributed: bool | None = None):
         device,
         distributed=distributed,
     )
+    if algo.rollout_kind is RolloutKind.DIFFERENTIABLE:
+        raise ValueError(
+            "Differentiable algorithms require trainer.learning_env; "
+            "simulator gym_config environments use standard rollouts."
+        )
 
     # Build Trainer
     event_modules = [
@@ -373,6 +568,9 @@ def train_from_config(config_path: str, distributed: bool | None = None):
         distributed=distributed,
         rank=rank,
         world_size=world_size,
+        eval_seed=int(trainer_cfg.get("eval_seed", seed + 10_000)),
+        best_eval_metric=trainer_cfg.get("best_eval_metric", "eval/avg_reward"),
+        best_eval_mode=trainer_cfg.get("best_eval_mode", "max"),
     )
 
     if rank == 0:
diff --git a/embodichain/learning/rl/utils/helper.py b/embodichain/learning/rl/utils/helper.py
index e443a6491..7077ddd3f 100644
--- a/embodichain/learning/rl/utils/helper.py
+++ b/embodichain/learning/rl/utils/helper.py
@@ -56,12 +56,13 @@ def _collect_tensors(data: TensorDict) -> None:
 
 
 def dict_to_tensordict(
-    obs_dict: TensorDict | Mapping[str, Any], device: torch.device | str
+    obs_dict: torch.Tensor | TensorDict | Mapping[str, Any],
+    device: torch.device | str,
 ) -> TensorDict:
     """Convert an environment observation mapping into a TensorDict.
 
     Args:
-        obs_dict: Environment observation returned by `reset()` or `step()`.
+        obs_dict: Tensor or mapping returned by ``reset()`` or ``step()``.
         device: Target device for the resulting TensorDict.
 
     Returns:
@@ -69,8 +70,13 @@ def dict_to_tensordict(
     """
     if isinstance(obs_dict, TensorDict):
         return obs_dict.to(device)
+    if isinstance(obs_dict, torch.Tensor):
+        tensor = obs_dict.to(device)
+        batch_size = [tensor.shape[0]] if tensor.ndim > 1 else []
+        return TensorDict({"obs": tensor}, batch_size=batch_size, device=device)
     if not isinstance(obs_dict, Mapping):
         raise TypeError(
-            f"Expected observation mapping or TensorDict, got {type(obs_dict)!r}."
+            f"Expected tensor, observation mapping, or TensorDict, "
+            f"got {type(obs_dict)!r}."
         )
     return TensorDict.from_dict(dict(obs_dict), device=device)
diff --git a/embodichain/learning/rl/utils/trainer.py b/embodichain/learning/rl/utils/trainer.py
index a4e824ee0..678696346 100644
--- a/embodichain/learning/rl/utils/trainer.py
+++ b/embodichain/learning/rl/utils/trainer.py
@@ -16,22 +16,20 @@
 
 from __future__ import annotations
 
-from typing import Any, Dict
 import time
+from collections import deque
+from typing import Any
+
 import numpy as np
 import torch
 import wandb
-from embodichain.utils import logger
-from torch.utils.tensorboard import SummaryWriter
-from collections import deque
 from tensordict import TensorDict
-from typing import Dict
+from torch.utils.tensorboard import SummaryWriter
 
-from embodichain.lab.gym.envs.managers.action_manager import ActionManager
+from embodichain.lab.gym.envs.managers.event_manager import EventManager
 from embodichain.learning.rl.buffer import RolloutBuffer
 from embodichain.learning.rl.collector import SyncCollector
-from embodichain.lab.gym.envs.managers.event_manager import EventManager
-from .helper import flatten_dict_observation
+from embodichain.learning.rl.evaluation import evaluate_episodes
 
 
 class Trainer:
@@ -57,7 +55,12 @@ def __init__(
         distributed: bool = False,
         rank: int = 0,
         world_size: int = 1,
+        eval_seed: int | None = None,
+        best_eval_metric: str = "eval/avg_reward",
+        best_eval_mode: str = "max",
     ):
+        if best_eval_mode not in {"min", "max"}:
+            raise ValueError("best_eval_mode must be 'min' or 'max'.")
         self.policy = policy
         self.distributed = distributed
         self.rank = rank
@@ -74,13 +77,17 @@ def __init__(
         self.exp_name = exp_name
         self.use_wandb = use_wandb
         self.num_eval_episodes = num_eval_episodes
+        self.eval_seed = eval_seed
+        self.best_eval_metric = best_eval_metric
+        self.best_eval_mode = best_eval_mode
+        self.best_eval_value: float | None = None
+        self.best_checkpoint_path: str | None = None
 
         if event_cfg is not None:
             self.event_manager = EventManager(event_cfg, env=self.env)
         if eval_event_cfg is not None:
             self.eval_event_manager = EventManager(eval_event_cfg, env=self.eval_env)
 
-        # Get device from algorithm
         self.device = self.algorithm.device
         self.global_step = 0
         self.start_time = time.time()
@@ -91,6 +98,8 @@ def __init__(
         self.last_eval_metrics: dict[str, float] = {}
         self.last_train_metrics: dict[str, float] = {}
         self.latest_checkpoint_path: str | None = None
+        self._next_eval_step = eval_freq if eval_freq > 0 else None
+        self._next_save_step = save_freq if save_freq > 0 else None
         num_envs = getattr(self.env, "num_envs", None)
         if num_envs is None:
             raise RuntimeError("Env must expose num_envs for trainer statistics.")
@@ -117,11 +126,9 @@ def __init__(
             ),
         )
 
-        # episode stats tracked on device to avoid repeated CPU round-trips
         self.curr_ret = torch.zeros(num_envs, dtype=torch.float32, device=self.device)
         self.curr_len = torch.zeros(num_envs, dtype=torch.int32, device=self.device)
 
-    # ---- lightweight helpers for dense logging ----
     @staticmethod
     def _mean_scalar(x) -> float:
         if hasattr(x, "detach"):
@@ -152,7 +159,7 @@ def _pack_log_dict(self, prefix: str, data: dict) -> dict:
                 continue
         return out
 
-    def train(self, total_timesteps: int) -> Dict[str, Any]:
+    def train(self, total_timesteps: int) -> dict[str, Any]:
         if self.rank == 0:
             print(f"Start training, total steps: {total_timesteps}")
         while self.global_step < total_timesteps:
@@ -160,13 +167,20 @@ def train(self, total_timesteps: int) -> Dict[str, Any]:
             losses = self.algorithm.update(self.buffer.get(flatten=False))
             self._log_train(losses)
             if (
-                self.eval_freq > 0
+                self._next_eval_step is not None
                 and self.eval_env is not None
-                and self.global_step % self.eval_freq == 0
+                and self.global_step >= self._next_eval_step
             ):
                 self._eval_once(num_episodes=self.num_eval_episodes)
-            if self.global_step % self.save_freq == 0:
+                while self._next_eval_step <= self.global_step:
+                    self._next_eval_step += self.eval_freq
+            if (
+                self._next_save_step is not None
+                and self.global_step >= self._next_save_step
+            ):
                 self.save_checkpoint()
+                while self._next_save_step <= self.global_step:
+                    self._next_save_step += self.save_freq
         return self.get_summary()
 
     @torch.no_grad()
@@ -285,7 +299,7 @@ def _sync_episode_stats(self) -> None:
             self.ret_window.extend(all_ret[start:])
             self.len_window.extend(all_len[start:])
 
-    def _log_train(self, losses: Dict[str, float]):
+    def _log_train(self, losses: dict[str, float]):
         elapsed = max(1e-6, time.time() - self.start_time)
         sps = self.global_step / elapsed
         avgR = np.mean(self.ret_window) if len(self.ret_window) > 0 else float("nan")
@@ -316,13 +330,10 @@ def _log_train(self, losses: Dict[str, float]):
                     float(np.mean(self.len_window)),
                     self.global_step,
                 )
-        # console and external logging are rank-0 only in distributed mode.
         if self.rank == 0:
             print(
                 f"[train] step={self.global_step} sps={sps:.0f} avgReward(100)={avgR:.3f} avgLength(100)={avgL:.1f}"
             )
-
-            # wandb (mirror TB logs)
             if self.use_wandb:
                 log_dict = {f"train/{k}": v for k, v in losses.items()}
                 log_dict["charts/SPS"] = sps
@@ -333,163 +344,87 @@ def _log_train(self, losses: Dict[str, float]):
                 wandb.log(log_dict, step=self.global_step)
 
     @torch.no_grad()
-    def _eval_once(self, num_episodes: int = 5) -> Dict[str, float]:
-        """Run evaluation for specified number of episodes.
-
-        Each episode runs all parallel environments until completion, allowing
-        environments to finish at different times.
-
-        Args:
-            num_episodes: Number of episodes to evaluate
-        """
-        self.policy.eval()
-        episode_returns = []
-        episode_lengths = []
-        episode_successes = []
-        metric_values: dict[str, list[float]] = {}
-
-        # Evaluation does not consume the training rollout buffer; binding it here can
-        # overflow the shared RL buffer when eval episodes are longer than buffer_size.
-        for _ in range(num_episodes):
-            # Reset and initialize episode tracking
-            obs, _ = self.eval_env.reset()
-            obs = flatten_dict_observation(obs)
-            num_envs = obs.shape[0] if obs.ndim == 2 else 1
-
-            done_mask = torch.zeros(num_envs, dtype=torch.bool, device=self.device)
-            cumulative_reward = torch.zeros(
-                num_envs, dtype=torch.float32, device=self.device
-            )
-            step_count = torch.zeros(num_envs, dtype=torch.int32, device=self.device)
-
-            # Run episode until all environments complete
-            while not done_mask.all():
-                # Get deterministic actions from policy
-                action_td = TensorDict(
-                    {"obs": obs},
-                    batch_size=[num_envs],
-                    device=self.device,
-                )
-                action_td = self.policy.get_action(action_td, deterministic=True)
-                actions = action_td["action"]
-                am: ActionManager = getattr(self.eval_env, "action_manager", None)
-                if am is None:
-                    action_in = actions
-                else:
-                    action_in = am.convert_policy_action_to_env_action(actions)
-
-                # Environment step
-                obs, reward, terminated, truncated, info = self.eval_env.step(action_in)
-                obs = (
-                    flatten_dict_observation(obs)
-                    if isinstance(obs, TensorDict)
-                    else obs
-                )
+    def _eval_once(self, num_episodes: int = 5) -> dict[str, float]:
+        """Evaluate ``num_episodes`` completed asynchronous episodes."""
 
-                # Update statistics only for still-running environments
-                done = terminated | truncated
-                still_running = ~done_mask
-                cumulative_reward[still_running] += reward[still_running].float()
-                step_count[still_running] += 1
-                newly_done = done & (~done_mask)
-                if newly_done.any():
-                    if isinstance(info, dict) and "success" in info:
-                        successes = info["success"][newly_done].detach().cpu().tolist()
-                        episode_successes.extend([float(v) for v in successes])
-                    if isinstance(info, dict) and "metrics" in info:
-                        for key, value in info["metrics"].items():
-                            values = value[newly_done].detach().cpu().tolist()
-                            metric_values.setdefault(key, []).extend(
-                                [float(v) for v in values]
-                            )
-                done_mask |= done
-
-                # Trigger evaluation events (e.g., video recording)
-                if hasattr(self, "eval_event_manager"):
-                    if "interval" in self.eval_event_manager.available_modes:
-                        self.eval_event_manager.apply(mode="interval")
-
-            # Collect episode statistics
-            episode_returns.extend(cumulative_reward.cpu().tolist())
-            episode_lengths.extend(step_count.cpu().tolist())
-
-        # Finalize evaluation functors (e.g., video recording)
-        if hasattr(self, "eval_event_manager"):
-            for functor_cfg in self.eval_event_manager._mode_functor_cfgs.get(
-                "interval", []
-            ):
-                functor = functor_cfg.func
-                save_path = functor_cfg.params.get("save_path", "./outputs/videos/eval")
-
-                if hasattr(functor, "flush"):
-                    functor.flush(save_path)
-                if hasattr(functor, "finalize"):
-                    functor.finalize(save_path)
-
-        # Log evaluation metrics
-        if self.writer and episode_returns:
-            self.writer.add_scalar(
-                "eval/avg_reward", float(np.mean(episode_returns)), self.global_step
-            )
-            self.writer.add_scalar(
-                "eval/avg_length", float(np.mean(episode_lengths)), self.global_step
-            )
-            if episode_successes:
-                self.writer.add_scalar(
-                    "eval/success_rate",
-                    float(np.mean(episode_successes)),
-                    self.global_step,
-                )
+        def on_step(_: dict[str, Any]) -> None:
+            if hasattr(self, "eval_event_manager"):
+                if "interval" in self.eval_event_manager.available_modes:
+                    self.eval_event_manager.apply(mode="interval")
 
-        summary = {
-            "global_step": float(self.global_step),
-            "eval/avg_reward": (
-                float(np.mean(episode_returns)) if episode_returns else float("nan")
-            ),
-            "eval/avg_length": (
-                float(np.mean(episode_lengths)) if episode_lengths else float("nan")
-            ),
-            "eval/success_rate": (
-                float(np.mean(episode_successes)) if episode_successes else float("nan")
-            ),
-        }
-        for key, values in metric_values.items():
-            if values:
-                summary[f"eval/metrics/{key}"] = float(np.mean(values))
+        metrics = evaluate_episodes(
+            policy=self.policy,
+            env=self.eval_env,
+            num_episodes=num_episodes,
+            device=self.device,
+            seed=self.eval_seed,
+            on_step=on_step,
+        )
+        summary = {"global_step": float(self.global_step), **metrics}
         self.eval_history.append(summary)
         self.last_eval_metrics = summary
+        if self.writer:
+            for key, value in metrics.items():
+                if np.isfinite(value):
+                    self.writer.add_scalar(key, value, self.global_step)
         if self.rank == 0 and self.use_wandb:
-            log_dict = {
-                key: value
-                for key, value in summary.items()
-                if key != "global_step" and not np.isnan(value)
-            }
-            if log_dict:
-                wandb.log(log_dict, step=self.global_step)
+            wandb.log(
+                {key: value for key, value in metrics.items() if np.isfinite(value)},
+                step=self.global_step,
+            )
+        candidate = metrics.get(self.best_eval_metric)
+        if candidate is not None and np.isfinite(candidate):
+            improved = self.best_eval_value is None or (
+                candidate > self.best_eval_value
+                if self.best_eval_mode == "max"
+                else candidate < self.best_eval_value
+            )
+            if improved:
+                self.best_eval_value = candidate
+                best_path = f"{self.checkpoint_dir}/{self.exp_name}_best.pt"
+                self.best_checkpoint_path = self.save_checkpoint(best_path)
+        self._finalize_eval_events()
         return summary
 
-    def save_checkpoint(self) -> str | None:
-        # minimal model-only checkpoint; trainer/algorithm states can be added
+    def _finalize_eval_events(self) -> None:
+        """Flush evaluation event functors such as video recorders."""
+        if not hasattr(self, "eval_event_manager"):
+            return
+        for functor_cfg in self.eval_event_manager._mode_functor_cfgs.get(
+            "interval", []
+        ):
+            functor = functor_cfg.func
+            save_path = functor_cfg.params.get("save_path", "./outputs/videos/eval")
+            if hasattr(functor, "flush"):
+                functor.flush(save_path)
+            if hasattr(functor, "finalize"):
+                functor.finalize(save_path)
+
+    def save_checkpoint(self, path: str | None = None) -> str | None:
+        """Save policy, optimizer (when available), and trainer counters."""
         if self.rank != 0:
             return None
-        path = f"{self.checkpoint_dir}/{self.exp_name}_step_{self.global_step}.pt"
+        if path is None:
+            path = f"{self.checkpoint_dir}/{self.exp_name}_step_{self.global_step}.pt"
         policy_state = (
             self.policy.module.state_dict()
             if hasattr(self.policy, "module")
             else self.policy.state_dict()
         )
-        torch.save(
-            {
-                "global_step": self.global_step,
-                "policy": policy_state,
-            },
-            path,
-        )
+        checkpoint = {
+            "global_step": self.global_step,
+            "policy": policy_state,
+            "best_eval_value": self.best_eval_value,
+        }
+        optimizer = getattr(self.algorithm, "optimizer", None)
+        if optimizer is not None:
+            checkpoint["optimizer"] = optimizer.state_dict()
+        torch.save(checkpoint, path)
         self.latest_checkpoint_path = path
         print(f"Checkpoint saved: {path}")
         return path
 
-    def get_summary(self) -> Dict[str, Any]:
+    def get_summary(self) -> dict[str, Any]:
         elapsed = max(1e-6, time.time() - self.start_time)
         return {
             "global_step": int(self.global_step),
@@ -500,4 +435,5 @@ def get_summary(self) -> Dict[str, Any]:
             "train_history": list(self.train_history),
             "eval_history": list(self.eval_history),
             "latest_checkpoint_path": self.latest_checkpoint_path,
+            "best_checkpoint_path": self.best_checkpoint_path,
         }
diff --git a/embodichain_tasks/configs/agents/rl/basic/point_mass/train_apg.yaml b/embodichain_tasks/configs/agents/rl/basic/point_mass/train_apg.yaml
new file mode 100644
index 000000000..38f4b6dc6
--- /dev/null
+++ b/embodichain_tasks/configs/agents/rl/basic/point_mass/train_apg.yaml
@@ -0,0 +1,42 @@
+trainer:
+  exp_name: point_mass_apg
+  learning_env:
+    name: PointMassRL
+    cfg:
+      max_episode_steps: 100
+      success_threshold: 0.03
+  seed: 42
+  device: cuda:0
+  num_envs: 64
+  iterations: 500
+  segment_length: 100
+  update_horizon: 100
+  enable_eval: true
+  eval_freq: 32000
+  num_eval_envs: 32
+  num_eval_episodes: 128
+  eval_seed: 10042
+  best_eval_metric: eval/success_rate
+  best_eval_mode: max
+  save_frequency_updates: 50
+  use_wandb: false
+  wandb_project_name: embodichain-point-mass
+
+policy:
+  name: actor_only
+  initial_log_std: -2.0
+  actor:
+    type: mlp
+    network_cfg:
+      hidden_sizes: [64, 64]
+      activation: tanh
+      last_activation: tanh
+      orthogonal_init: [1.414, 1.414, 0.01]
+
+algorithm:
+  name: apg
+  cfg:
+    learning_rate: 0.00025
+    gamma: 0.99
+    max_grad_norm: 0.5
+    ent_coef: 0.0
diff --git a/embodichain_tasks/configs/agents/rl/basic/point_mass/train_ppo.yaml b/embodichain_tasks/configs/agents/rl/basic/point_mass/train_ppo.yaml
new file mode 100644
index 000000000..e18b8f8db
--- /dev/null
+++ b/embodichain_tasks/configs/agents/rl/basic/point_mass/train_ppo.yaml
@@ -0,0 +1,52 @@
+trainer:
+  exp_name: point_mass_ppo
+  learning_env:
+    name: PointMassRL
+    cfg:
+      max_episode_steps: 100
+      success_threshold: 0.03
+  seed: 42
+  device: cuda:0
+  num_envs: 64
+  iterations: 500
+  buffer_size: 256
+  enable_eval: true
+  eval_freq: 32768
+  num_eval_envs: 32
+  num_eval_episodes: 128
+  eval_seed: 10042
+  best_eval_metric: eval/success_rate
+  best_eval_mode: max
+  save_freq: 262144
+  use_wandb: false
+  wandb_project_name: embodichain-point-mass
+
+policy:
+  name: actor_critic
+  initial_log_std: -1.0
+  actor:
+    type: mlp
+    network_cfg:
+      hidden_sizes: [64, 64]
+      activation: tanh
+      last_activation: tanh
+      orthogonal_init: [1.414, 1.414, 0.01]
+  critic:
+    type: mlp
+    network_cfg:
+      hidden_sizes: [64, 64]
+      activation: tanh
+      orthogonal_init: [1.414, 1.414, 1.0]
+
+algorithm:
+  name: ppo
+  cfg:
+    learning_rate: 0.00025
+    n_epochs: 10
+    batch_size: 4096
+    gamma: 0.99
+    gae_lambda: 0.95
+    clip_coef: 0.2
+    ent_coef: 0.01
+    vf_coef: 0.5
+    max_grad_norm: 0.5
diff --git a/embodichain_tasks/embodichain_tasks/rl/basic/__init__.py b/embodichain_tasks/embodichain_tasks/rl/basic/__init__.py
index e69de29bb..ed913c16f 100644
--- a/embodichain_tasks/embodichain_tasks/rl/basic/__init__.py
+++ b/embodichain_tasks/embodichain_tasks/rl/basic/__init__.py
@@ -0,0 +1,24 @@
+# ----------------------------------------------------------------------------
+# 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.
+# ----------------------------------------------------------------------------
+
+"""Basic reinforcement-learning tasks."""
+
+from __future__ import annotations
+
+from .cart_pole import CartPoleEnv
+from .point_mass import PointMassEnv
+
+__all__ = ["CartPoleEnv", "PointMassEnv"]
diff --git a/embodichain_tasks/embodichain_tasks/rl/basic/point_mass.py b/embodichain_tasks/embodichain_tasks/rl/basic/point_mass.py
new file mode 100644
index 000000000..0e018fe64
--- /dev/null
+++ b/embodichain_tasks/embodichain_tasks/rl/basic/point_mass.py
@@ -0,0 +1,296 @@
+# ----------------------------------------------------------------------------
+# 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.
+# ----------------------------------------------------------------------------
+
+"""Differentiable two-dimensional point-mass navigation task."""
+
+from __future__ import annotations
+
+from collections.abc import Mapping, Sequence
+from typing import Any
+
+import gymnasium as gym
+import numpy as np
+import torch
+
+from embodichain.learning.rl.env import register_learning_env
+
+__all__ = ["PointMassEnv"]
+
+
+@register_learning_env("PointMassRL", override=True)
+class PointMassEnv:
+    """Navigate a damped point mass to a goal while avoiding two obstacles.
+
+    The implementation is always differentiable. Standard collectors execute
+    it under ``torch.no_grad()``, while APG retains the path from policy action
+    through dynamics and reward.
+    """
+
+    observation_dim = 14
+    action_dim = 2
+
+    def __init__(
+        self,
+        num_envs: int = 64,
+        device: torch.device | str = "cpu",
+        *,
+        max_episode_steps: int = 100,
+        workspace_half: float = 1.0,
+        mass: float = 1.0,
+        force_scale: float = 5.0,
+        linear_damping: float = 5.0,
+        dt: float = 1.0 / 60.0,
+        success_threshold: float = 0.03,
+        obstacle_radius_range: Sequence[float] = (0.08, 0.15),
+        obstacle_min_goal_distance: float = 0.25,
+    ) -> None:
+        if num_envs <= 0:
+            raise ValueError("num_envs must be positive.")
+        if max_episode_steps <= 0:
+            raise ValueError("max_episode_steps must be positive.")
+        if len(obstacle_radius_range) != 2:
+            raise ValueError("obstacle_radius_range must contain two values.")
+        radius_min, radius_max = map(float, obstacle_radius_range)
+        if radius_min <= 0 or radius_max < radius_min:
+            raise ValueError("Invalid obstacle_radius_range.")
+
+        self.num_envs = int(num_envs)
+        self.device = torch.device(device)
+        self.max_episode_steps = int(max_episode_steps)
+        self.workspace_half = float(workspace_half)
+        self.mass = float(mass)
+        self.force_scale = float(force_scale)
+        self.linear_damping = float(linear_damping)
+        self.dt = float(dt)
+        self.success_threshold = float(success_threshold)
+        self.obstacle_radius_range = (radius_min, radius_max)
+        self.obstacle_min_goal_distance = float(obstacle_min_goal_distance)
+
+        self.single_observation_space = gym.spaces.Box(
+            low=-np.inf,
+            high=np.inf,
+            shape=(self.observation_dim,),
+            dtype=np.float32,
+        )
+        self.single_action_space = gym.spaces.Box(
+            low=-1.0,
+            high=1.0,
+            shape=(self.action_dim,),
+            dtype=np.float32,
+        )
+        self.observation_space = self.single_observation_space
+        self.action_space = self.single_action_space
+
+        self._generator = torch.Generator(device=self.device)
+        self._generator.manual_seed(torch.seed())
+        self.position = torch.zeros(self.num_envs, 2, device=self.device)
+        self.velocity = torch.zeros_like(self.position)
+        self.goal_position = torch.zeros_like(self.position)
+        self.obstacle_position = torch.zeros(self.num_envs, 2, 2, device=self.device)
+        self.obstacle_radius = torch.zeros(self.num_envs, 2, device=self.device)
+        self.last_action = torch.zeros(
+            self.num_envs, self.action_dim, device=self.device
+        )
+        self.episode_step = torch.zeros(
+            self.num_envs, dtype=torch.long, device=self.device
+        )
+        self.path_length = torch.zeros(self.num_envs, device=self.device)
+        self.collision_count = torch.zeros(
+            self.num_envs, dtype=torch.long, device=self.device
+        )
+
+    def reset(
+        self,
+        *,
+        seed: int | None = None,
+        options: Mapping[str, Any] | None = None,
+    ) -> tuple[torch.Tensor, dict[str, Any]]:
+        """Reset all environments, or selected IDs from ``options``."""
+        if seed is not None:
+            self._generator.manual_seed(int(seed))
+        env_ids = None if options is None else options.get("env_ids")
+        ids = self._normalize_env_ids(env_ids)
+        self._reset_ids(ids)
+        return self._observation(), {}
+
+    def step(self, action: torch.Tensor) -> tuple[
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+        dict[str, Any],
+    ]:
+        """Advance the differentiable dynamics and auto-reset finished rows."""
+        action = torch.as_tensor(action, dtype=torch.float32, device=self.device).clamp(
+            -1.0, 1.0
+        )
+        if tuple(action.shape) != (self.num_envs, self.action_dim):
+            raise ValueError(
+                f"Expected action shape {(self.num_envs, self.action_dim)}, "
+                f"got {tuple(action.shape)}."
+            )
+
+        previous_position = self.position
+        damping = 1.0 - self.linear_damping * self.dt
+        self.velocity = (
+            self.velocity * damping + action * self.force_scale / self.mass * self.dt
+        )
+        self.position = (self.position + self.velocity * self.dt).clamp(
+            -self.workspace_half, self.workspace_half
+        )
+        self.episode_step = self.episode_step + 1
+        self.path_length = (
+            self.path_length + (self.position - previous_position).norm(dim=-1).detach()
+        )
+
+        center_distance = (self.position[:, None, :] - self.obstacle_position).norm(
+            dim=-1
+        )
+        penetration = torch.clamp(self.obstacle_radius - center_distance, min=0.0)
+        collided = (penetration > 0).any(dim=-1)
+        self.collision_count = self.collision_count + collided.long()
+
+        distance = (self.position - self.goal_position).norm(dim=-1)
+        reward = (
+            -distance
+            + 0.5 * torch.exp(-(distance.square()) / (2.0 * 0.05**2))
+            - 2.0 * penetration.square().sum(dim=-1)
+            - 0.01 * self.velocity.square().sum(dim=-1)
+            - 0.001 * (action - self.last_action).square().sum(dim=-1)
+        )
+        terminated = distance < self.success_threshold
+        truncated = self.episode_step >= self.max_episode_steps
+        done = terminated | truncated
+
+        terminal_distance = distance.detach()
+        terminal_path_length = self.path_length.detach().clone()
+        terminal_collisions = self.collision_count.detach().clone()
+        info = {
+            "success": terminated.detach(),
+            "metrics": {
+                "final_distance": terminal_distance,
+                "path_length": terminal_path_length,
+                "collision_count": terminal_collisions,
+            },
+        }
+
+        self.last_action = action.detach()
+        if bool(done.any()):
+            self._reset_ids(torch.nonzero(done, as_tuple=False).squeeze(-1))
+        return self._observation(), reward, terminated, truncated, info
+
+    def detach_state(self) -> torch.Tensor:
+        """Detach dynamic state at a truncated-backpropagation boundary."""
+        self.position = self.position.detach()
+        self.velocity = self.velocity.detach()
+        self.last_action = self.last_action.detach()
+        return self._observation()
+
+    def close(self) -> None:
+        """No-op; the tensor-only task owns no external resources."""
+
+    def _observation(self) -> torch.Tensor:
+        return torch.cat(
+            (
+                self.position,
+                self.velocity,
+                self.goal_position,
+                self.last_action,
+                self.obstacle_position.reshape(self.num_envs, -1),
+                self.obstacle_radius,
+            ),
+            dim=-1,
+        )
+
+    def _normalize_env_ids(
+        self, env_ids: torch.Tensor | Sequence[int] | None
+    ) -> torch.Tensor:
+        if env_ids is None:
+            return torch.arange(self.num_envs, device=self.device)
+        return torch.as_tensor(env_ids, dtype=torch.long, device=self.device).reshape(
+            -1
+        )
+
+    def _reset_ids(self, ids: torch.Tensor) -> None:
+        if ids.numel() == 0:
+            return
+        count = int(ids.numel())
+        goal = self._uniform(count, 2, scale=0.7 * self.workspace_half)
+        start = self._uniform(count, 2, scale=0.6 * self.workspace_half)
+        obstacle_position = self._sample_obstacles(goal)
+        radius_min, radius_max = self.obstacle_radius_range
+        obstacle_radius = radius_min + (radius_max - radius_min) * torch.rand(
+            count, 2, generator=self._generator, device=self.device
+        )
+
+        mask = torch.zeros(self.num_envs, dtype=torch.bool, device=self.device)
+        mask[ids] = True
+        mask_2d = mask[:, None]
+        mask_3d = mask[:, None, None]
+
+        full_goal = self.goal_position.detach().clone()
+        full_goal[ids] = goal
+        full_start = self.position.detach().clone()
+        full_start[ids] = start
+        full_obstacles = self.obstacle_position.detach().clone()
+        full_obstacles[ids] = obstacle_position
+        full_radii = self.obstacle_radius.detach().clone()
+        full_radii[ids] = obstacle_radius
+
+        self.position = torch.where(mask_2d, full_start, self.position)
+        self.velocity = torch.where(
+            mask_2d, torch.zeros_like(self.velocity), self.velocity
+        )
+        self.goal_position = torch.where(mask_2d, full_goal, self.goal_position)
+        self.obstacle_position = torch.where(
+            mask_3d, full_obstacles, self.obstacle_position
+        )
+        self.obstacle_radius = torch.where(mask_2d, full_radii, self.obstacle_radius)
+        self.last_action = torch.where(
+            mask_2d, torch.zeros_like(self.last_action), self.last_action
+        )
+        self.episode_step = torch.where(
+            mask, torch.zeros_like(self.episode_step), self.episode_step
+        )
+        self.path_length = torch.where(
+            mask, torch.zeros_like(self.path_length), self.path_length
+        )
+        self.collision_count = torch.where(
+            mask, torch.zeros_like(self.collision_count), self.collision_count
+        )
+
+    def _sample_obstacles(self, goal: torch.Tensor) -> torch.Tensor:
+        count = goal.shape[0]
+        result = self._uniform(count, 2, 2, scale=0.8 * self.workspace_half)
+        valid = (result - goal[:, None, :]).norm(
+            dim=-1
+        ) > self.obstacle_min_goal_distance
+        for _ in range(50):
+            if bool(valid.all()):
+                break
+            candidates = self._uniform(count, 2, 2, scale=0.8 * self.workspace_half)
+            result = torch.where(valid[:, :, None], result, candidates)
+            valid = (result - goal[:, None, :]).norm(
+                dim=-1
+            ) > self.obstacle_min_goal_distance
+        fallback = -0.5 * goal[:, None, :].expand(-1, 2, -1)
+        return torch.where(valid[:, :, None], result, fallback)
+
+    def _uniform(self, *shape: int, scale: float) -> torch.Tensor:
+        return (
+            torch.rand(*shape, generator=self._generator, device=self.device) * 2.0
+            - 1.0
+        ) * scale
diff --git a/scripts/benchmark/rl/algorithms/apg.yaml b/scripts/benchmark/rl/algorithms/apg.yaml
new file mode 100644
index 000000000..8f0a579cf
--- /dev/null
+++ b/scripts/benchmark/rl/algorithms/apg.yaml
@@ -0,0 +1,19 @@
+name: apg
+config:
+  policy:
+    name: actor_only
+    initial_log_std: -2.0
+    actor:
+      type: mlp
+      network_cfg:
+        hidden_sizes: [64, 64]
+        activation: tanh
+        last_activation: tanh
+        orthogonal_init: [1.414, 1.414, 0.01]
+  algorithm:
+    name: apg
+    cfg:
+      learning_rate: 0.00025
+      gamma: 0.99
+      max_grad_norm: 0.5
+      ent_coef: 0.0
diff --git a/scripts/benchmark/rl/runtime.py b/scripts/benchmark/rl/runtime.py
index a3b4b2055..351bdb819 100644
--- a/scripts/benchmark/rl/runtime.py
+++ b/scripts/benchmark/rl/runtime.py
@@ -28,11 +28,18 @@
 from torch.utils.tensorboard import SummaryWriter
 
 from embodichain.learning.rl.algo import build_algo
+from embodichain.learning.rl.differentiable_trainer import (
+    DifferentiableTrainer,
+    DifferentiableTrainerCfg,
+)
+from embodichain.learning.rl.env import build_learning_env
+from embodichain.learning.rl.evaluation import evaluate_episodes
 from embodichain.learning.rl.models import build_mlp_from_cfg, build_policy
+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.lab.gym.envs.managers.cfg import EventCfg
-from embodichain.lab.gym.utils.registration import build_env
+from embodichain.lab.gym.utils.registration import build_env, discover_task_packages
 from embodichain.lab.gym.utils.gym_utils import get_manager_modules, config_to_cfg
 from embodichain.lab.sim import SimulationManagerCfg
 from embodichain.utils.module_utils import find_function_from_modules
@@ -184,8 +191,13 @@ def build_policy_from_env(policy_block: dict[str, Any], env, device: torch.devic
     sample_obs, _ = env.reset()
     sample_obs_td = dict_to_tensordict(sample_obs, device)
     obs_dim = flatten_dict_observation(sample_obs_td).shape[-1]
-    flat_obs_space = env.flattened_observation_space
-    env_action_dim = env.action_space.shape[-1]
+    flat_obs_space = getattr(env, "flattened_observation_space", None)
+    if flat_obs_space is None:
+        flat_obs_space = env.single_observation_space
+    action_space = getattr(env, "action_space", None)
+    if action_space is None:
+        action_space = env.single_action_space
+    env_action_dim = action_space.shape[-1]
 
     policy_name = policy_block["name"].lower()
     if policy_name == "actor_critic":
@@ -194,7 +206,7 @@ def build_policy_from_env(policy_block: dict[str, Any], env, device: torch.devic
         return build_policy(
             policy_block,
             flat_obs_space,
-            env.action_space,
+            action_space,
             device,
             actor=actor,
             critic=critic,
@@ -204,11 +216,127 @@ def build_policy_from_env(policy_block: dict[str, Any], env, device: torch.devic
         return build_policy(
             policy_block,
             flat_obs_space,
-            env.action_space,
+            action_space,
             device,
             actor=actor,
         )
-    return build_policy(policy_block, flat_obs_space, env.action_space, device)
+    return build_policy(policy_block, flat_obs_space, action_space, device)
+
+
+def _train_learning_config(
+    cfg_json: dict[str, Any],
+    output_dir: str | Path,
+) -> dict[str, Any]:
+    trainer_cfg = deepcopy(cfg_json["trainer"])
+    policy_block = deepcopy(cfg_json["policy"])
+    algo_block = deepcopy(cfg_json["algorithm"])
+    device = resolve_device(trainer_cfg.get("device", "cpu"))
+    set_random_seed(int(trainer_cfg.get("seed", 1)), device)
+    discover_task_packages()
+
+    env_block = trainer_cfg["learning_env"]
+    env_name = env_block if isinstance(env_block, str) else env_block["name"]
+    env_kwargs = {} if isinstance(env_block, str) else dict(env_block.get("cfg", {}))
+    num_envs = int(trainer_cfg.get("num_envs", 64))
+    env = build_learning_env(env_name, num_envs=num_envs, device=device, **env_kwargs)
+    eval_env = None
+    if trainer_cfg.get("enable_eval", True):
+        eval_env = build_learning_env(
+            env_name,
+            num_envs=int(trainer_cfg.get("num_eval_envs", 8)),
+            device=device,
+            **env_kwargs,
+        )
+
+    output_root = Path(output_dir)
+    checkpoint_dir = output_root / "checkpoints"
+    checkpoint_dir.mkdir(parents=True, exist_ok=True)
+    writer = SummaryWriter(str(output_root / "logs"))
+    try:
+        policy = build_policy_from_env(policy_block, env, device)
+        if "initial_log_std" in policy_block and hasattr(policy, "log_std"):
+            with torch.no_grad():
+                policy.log_std.fill_(float(policy_block["initial_log_std"]))
+        algorithm = build_algo(algo_block["name"], algo_block["cfg"], policy, device)
+        iterations = int(trainer_cfg.get("iterations", 1))
+        eval_freq = int(trainer_cfg.get("eval_freq", 0))
+        trainer_class = get_trainer_class(algorithm)
+        best_eval_metric = trainer_cfg.get("best_eval_metric", "eval/avg_reward")
+        best_eval_mode = trainer_cfg.get("best_eval_mode", "max")
+        if trainer_class is DifferentiableTrainer:
+            segment_length = int(trainer_cfg.get("segment_length", 25))
+            update_horizon = int(trainer_cfg.get("update_horizon", segment_length))
+            trainer = DifferentiableTrainer(
+                cfg=DifferentiableTrainerCfg(
+                    segment_length=segment_length,
+                    update_horizon=update_horizon,
+                    deterministic_actions=bool(
+                        trainer_cfg.get("deterministic_actions", False)
+                    ),
+                    checkpoint_dir=str(checkpoint_dir),
+                    experiment_name=str(trainer_cfg.get("exp_name", "benchmark_run")),
+                    save_frequency_updates=int(
+                        trainer_cfg.get("save_frequency_updates", 0)
+                    ),
+                    eval_frequency_steps=eval_freq,
+                    num_eval_episodes=int(trainer_cfg.get("num_eval_episodes", 5)),
+                    eval_seed=trainer_cfg.get("eval_seed"),
+                    best_eval_metric=best_eval_metric,
+                    best_eval_mode=best_eval_mode,
+                ),
+                env=env,
+                policy=policy,
+                algorithm=algorithm,
+                writer=writer,
+                eval_env=eval_env,
+            )
+            total_steps = iterations * update_horizon * num_envs
+        else:
+            buffer_size = int(trainer_cfg.get("buffer_size", 256))
+            trainer = Trainer(
+                policy=policy,
+                env=env,
+                algorithm=algorithm,
+                buffer_size=buffer_size,
+                batch_size=int(algorithm.cfg.batch_size),
+                writer=writer,
+                eval_freq=eval_freq,
+                save_freq=int(trainer_cfg.get("save_freq", 0)),
+                checkpoint_dir=str(checkpoint_dir),
+                exp_name=str(trainer_cfg.get("exp_name", "benchmark_run")),
+                use_wandb=False,
+                eval_env=eval_env,
+                num_eval_episodes=int(trainer_cfg.get("num_eval_episodes", 5)),
+                eval_seed=trainer_cfg.get("eval_seed"),
+                best_eval_metric=best_eval_metric,
+                best_eval_mode=best_eval_mode,
+            )
+            total_steps = iterations * buffer_size * num_envs
+        start_time = time.perf_counter()
+        summary = trainer.train(total_steps)
+        wall_time = time.perf_counter() - start_time
+        checkpoint_path = trainer.save_checkpoint()
+    finally:
+        writer.close()
+        if eval_env is not None:
+            eval_env.close()
+        env.close()
+
+    peak_gpu_memory_mb = (
+        torch.cuda.max_memory_allocated(device=device) / (1024.0 * 1024.0)
+        if device.type == "cuda"
+        else 0.0
+    )
+    summary.update(
+        {
+            "checkpoint_path": checkpoint_path,
+            "output_dir": str(output_root),
+            "wall_time_sec": float(wall_time),
+            "training_fps": float(total_steps / max(wall_time, 1e-6)),
+            "peak_gpu_memory_mb": float(peak_gpu_memory_mb),
+        }
+    )
+    return summary
 
 
 def train_with_config(
@@ -216,6 +344,8 @@ def train_with_config(
     output_dir: str | Path,
 ) -> dict[str, Any]:
     """Train an RL configuration and return a structured summary."""
+    if "learning_env" in cfg_json["trainer"]:
+        return _train_learning_config(cfg_json, output_dir)
     trainer_cfg = deepcopy(cfg_json["trainer"])
     policy_block = deepcopy(cfg_json["policy"])
     algo_block = deepcopy(cfg_json["algorithm"])
@@ -321,105 +451,79 @@ def evaluate_checkpoint(
     policy_block = deepcopy(cfg_json["policy"])
 
     device = resolve_device(trainer_cfg.get("device", "cpu"))
-    gym_config_data, gym_env_cfg = _build_env_cfg(
-        gym_config_path=trainer_cfg["gym_config"],
-        num_envs=num_envs if num_envs is not None else trainer_cfg.get("num_eval_envs"),
-        headless=True,
-        device=device,
-        gpu_id=int(trainer_cfg.get("gpu_id", 0)),
-    )
+    is_learning_env = "learning_env" in trainer_cfg
+    if is_learning_env:
+        discover_task_packages()
+        env_block = trainer_cfg["learning_env"]
+        env_name = env_block if isinstance(env_block, str) else env_block["name"]
+        env_kwargs = (
+            {} if isinstance(env_block, str) else dict(env_block.get("cfg", {}))
+        )
+        eval_num_envs = int(
+            num_envs if num_envs is not None else trainer_cfg.get("num_eval_envs", 4)
+        )
+    else:
+        gym_config_data, gym_env_cfg = _build_env_cfg(
+            gym_config_path=trainer_cfg["gym_config"],
+            num_envs=(
+                num_envs if num_envs is not None else trainer_cfg.get("num_eval_envs")
+            ),
+            headless=True,
+            device=device,
+            gpu_id=int(trainer_cfg.get("gpu_id", 0)),
+        )
     env = None
     try:
-        env = build_env(gym_config_data["id"], base_env_cfg=gym_env_cfg)
+        if is_learning_env:
+            env = build_learning_env(
+                env_name,
+                num_envs=eval_num_envs,
+                device=device,
+                **env_kwargs,
+            )
+        else:
+            env = build_env(gym_config_data["id"], base_env_cfg=gym_env_cfg)
         policy = build_policy_from_env(policy_block, env, device)
         eval_rollout_buffer = None
         if hasattr(env, "set_rollout_buffer"):
             eval_rollout_buffer = _allocate_eval_rollout_buffer(env, policy, device)
+            env.set_rollout_buffer(eval_rollout_buffer)
 
         checkpoint = torch.load(checkpoint_path, map_location=device)
         policy.load_state_dict(checkpoint["policy"])
         policy.eval()
 
-        target_episodes = int(num_episodes)
-        completed = 0
-        cumulative_reward = torch.zeros(
-            env.num_envs, dtype=torch.float32, device=device
-        )
-        step_count = torch.zeros(env.num_envs, dtype=torch.int32, device=device)
-
-        returns: list[float] = []
-        lengths: list[int] = []
-        successes: list[float] = []
-        metric_values: dict[str, list[float]] = {}
-        env_step_count = 0
-        env_step_time = 0.0
-
-        if eval_rollout_buffer is not None:
-            env.set_rollout_buffer(eval_rollout_buffer)
-        obs, _ = env.reset()
-        while completed < target_episodes:
-            flat_obs = flatten_dict_observation(obs)
-            action_td = TensorDict(
-                {"obs": flat_obs},
-                batch_size=[env.num_envs],
-                device=device,
-            )
-            action_td = policy.get_action(action_td, deterministic=True)
-            action_manager = getattr(env, "action_manager", None)
-            if action_manager is None:
-                action_in = action_td["action"]
-            else:
-                action_in = action_manager.convert_policy_action_to_env_action(
-                    action_td["action"]
-                )
-
+        def on_step(_: dict[str, Any]) -> None:
             if eval_rollout_buffer is not None:
                 _compact_eval_rollout_buffer(env, eval_rollout_buffer)
-                eval_rollout_buffer["action"][:, env.current_rollout_step].copy_(
-                    action_td["action"]
-                )
-            step_start = time.perf_counter()
-            obs, reward, terminated, truncated, info = env.step(action_in)
-            env_step_time += time.perf_counter() - step_start
-            env_step_count += env.num_envs
-
-            done = terminated | truncated
-            cumulative_reward += reward.float()
-            step_count += 1
-
-            newly_done = done.nonzero(as_tuple=False).squeeze(-1)
-            for env_id in newly_done.tolist():
-                if completed >= target_episodes:
-                    break
-                returns.append(float(cumulative_reward[env_id].item()))
-                lengths.append(int(step_count[env_id].item()))
-                if "success" in info:
-                    successes.append(float(info["success"][env_id].item()))
-                if "metrics" in info:
-                    for key, value in info["metrics"].items():
-                        metric_values.setdefault(key, []).append(
-                            float(value[env_id].item())
-                        )
-                cumulative_reward[env_id] = 0.0
-                step_count[env_id] = 0
-                completed += 1
+
+        start_time = time.perf_counter()
+        unified_metrics = evaluate_episodes(
+            policy=policy,
+            env=env,
+            num_episodes=int(num_episodes),
+            device=device,
+            seed=trainer_cfg.get("eval_seed"),
+            on_step=on_step,
+        )
+        elapsed = time.perf_counter() - start_time
+        avg_length = unified_metrics["eval/avg_length"]
+        return {
+            "num_episodes": int(num_episodes),
+            "avg_reward": unified_metrics["eval/avg_reward"],
+            "avg_episode_length": avg_length,
+            "success_rate": unified_metrics["eval/success_rate"],
+            "environment_fps": float(num_episodes * avg_length / max(elapsed, 1e-6)),
+            "metrics": {
+                key.removeprefix("eval/metrics/"): value
+                for key, value in unified_metrics.items()
+                if key.startswith("eval/metrics/")
+            },
+        }
     finally:
         if env is not None:
             env.close()
 
-    return {
-        "num_episodes": completed,
-        "avg_reward": float(np.mean(returns)) if returns else float("nan"),
-        "avg_episode_length": float(np.mean(lengths)) if lengths else float("nan"),
-        "success_rate": float(np.mean(successes)) if successes else float("nan"),
-        "environment_fps": float(env_step_count / max(env_step_time, 1e-6)),
-        "metrics": {
-            key: float(np.mean(values))
-            for key, values in metric_values.items()
-            if values
-        },
-    }
-
 
 def dump_json(data: dict[str, Any], path: str | Path) -> Path:
     """Write a JSON artifact to disk."""
diff --git a/scripts/benchmark/rl/suites/point_mass.yaml b/scripts/benchmark/rl/suites/point_mass.yaml
new file mode 100644
index 000000000..cae7e85d8
--- /dev/null
+++ b/scripts/benchmark/rl/suites/point_mass.yaml
@@ -0,0 +1,21 @@
+tasks:
+  - point_mass
+algorithms:
+  - apg
+  - ppo
+seeds:
+  - 0
+  - 1
+  - 2
+protocol:
+  device: cuda:0
+  headless: true
+  iterations: 500
+  buffer_size: 256
+  num_envs: 64
+  num_eval_envs: 32
+  evaluation_interval: 32000
+  evaluation_episodes: 128
+  threshold_sustain_count: 3
+  final_eval_window: 3
+  save_interval: 262144
diff --git a/scripts/benchmark/rl/tasks/point_mass.yaml b/scripts/benchmark/rl/tasks/point_mass.yaml
new file mode 100644
index 000000000..a9b6eb1ce
--- /dev/null
+++ b/scripts/benchmark/rl/tasks/point_mass.yaml
@@ -0,0 +1,24 @@
+name: point_mass
+env_id: PointMassRL
+success_threshold: 0.95
+base_config:
+  trainer:
+    learning_env:
+      name: PointMassRL
+      cfg:
+        max_episode_steps: 100
+        success_threshold: 0.03
+    exp_name: point_mass
+    device: cuda:0
+    headless: true
+    num_envs: 64
+    iterations: 500
+    buffer_size: 256
+    segment_length: 100
+    update_horizon: 100
+    enable_eval: true
+    num_eval_envs: 32
+    num_eval_episodes: 128
+    best_eval_metric: eval/success_rate
+    best_eval_mode: max
+    use_wandb: false
diff --git a/tests/learning/test_apg.py b/tests/learning/test_apg.py
index eebfd8d6b..3cd509853 100644
--- a/tests/learning/test_apg.py
+++ b/tests/learning/test_apg.py
@@ -109,14 +109,17 @@ def _objective_at(weight: float) -> float:
     return segmented_discounted_return(rollout, gamma=0.9).mean().detach().item()
 
 
-def test_apg_is_not_registered_with_standard_trainer() -> None:
-    assert "apg" not in get_registered_algo_names()
-    with pytest.raises(ValueError, match="differentiable rollouts"):
+def test_apg_is_registered_with_differentiable_rollout_kind() -> None:
+    assert "apg" in get_registered_algo_names()
+    algorithm = build_algo("apg", {}, _make_policy(), torch.device("cpu"))
+    assert algorithm.rollout_kind.value == "differentiable"
+    with pytest.raises(ValueError, match="do not support distributed"):
         build_algo(
             "apg",
             {},
             _make_policy(),
             torch.device("cpu"),
+            distributed=True,
         )
 
 
diff --git a/tests/learning/test_differentiable_env.py b/tests/learning/test_differentiable_env.py
index 5d109cc89..bca7de76d 100644
--- a/tests/learning/test_differentiable_env.py
+++ b/tests/learning/test_differentiable_env.py
@@ -73,6 +73,9 @@ def detach_state(self) -> torch.Tensor:
         self._state = self._state.detach()
         return self._state
 
+    def close(self) -> None:
+        return None
+
 
 def test_structural_protocol_accepts_differentiable_environment() -> None:
     env = _MockDifferentiableEnv()
diff --git a/tests/learning/test_differentiable_trainer.py b/tests/learning/test_differentiable_trainer.py
index 3e9cf38fd..ea601d851 100644
--- a/tests/learning/test_differentiable_trainer.py
+++ b/tests/learning/test_differentiable_trainer.py
@@ -69,6 +69,9 @@ def detach_state(self) -> torch.Tensor:
         self._state = self._state.detach()
         return self._state
 
+    def close(self) -> None:
+        return None
+
 
 def _make_components(
     ent_coef: float = 0.0,
diff --git a/tests/learning/test_evaluation.py b/tests/learning/test_evaluation.py
new file mode 100644
index 000000000..9afa8560b
--- /dev/null
+++ b/tests/learning/test_evaluation.py
@@ -0,0 +1,90 @@
+# ----------------------------------------------------------------------------
+# 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.
+# ----------------------------------------------------------------------------
+
+from __future__ import annotations
+
+import gymnasium as gym
+import torch
+
+from embodichain.learning.rl.evaluation import evaluate_episodes
+
+
+class _DeterministicPolicy(torch.nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        self.anchor = torch.nn.Parameter(torch.zeros(()))
+
+    def get_action(self, tensordict, deterministic: bool = False):
+        assert deterministic
+        tensordict["action"] = torch.zeros(
+            tensordict.batch_size[0], 1, device=tensordict.device
+        )
+        return tensordict
+
+
+class _AsyncAutoResetEnv:
+    num_envs = 3
+    device = torch.device("cpu")
+    single_observation_space = gym.spaces.Box(-1.0, 1.0, (1,))
+    single_action_space = gym.spaces.Box(-1.0, 1.0, (1,))
+
+    def __init__(self) -> None:
+        self.steps = torch.zeros(3, dtype=torch.long)
+        self.horizons = torch.tensor([1, 2, 3])
+        self.last_seed = None
+
+    def reset(self, *, seed=None, options=None):
+        self.last_seed = seed
+        self.steps.zero_()
+        return self.steps[:, None].float(), {}
+
+    def step(self, action):
+        self.steps += 1
+        done = self.steps >= self.horizons
+        terminal_steps = self.steps.clone()
+        info = {
+            "success": done.clone(),
+            "metrics": {"terminal_step": terminal_steps.float()},
+        }
+        self.steps[done] = 0
+        return (
+            self.steps[:, None].float(),
+            torch.ones(3),
+            done,
+            torch.zeros(3, dtype=torch.bool),
+            info,
+        )
+
+
+def test_evaluate_episodes_counts_actual_completions_and_restores_mode() -> None:
+    policy = _DeterministicPolicy()
+    policy.train()
+    env = _AsyncAutoResetEnv()
+
+    result = evaluate_episodes(
+        policy=policy,
+        env=env,
+        num_episodes=4,
+        device="cpu",
+        seed=123,
+    )
+
+    assert policy.training
+    assert env.last_seed == 123
+    assert result["eval/avg_reward"] == 1.25
+    assert result["eval/avg_length"] == 1.25
+    assert result["eval/success_rate"] == 1.0
+    assert result["eval/metrics/terminal_step"] == 1.25
diff --git a/tests/learning/test_point_mass.py b/tests/learning/test_point_mass.py
new file mode 100644
index 000000000..7850b6f3f
--- /dev/null
+++ b/tests/learning/test_point_mass.py
@@ -0,0 +1,226 @@
+# ----------------------------------------------------------------------------
+# 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.
+# ----------------------------------------------------------------------------
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import pytest
+import torch
+import torch.nn as nn
+
+from embodichain.learning.rl.algo import PPO, PPOCfg
+from embodichain.learning.rl.buffer import RolloutBuffer
+from embodichain.learning.rl.collector import SyncCollector
+from embodichain.learning.rl.env import (
+    DifferentiableVecEnv,
+    build_learning_env,
+)
+from embodichain.learning.rl.models import ActorCritic
+from embodichain.learning.rl.train import train_from_config
+from embodichain.learning.rl.utils.trainer import Trainer
+from embodichain_tasks.rl.basic.point_mass import PointMassEnv
+
+
+def test_point_mass_is_registered_differentiable_environment() -> None:
+    env = build_learning_env("PointMassRL", num_envs=4, device="cpu")
+
+    assert isinstance(env, PointMassEnv)
+    assert isinstance(env, DifferentiableVecEnv)
+    assert env.single_observation_space.shape == (14,)
+    assert env.single_action_space.shape == (2,)
+
+
+def test_point_mass_reset_seed_is_reproducible() -> None:
+    env = PointMassEnv(num_envs=8)
+
+    first, _ = env.reset(seed=17)
+    env.step(torch.ones(8, 2))
+    second, _ = env.reset(seed=17)
+
+    torch.testing.assert_close(first, second)
+
+
+def test_point_mass_step_preserves_action_gradient() -> None:
+    env = PointMassEnv(num_envs=4)
+    env.reset(seed=3)
+    action = torch.randn(4, 2, requires_grad=True)
+
+    next_observation, reward, _, _, _ = env.step(action)
+    reward.sum().backward()
+
+    assert next_observation.grad_fn is not None
+    assert action.grad is not None
+    assert torch.isfinite(action.grad).all()
+    assert float(action.grad.norm()) > 0.0
+
+
+def test_point_mass_auto_reset_keeps_terminal_metrics() -> None:
+    env = PointMassEnv(num_envs=2)
+    env.reset(seed=5)
+    env.velocity.zero_()
+    env.goal_position[0] = env.position[0]
+
+    next_observation, _, terminated, truncated, info = env.step(torch.zeros(2, 2))
+
+    assert bool(terminated[0])
+    assert not bool(truncated[0])
+    assert float(info["metrics"]["final_distance"][0]) < env.success_threshold
+    assert env.episode_step[0] == 0
+    assert next_observation.shape == (2, 14)
+
+
+def test_point_mass_detach_state_keeps_values() -> None:
+    env = PointMassEnv(num_envs=3)
+    env.reset(seed=9)
+    action = torch.randn(3, 2, requires_grad=True)
+    observation, *_ = env.step(action)
+
+    detached = env.detach_state()
+
+    torch.testing.assert_close(observation, detached)
+    assert detached.grad_fn is None
+
+
+@pytest.mark.parametrize(
+    ("algorithm_name", "policy_name"),
+    [("apg", "actor_only"), ("ppo", "actor_critic")],
+)
+def test_unified_train_entry_runs_apg_and_ppo(
+    tmp_path, monkeypatch, algorithm_name: str, policy_name: str
+) -> None:
+    monkeypatch.chdir(tmp_path)
+    policy = {
+        "name": policy_name,
+        "actor": {
+            "type": "mlp",
+            "network_cfg": {"hidden_sizes": [8], "activation": "tanh"},
+        },
+    }
+    algorithm_cfg = {
+        "learning_rate": 1e-3,
+        "batch_size": 8,
+        "gamma": 0.99,
+        "max_grad_norm": 1.0,
+    }
+    if policy_name == "actor_critic":
+        policy["critic"] = {
+            "type": "mlp",
+            "network_cfg": {"hidden_sizes": [8], "activation": "tanh"},
+        }
+        algorithm_cfg.update(
+            {
+                "n_epochs": 1,
+                "gae_lambda": 0.95,
+                "clip_coef": 0.2,
+                "ent_coef": 0.0,
+                "vf_coef": 0.5,
+            }
+        )
+    config = {
+        "trainer": {
+            "exp_name": f"smoke_{algorithm_name}",
+            "learning_env": {
+                "name": "PointMassRL",
+                "cfg": {"max_episode_steps": 4},
+            },
+            "device": "cpu",
+            "num_envs": 2,
+            "iterations": 1,
+            "buffer_size": 4,
+            "segment_length": 2,
+            "update_horizon": 4,
+            "enable_eval": False,
+            "use_wandb": False,
+        },
+        "policy": policy,
+        "algorithm": {"name": algorithm_name, "cfg": algorithm_cfg},
+    }
+    config_path = tmp_path / f"{algorithm_name}.json"
+    config_path.write_text(json.dumps(config), encoding="utf-8")
+
+    summary = train_from_config(str(config_path))
+
+    assert summary["global_step"] == 8
+    assert summary["latest_checkpoint_path"] is not None
+
+
+def test_sync_collector_accepts_tensor_point_mass_observations() -> None:
+    env = PointMassEnv(num_envs=2, max_episode_steps=4)
+    actor = nn.Linear(14, 2)
+    critic = nn.Linear(14, 1)
+    policy = ActorCritic(14, 2, env.device, actor=actor, critic=critic)
+    buffer = RolloutBuffer(
+        num_envs=2,
+        rollout_len=3,
+        obs_dim=14,
+        action_dim=2,
+        device=env.device,
+    )
+    collector = SyncCollector(env=env, policy=policy, device=env.device)
+    rollout = collector.collect(num_steps=3, rollout=buffer.start_rollout())
+
+    assert rollout["obs"].shape == (2, 4, 14)
+    assert torch.isfinite(rollout["reward"][:, :3]).all()
+    assert torch.isfinite(rollout["action"][:, :3]).all()
+
+
+def test_trainer_saves_best_checkpoint_from_eval(tmp_path: Path) -> None:
+    env = PointMassEnv(num_envs=2, max_episode_steps=4)
+    eval_env = PointMassEnv(num_envs=2, max_episode_steps=4)
+    actor = nn.Linear(14, 2)
+    critic = nn.Linear(14, 1)
+    policy = ActorCritic(14, 2, env.device, actor=actor, critic=critic)
+    algorithm = PPO(
+        PPOCfg(
+            device="cpu",
+            learning_rate=1e-3,
+            n_epochs=1,
+            batch_size=8,
+            gamma=0.99,
+            gae_lambda=0.95,
+            clip_coef=0.2,
+            ent_coef=0.0,
+            vf_coef=0.5,
+        ),
+        policy,
+    )
+    trainer = Trainer(
+        policy=policy,
+        env=env,
+        algorithm=algorithm,
+        buffer_size=4,
+        batch_size=8,
+        writer=None,
+        eval_freq=8,
+        save_freq=0,
+        checkpoint_dir=str(tmp_path),
+        exp_name="point_mass_best",
+        use_wandb=False,
+        eval_env=eval_env,
+        num_eval_episodes=2,
+        eval_seed=7,
+        best_eval_metric="eval/avg_reward",
+        best_eval_mode="max",
+    )
+
+    summary = trainer.train(total_timesteps=8)
+
+    assert summary["last_eval_metrics"]
+    assert "eval/avg_reward" in summary["last_eval_metrics"]
+    assert summary["best_checkpoint_path"] is not None
+    assert Path(summary["best_checkpoint_path"]).is_file()
diff --git a/tests/learning/test_routing.py b/tests/learning/test_routing.py
new file mode 100644
index 000000000..72278c3a7
--- /dev/null
+++ b/tests/learning/test_routing.py
@@ -0,0 +1,48 @@
+# ----------------------------------------------------------------------------
+# 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.
+# ----------------------------------------------------------------------------
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+
+from embodichain.learning.rl.algo import APG, APGCfg, PPO, PPOCfg, RolloutKind
+from embodichain.learning.rl.differentiable_trainer import DifferentiableTrainer
+from embodichain.learning.rl.models import ActorCritic, ActorOnly
+from embodichain.learning.rl.routing import get_trainer_class
+from embodichain.learning.rl.utils.trainer import Trainer
+
+
+def test_get_trainer_class_routes_by_rollout_kind() -> None:
+    actor = nn.Linear(2, 1)
+    apg_policy = ActorOnly(2, 1, torch.device("cpu"), actor=actor)
+    ppo_policy = ActorCritic(
+        2,
+        1,
+        torch.device("cpu"),
+        actor=nn.Linear(2, 1),
+        critic=nn.Linear(2, 1),
+    )
+    apg = APG(APGCfg(device="cpu"), apg_policy)
+    ppo = PPO(PPOCfg(device="cpu"), ppo_policy)
+
+    assert apg.rollout_kind is RolloutKind.DIFFERENTIABLE
+    assert ppo.rollout_kind is RolloutKind.STANDARD
+    assert get_trainer_class(apg) is DifferentiableTrainer
+    assert get_trainer_class(ppo) is Trainer
+    assert get_trainer_class(apg) is get_trainer_class(
+        type("Algo", (), {"rollout_kind": "differentiable"})()
+    )

From 3e72d438313f9a018d03f59346172e1d5dbc9025 Mon Sep 17 00:00:00 2001
From: yangchen73 <122090643@link.cuhk.edu.cn>
Date: Sun, 2 Aug 2026 02:53:08 +0800
Subject: [PATCH 2/7] Let optimizer and learning rate configurable

---
 .../topics/rl-learning/rl-learning.md         |   4 +-
 embodichain/learning/rl/algo/__init__.py      |  17 +-
 embodichain/learning/rl/algo/apg.py           |  32 +--
 embodichain/learning/rl/algo/base.py          |  53 ++++-
 embodichain/learning/rl/algo/grpo.py          |   4 +-
 embodichain/learning/rl/algo/ppo.py           |   4 +-
 .../learning/rl/differentiable_trainer.py     |  41 +++-
 .../experimental/newton/train_planar_reach.py |   3 +-
 embodichain/learning/rl/utils/__init__.py     |  24 ++-
 embodichain/learning/rl/utils/config.py       |  32 ++-
 embodichain/learning/rl/utils/optimizer.py    | 195 ++++++++++++++++++
 embodichain/learning/rl/utils/trainer.py      |  19 +-
 .../rl/basic/cart_pole/train_config.json      |   5 +-
 .../rl/basic/cart_pole/train_config.yaml      |   4 +-
 .../rl/basic/cart_pole/train_config_grpo.json |   5 +-
 .../rl/basic/cart_pole/train_config_grpo.yaml |   4 +-
 .../agents/rl/basic/point_mass/train_apg.yaml |   4 +-
 .../agents/rl/basic/point_mass/train_ppo.yaml |   4 +-
 .../agents/rl/push_cube/train_config.json     |   5 +-
 .../rl/push_cube/train_config_grpo.json       |   5 +-
 scripts/benchmark/rl/algorithms/apg.yaml      |   4 +-
 scripts/benchmark/rl/algorithms/grpo.yaml     |   4 +-
 scripts/benchmark/rl/algorithms/ppo.yaml      |   4 +-
 tests/learning/test_apg.py                    |  70 ++++++-
 tests/learning/test_differentiable_trainer.py |  60 +++++-
 tests/learning/test_newton_planar_reach.py    |   3 +-
 tests/learning/test_optimizer.py              | 109 ++++++++++
 tests/learning/test_point_mass.py             |   5 +-
 28 files changed, 654 insertions(+), 69 deletions(-)
 create mode 100644 embodichain/learning/rl/utils/optimizer.py
 create mode 100644 tests/learning/test_optimizer.py

diff --git a/agent_context/topics/rl-learning/rl-learning.md b/agent_context/topics/rl-learning/rl-learning.md
index 65cc28f85..bf53f73f1 100644
--- a/agent_context/topics/rl-learning/rl-learning.md
+++ b/agent_context/topics/rl-learning/rl-learning.md
@@ -50,7 +50,7 @@ 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, trainer counters, and best-evaluation
+- 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.
@@ -125,7 +125,7 @@ Trainer(policy, env, algorithm, ...)
 **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]`.
diff --git a/embodichain/learning/rl/algo/__init__.py b/embodichain/learning/rl/algo/__init__.py
index 5fb312ea7..7551d2b0c 100644
--- a/embodichain/learning/rl/algo/__init__.py
+++ b/embodichain/learning/rl/algo/__init__.py
@@ -23,6 +23,11 @@
 import torch
 from torch.nn.parallel import DistributedDataParallel as DDP
 
+from embodichain.learning.rl.utils import (
+    coerce_lr_scheduler_cfg,
+    coerce_optimizer_cfg,
+)
+
 from .apg import APG, APGCfg, segmented_discounted_return
 from .base import BaseAlgorithm, RolloutKind
 from .common import compute_gae
@@ -40,6 +45,16 @@ def get_registered_algo_names() -> list[str]:
     return list(_ALGO_REGISTRY.keys())
 
 
+def _normalize_algo_cfg_kwargs(cfg_kwargs: Dict[str, Any]) -> Dict[str, Any]:
+    """Coerce nested optimizer/scheduler mappings from YAML/JSON."""
+    normalized = dict(cfg_kwargs)
+    if "optimizer" in normalized:
+        normalized["optimizer"] = coerce_optimizer_cfg(normalized["optimizer"])
+    if "lr_scheduler" in normalized:
+        normalized["lr_scheduler"] = coerce_lr_scheduler_cfg(normalized["lr_scheduler"])
+    return normalized
+
+
 def build_algo(
     name: str,
     cfg_kwargs: Dict[str, Any],
@@ -54,7 +69,7 @@ def build_algo(
             f"Algorithm '{name}' not found. Available: {get_registered_algo_names()}"
         )
     CfgCls, AlgoCls = _ALGO_REGISTRY[key]
-    cfg = CfgCls(device=str(device), **cfg_kwargs)
+    cfg = CfgCls(device=str(device), **_normalize_algo_cfg_kwargs(cfg_kwargs))
     if distributed:
         if AlgoCls.rollout_kind is RolloutKind.DIFFERENTIABLE:
             raise ValueError(
diff --git a/embodichain/learning/rl/algo/apg.py b/embodichain/learning/rl/algo/apg.py
index 76f1f33d3..279320593 100644
--- a/embodichain/learning/rl/algo/apg.py
+++ b/embodichain/learning/rl/algo/apg.py
@@ -53,10 +53,9 @@ def segmented_discounted_return(
 
 @configclass
 class APGCfg(AlgorithmCfg):
-    """Configuration for analytic policy-gradient updates.
+    """Analytic policy-gradient config.
 
-    ``gamma`` is applied within each truncated-backpropagation segment. The
-    discount restarts after a terminated or truncated transition.
+    ``gamma`` applies within each TBPTT segment and restarts after done.
     """
 
     ent_coef: float = 0.0
@@ -72,7 +71,7 @@ def __init__(self, cfg: APGCfg, policy: torch.nn.Module) -> None:
         self.cfg = cfg
         self.policy = policy
         self.device = torch.device(cfg.device)
-        self.optimizer = torch.optim.Adam(policy.parameters(), lr=cfg.learning_rate)
+        self._setup_optimization(cfg, policy.parameters())
         self._update_active = False
         self._update_valid = True
         self._discount: torch.Tensor | None = None
@@ -82,14 +81,7 @@ def __init__(self, cfg: APGCfg, policy: torch.nn.Module) -> None:
         self._num_accumulated_steps = 0
 
     def update(self, rollout: DifferentiableRollout) -> Dict[str, float]:
-        """Apply one pathwise-gradient update from a rollout segment.
-
-        Args:
-            rollout: Graph-preserving rollout used for the update.
-
-        Returns:
-            Scalar metrics describing the optimizer update.
-        """
+        """Apply one pathwise-gradient update from a rollout segment."""
         self.begin_update()
         try:
             self.accumulate_segment(rollout)
@@ -99,7 +91,6 @@ def update(self, rollout: DifferentiableRollout) -> Dict[str, float]:
             raise
 
     def begin_update(self) -> None:
-        """Start accumulating segment gradients for one optimizer update."""
         if self._update_active:
             raise RuntimeError("An APG optimizer update is already active.")
         self.optimizer.zero_grad(set_to_none=True)
@@ -112,11 +103,7 @@ def begin_update(self) -> None:
         self._num_accumulated_steps = 0
 
     def accumulate_segment(self, rollout: DifferentiableRollout) -> None:
-        """Accumulate gradients from one TBPTT segment without updating policy.
-
-        Args:
-            rollout: Graph-preserving segment from the active update horizon.
-        """
+        """Accumulate gradients from one TBPTT segment without stepping the optimizer."""
         if not self._update_active:
             raise RuntimeError("Call begin_update() before accumulating a segment.")
         if rollout.num_steps == 0:
@@ -157,11 +144,7 @@ def accumulate_segment(self, rollout: DifferentiableRollout) -> None:
             return
 
     def finish_update(self) -> Dict[str, float]:
-        """Clip accumulated gradients and perform one optimizer step.
-
-        Returns:
-            Metrics aggregated across all accumulated segments.
-        """
+        """Clip gradients and apply one optimizer step."""
         if not self._update_active:
             raise RuntimeError("Call begin_update() before finishing an update.")
         if self._num_accumulated_steps == 0:
@@ -180,6 +163,7 @@ def finish_update(self) -> Dict[str, float]:
             self.cfg.max_grad_norm,
         )
         self.optimizer.step()
+        self._step_scheduler()
         metrics = self._accumulated_metrics(
             grad_norm=float(grad_norm.detach()),
             skipped_update=0.0,
@@ -188,7 +172,6 @@ def finish_update(self) -> Dict[str, float]:
         return metrics
 
     def cancel_update(self) -> None:
-        """Discard gradients accumulated for the active optimizer update."""
         self.optimizer.zero_grad(set_to_none=True)
         self._update_active = False
 
@@ -229,4 +212,5 @@ def _accumulated_metrics(
             "entropy": self._entropy_total,
             "grad_norm": grad_norm,
             "skipped_update": skipped_update,
+            "learning_rate": self.current_learning_rate(),
         }
diff --git a/embodichain/learning/rl/algo/base.py b/embodichain/learning/rl/algo/base.py
index 57235fbdc..f9b59b3f6 100644
--- a/embodichain/learning/rl/algo/base.py
+++ b/embodichain/learning/rl/algo/base.py
@@ -17,11 +17,22 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
+from collections.abc import Iterable
 from enum import Enum
 from typing import Dict, Generic, TypeVar
 
 import torch
 
+from embodichain.learning.rl.utils import (
+    AlgorithmCfg,
+    bind_scheduler_horizon,
+    build_lr_scheduler,
+    build_optimizer,
+    coerce_lr_scheduler_cfg,
+    coerce_optimizer_cfg,
+    scheduler_needs_horizon,
+)
+
 __all__ = ["BaseAlgorithm", "RolloutKind"]
 
 RolloutT = TypeVar("RolloutT")
@@ -35,15 +46,49 @@ class RolloutKind(str, Enum):
 
 
 class BaseAlgorithm(ABC, Generic[RolloutT]):
-    """Base class for RL algorithms.
-
-    Algorithms only implement policy updates over collected rollouts.
-    """
+    """Base class for RL algorithms."""
 
     device: torch.device
     rollout_kind = RolloutKind.STANDARD
+    optimizer: torch.optim.Optimizer
+    lr_scheduler: torch.optim.lr_scheduler.LRScheduler | None
 
     @abstractmethod
     def update(self, rollout: RolloutT) -> Dict[str, float]:
         """Update policy using collected data and return training losses."""
         raise NotImplementedError
+
+    def _setup_optimization(
+        self,
+        cfg: AlgorithmCfg,
+        parameters: Iterable[torch.nn.Parameter],
+    ) -> None:
+        cfg.optimizer = coerce_optimizer_cfg(cfg.optimizer)
+        cfg.lr_scheduler = coerce_lr_scheduler_cfg(cfg.lr_scheduler)
+        self._lr_scheduler_cfg = cfg.lr_scheduler
+        self.optimizer = build_optimizer(parameters, cfg.optimizer)
+        self.lr_scheduler = None
+        if self._lr_scheduler_cfg.name is not None and not scheduler_needs_horizon(
+            self._lr_scheduler_cfg
+        ):
+            self.lr_scheduler = build_lr_scheduler(
+                self.optimizer, self._lr_scheduler_cfg
+            )
+
+    def bind_schedule(self, *, total_updates: int) -> None:
+        """Bind horizon-dependent LR schedules from the training budget."""
+        if total_updates <= 0:
+            raise ValueError("total_updates must be positive.")
+        if not scheduler_needs_horizon(self._lr_scheduler_cfg):
+            return
+        self._lr_scheduler_cfg = bind_scheduler_horizon(
+            self._lr_scheduler_cfg, total_updates
+        )
+        self.lr_scheduler = build_lr_scheduler(self.optimizer, self._lr_scheduler_cfg)
+
+    def current_learning_rate(self) -> float:
+        return float(self.optimizer.param_groups[0]["lr"])
+
+    def _step_scheduler(self) -> None:
+        if self.lr_scheduler is not None:
+            self.lr_scheduler.step()
diff --git a/embodichain/learning/rl/algo/grpo.py b/embodichain/learning/rl/algo/grpo.py
index f9849ffc8..3dcacf20e 100644
--- a/embodichain/learning/rl/algo/grpo.py
+++ b/embodichain/learning/rl/algo/grpo.py
@@ -55,7 +55,7 @@ def __init__(self, cfg: GRPOCfg, policy):
         self.cfg = cfg
         self.policy = policy
         self.device = torch.device(cfg.device)
-        self.optimizer = torch.optim.Adam(policy.parameters(), lr=cfg.learning_rate)
+        self._setup_optimization(cfg, policy.parameters())
         if self.cfg.kl_coef > 0.0:
             raw_policy = getattr(policy, "module", policy)
             self.ref_policy = deepcopy(raw_policy).to(self.device).eval()
@@ -198,8 +198,10 @@ def update(self, rollout: TensorDict) -> Dict[str, float]:
                 total_kl += kl.item() * weight
                 total_weight += weight
 
+        self._step_scheduler()
         return {
             "actor_loss": total_actor_loss / max(1.0, total_weight),
             "entropy": total_entropy / max(1.0, total_weight),
             "approx_ref_kl": total_kl / max(1.0, total_weight),
+            "learning_rate": self.current_learning_rate(),
         }
diff --git a/embodichain/learning/rl/algo/ppo.py b/embodichain/learning/rl/algo/ppo.py
index 0d8d5efc7..942e9da4c 100644
--- a/embodichain/learning/rl/algo/ppo.py
+++ b/embodichain/learning/rl/algo/ppo.py
@@ -47,7 +47,7 @@ def __init__(self, cfg: PPOCfg, policy):
         self.cfg = cfg
         self.policy = policy
         self.device = torch.device(cfg.device)
-        self.optimizer = torch.optim.Adam(policy.parameters(), lr=cfg.learning_rate)
+        self._setup_optimization(cfg, policy.parameters())
 
     def update(self, rollout: TensorDict) -> Dict[str, float]:
         """Update the policy using a collected rollout."""
@@ -108,8 +108,10 @@ def update(self, rollout: TensorDict) -> Dict[str, float]:
                 total_entropy += (-entropy_loss.item()) * bs
                 total_steps += bs
 
+        self._step_scheduler()
         return {
             "actor_loss": total_actor_loss / max(1, total_steps),
             "value_loss": total_value_loss / max(1, total_steps),
             "entropy": total_entropy / max(1, total_steps),
+            "learning_rate": self.current_learning_rate(),
         }
diff --git a/embodichain/learning/rl/differentiable_trainer.py b/embodichain/learning/rl/differentiable_trainer.py
index 46695c930..fe2472f85 100644
--- a/embodichain/learning/rl/differentiable_trainer.py
+++ b/embodichain/learning/rl/differentiable_trainer.py
@@ -32,6 +32,7 @@
 from embodichain.learning.rl.env import DifferentiableVecEnv
 from embodichain.learning.rl.evaluation import evaluate_episodes
 from embodichain.learning.rl.models import Policy
+from embodichain.learning.rl.utils import LRSchedulerCfg, build_lr_scheduler
 from embodichain.utils import configclass
 
 if TYPE_CHECKING:
@@ -130,6 +131,11 @@ def train(self, total_timesteps: int) -> dict[str, Any]:
         if total_timesteps < 0:
             raise ValueError("total_timesteps 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)
+
         self.policy.train()
         while self.global_step < total_timesteps:
             remaining_vector_steps = math.ceil(
@@ -205,17 +211,21 @@ def save_checkpoint(self, path: str | Path | None = None) -> str:
             )
         checkpoint_path = Path(path)
         checkpoint_path.parent.mkdir(parents=True, exist_ok=True)
-        torch.save(
-            {
-                "schema_version": _CHECKPOINT_SCHEMA_VERSION,
-                "global_step": self.global_step,
-                "num_updates": self.num_updates,
-                "policy": self.policy.state_dict(),
-                "optimizer": self.algorithm.optimizer.state_dict(),
-                "best_eval_value": self.best_eval_value,
-            },
-            checkpoint_path,
-        )
+        payload = {
+            "schema_version": _CHECKPOINT_SCHEMA_VERSION,
+            "global_step": self.global_step,
+            "num_updates": self.num_updates,
+            "policy": self.policy.state_dict(),
+            "optimizer": self.algorithm.optimizer.state_dict(),
+            "best_eval_value": self.best_eval_value,
+        }
+        if self.algorithm.lr_scheduler is not None:
+            payload["lr_scheduler"] = self.algorithm.lr_scheduler.state_dict()
+            payload["lr_scheduler_cfg"] = {
+                "name": self.algorithm._lr_scheduler_cfg.name,
+                "kwargs": dict(self.algorithm._lr_scheduler_cfg.kwargs),
+            }
+        torch.save(payload, checkpoint_path)
         self.latest_checkpoint_path = str(checkpoint_path)
         return self.latest_checkpoint_path
 
@@ -237,6 +247,15 @@ def load_checkpoint(self, path: str | Path) -> None:
             )
         self.policy.load_state_dict(checkpoint["policy"])
         self.algorithm.optimizer.load_state_dict(checkpoint["optimizer"])
+        sched_cfg_data = checkpoint.get("lr_scheduler_cfg")
+        if sched_cfg_data is not None and checkpoint.get("lr_scheduler") is not None:
+            bound_cfg = LRSchedulerCfg(**sched_cfg_data)
+            self.algorithm._lr_scheduler_cfg = bound_cfg
+            self.algorithm.lr_scheduler = build_lr_scheduler(
+                self.algorithm.optimizer,
+                bound_cfg,
+            )
+            self.algorithm.lr_scheduler.load_state_dict(checkpoint["lr_scheduler"])
         self.global_step = int(checkpoint["global_step"])
         self.num_updates = int(checkpoint["num_updates"])
         self.best_eval_value = checkpoint.get("best_eval_value")
diff --git a/embodichain/learning/rl/experimental/newton/train_planar_reach.py b/embodichain/learning/rl/experimental/newton/train_planar_reach.py
index 3e980f895..aaeefd8b0 100644
--- a/embodichain/learning/rl/experimental/newton/train_planar_reach.py
+++ b/embodichain/learning/rl/experimental/newton/train_planar_reach.py
@@ -34,6 +34,7 @@
     DifferentiableTrainerCfg,
 )
 from embodichain.learning.rl.models import ActorOnly, MLP
+from embodichain.learning.rl.utils import OptimizerCfg
 from embodichain.utils import configclass
 
 from .planar_reach import NewtonPlanarReachEnv, NewtonPlanarReachEnvCfg
@@ -115,7 +116,7 @@ def train_planar_reach(
     algorithm = APG(
         APGCfg(
             device=cfg.device,
-            learning_rate=cfg.learning_rate,
+            optimizer=OptimizerCfg(learning_rate=cfg.learning_rate),
             gamma=0.99,
             max_grad_norm=10.0,
         ),
diff --git a/embodichain/learning/rl/utils/__init__.py b/embodichain/learning/rl/utils/__init__.py
index e842d1460..bea889408 100644
--- a/embodichain/learning/rl/utils/__init__.py
+++ b/embodichain/learning/rl/utils/__init__.py
@@ -14,13 +14,33 @@
 # limitations under the License.
 # ----------------------------------------------------------------------------
 
-"""RL helper utilities: the ``AlgorithmCfg`` config and data-conversion helpers (``dict_to_tensordict``, ``flatten_dict_observation``)."""
+"""RL helper utilities: algorithm config, optimizers, and observation helpers."""
 
-from .config import AlgorithmCfg
+from .config import AlgorithmCfg, LRSchedulerCfg, OptimizerCfg
 from .helper import dict_to_tensordict, flatten_dict_observation
+from .optimizer import (
+    bind_scheduler_horizon,
+    build_lr_scheduler,
+    build_optimizer,
+    coerce_lr_scheduler_cfg,
+    coerce_optimizer_cfg,
+    get_registered_lr_scheduler_names,
+    get_registered_optimizer_names,
+    scheduler_needs_horizon,
+)
 
 __all__ = [
     "AlgorithmCfg",
+    "LRSchedulerCfg",
+    "OptimizerCfg",
+    "bind_scheduler_horizon",
+    "build_lr_scheduler",
+    "build_optimizer",
+    "coerce_lr_scheduler_cfg",
+    "coerce_optimizer_cfg",
     "dict_to_tensordict",
     "flatten_dict_observation",
+    "get_registered_lr_scheduler_names",
+    "get_registered_optimizer_names",
+    "scheduler_needs_horizon",
 ]
diff --git a/embodichain/learning/rl/utils/config.py b/embodichain/learning/rl/utils/config.py
index 963e2e0a3..d1d491c32 100644
--- a/embodichain/learning/rl/utils/config.py
+++ b/embodichain/learning/rl/utils/config.py
@@ -14,15 +14,43 @@
 # limitations under the License.
 # ----------------------------------------------------------------------------
 
+from __future__ import annotations
+
+from typing import Any
+
 from embodichain.utils import configclass
 
+__all__ = ["AlgorithmCfg", "LRSchedulerCfg", "OptimizerCfg"]
+
+
+@configclass
+class OptimizerCfg:
+    """Policy optimizer configuration."""
+
+    name: str = "adam"
+    learning_rate: float = 3e-4
+    kwargs: dict[str, Any] = dict()
+
+
+@configclass
+class LRSchedulerCfg:
+    """Optional LR scheduler. ``name=None`` disables scheduling.
+
+    Horizon keys (``total_iters`` / ``T_max``) may be omitted and bound later by
+    ``BaseAlgorithm.bind_schedule``.
+    """
+
+    name: str | None = None
+    kwargs: dict[str, Any] = dict()
+
 
 @configclass
 class AlgorithmCfg:
-    """Minimal algorithm configuration shared across RL algorithms."""
+    """Shared fields for RL algorithm configs."""
 
     device: str = "cuda"
-    learning_rate: float = 3e-4
+    optimizer: OptimizerCfg = OptimizerCfg()
+    lr_scheduler: LRSchedulerCfg = LRSchedulerCfg()
     batch_size: int = 64
     gamma: float = 0.99
     gae_lambda: float = 0.95
diff --git a/embodichain/learning/rl/utils/optimizer.py b/embodichain/learning/rl/utils/optimizer.py
new file mode 100644
index 000000000..6e6df2635
--- /dev/null
+++ b/embodichain/learning/rl/utils/optimizer.py
@@ -0,0 +1,195 @@
+# ----------------------------------------------------------------------------
+# 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.
+# ----------------------------------------------------------------------------
+
+"""Shared optimizer and LR-scheduler construction for RL algorithms."""
+
+from __future__ import annotations
+
+from collections.abc import Callable, Iterable, Mapping
+from typing import Any
+
+import torch
+
+from embodichain.learning.rl.utils.config import LRSchedulerCfg, OptimizerCfg
+
+__all__ = [
+    "bind_scheduler_horizon",
+    "build_lr_scheduler",
+    "build_optimizer",
+    "coerce_lr_scheduler_cfg",
+    "coerce_optimizer_cfg",
+    "get_registered_lr_scheduler_names",
+    "get_registered_optimizer_names",
+    "scheduler_needs_horizon",
+]
+
+_OPTIMIZER_REGISTRY: dict[str, type[torch.optim.Optimizer]] = {
+    "adam": torch.optim.Adam,
+    "adamw": torch.optim.AdamW,
+    "sgd": torch.optim.SGD,
+    "rmsprop": torch.optim.RMSprop,
+}
+
+_HORIZON_KEYS: dict[str, str] = {
+    "linear": "total_iters",
+    "cosine": "T_max",
+}
+
+
+def _build_linear(
+    optimizer: torch.optim.Optimizer, kwargs: dict[str, Any]
+) -> torch.optim.lr_scheduler.LRScheduler:
+    if "total_iters" not in kwargs:
+        raise ValueError(
+            "linear scheduler requires kwargs['total_iters'] "
+            "(or bind it via BaseAlgorithm.bind_schedule)."
+        )
+    kwargs.setdefault("start_factor", 1.0)
+    kwargs.setdefault("end_factor", 0.0)
+    return torch.optim.lr_scheduler.LinearLR(optimizer, **kwargs)
+
+
+def _build_cosine(
+    optimizer: torch.optim.Optimizer, kwargs: dict[str, Any]
+) -> torch.optim.lr_scheduler.LRScheduler:
+    if "T_max" not in kwargs:
+        raise ValueError(
+            "cosine scheduler requires kwargs['T_max'] "
+            "(or bind it via BaseAlgorithm.bind_schedule)."
+        )
+    return torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, **kwargs)
+
+
+def _build_step(
+    optimizer: torch.optim.Optimizer, kwargs: dict[str, Any]
+) -> torch.optim.lr_scheduler.LRScheduler:
+    if "step_size" not in kwargs:
+        raise ValueError("step scheduler requires kwargs['step_size'].")
+    return torch.optim.lr_scheduler.StepLR(optimizer, **kwargs)
+
+
+def _build_exponential(
+    optimizer: torch.optim.Optimizer, kwargs: dict[str, Any]
+) -> torch.optim.lr_scheduler.LRScheduler:
+    if "gamma" not in kwargs:
+        raise ValueError("exponential scheduler requires kwargs['gamma'].")
+    return torch.optim.lr_scheduler.ExponentialLR(optimizer, **kwargs)
+
+
+_SCHEDULER_BUILDERS: dict[
+    str,
+    Callable[
+        [torch.optim.Optimizer, dict[str, Any]], torch.optim.lr_scheduler.LRScheduler
+    ],
+] = {
+    "linear": _build_linear,
+    "cosine": _build_cosine,
+    "step": _build_step,
+    "exponential": _build_exponential,
+}
+
+
+def get_registered_optimizer_names() -> list[str]:
+    return sorted(_OPTIMIZER_REGISTRY)
+
+
+def get_registered_lr_scheduler_names() -> list[str]:
+    return sorted(_SCHEDULER_BUILDERS)
+
+
+def coerce_optimizer_cfg(
+    value: OptimizerCfg | Mapping[str, Any] | None,
+) -> OptimizerCfg:
+    if value is None:
+        return OptimizerCfg()
+    if isinstance(value, OptimizerCfg):
+        return value
+    if isinstance(value, Mapping):
+        return OptimizerCfg(**dict(value))
+    raise TypeError(f"Expected OptimizerCfg or mapping, got {type(value)!r}.")
+
+
+def coerce_lr_scheduler_cfg(
+    value: LRSchedulerCfg | Mapping[str, Any] | None,
+) -> LRSchedulerCfg:
+    if value is None:
+        return LRSchedulerCfg()
+    if isinstance(value, LRSchedulerCfg):
+        return value
+    if isinstance(value, Mapping):
+        return LRSchedulerCfg(**dict(value))
+    raise TypeError(f"Expected LRSchedulerCfg or mapping, got {type(value)!r}.")
+
+
+def scheduler_needs_horizon(cfg: LRSchedulerCfg | Mapping[str, Any] | None) -> bool:
+    cfg = coerce_lr_scheduler_cfg(cfg)
+    if cfg.name is None:
+        return False
+    horizon_key = _HORIZON_KEYS.get(cfg.name.lower())
+    return horizon_key is not None and horizon_key not in dict(cfg.kwargs)
+
+
+def bind_scheduler_horizon(
+    cfg: LRSchedulerCfg | Mapping[str, Any] | None,
+    total_updates: int,
+) -> LRSchedulerCfg:
+    """Fill ``total_iters`` / ``T_max`` from the training update budget."""
+    cfg = coerce_lr_scheduler_cfg(cfg)
+    if cfg.name is None:
+        return cfg
+    horizon_key = _HORIZON_KEYS.get(cfg.name.lower())
+    if horizon_key is None:
+        return cfg
+    kwargs = dict(cfg.kwargs)
+    kwargs.setdefault(horizon_key, total_updates)
+    return LRSchedulerCfg(name=cfg.name, kwargs=kwargs)
+
+
+def build_optimizer(
+    parameters: Iterable[torch.nn.Parameter],
+    cfg: OptimizerCfg | Mapping[str, Any] | None = None,
+) -> torch.optim.Optimizer:
+    opt_cfg = coerce_optimizer_cfg(cfg)
+    key = opt_cfg.name.lower()
+    if key not in _OPTIMIZER_REGISTRY:
+        available = ", ".join(get_registered_optimizer_names())
+        raise ValueError(
+            f"Unsupported optimizer '{opt_cfg.name}'. Available: {available}"
+        )
+    opt_kwargs = dict(opt_cfg.kwargs)
+    if "lr" in opt_kwargs:
+        raise ValueError(
+            "Pass learning rate via OptimizerCfg.learning_rate, not kwargs['lr']."
+        )
+    return _OPTIMIZER_REGISTRY[key](parameters, lr=opt_cfg.learning_rate, **opt_kwargs)
+
+
+def build_lr_scheduler(
+    optimizer: torch.optim.Optimizer,
+    cfg: LRSchedulerCfg | Mapping[str, Any] | None,
+) -> torch.optim.lr_scheduler.LRScheduler | None:
+    """Build a scheduler, or ``None`` when ``name`` is unset."""
+    cfg = coerce_lr_scheduler_cfg(cfg)
+    if cfg.name is None:
+        return None
+    key = cfg.name.lower()
+    builder = _SCHEDULER_BUILDERS.get(key)
+    if builder is None:
+        available = ", ".join(get_registered_lr_scheduler_names())
+        raise ValueError(
+            f"Unsupported lr_scheduler '{cfg.name}'. Available: {available}"
+        )
+    return builder(optimizer, dict(cfg.kwargs))
diff --git a/embodichain/learning/rl/utils/trainer.py b/embodichain/learning/rl/utils/trainer.py
index 678696346..a97a9ac89 100644
--- a/embodichain/learning/rl/utils/trainer.py
+++ b/embodichain/learning/rl/utils/trainer.py
@@ -16,6 +16,7 @@
 
 from __future__ import annotations
 
+import math
 import time
 from collections import deque
 from typing import Any
@@ -162,6 +163,11 @@ def _pack_log_dict(self, prefix: str, data: dict) -> dict:
     def train(self, total_timesteps: int) -> dict[str, Any]:
         if self.rank == 0:
             print(f"Start training, total steps: {total_timesteps}")
+        num_envs = int(self.env.num_envs)
+        steps_per_update = self.buffer_size * 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)
         while self.global_step < total_timesteps:
             self._collect_rollout()
             losses = self.algorithm.update(self.buffer.get(flatten=False))
@@ -187,12 +193,9 @@ def train(self, total_timesteps: int) -> dict[str, Any]:
     def _collect_rollout(self):
         """Collect a rollout with the synchronous collector."""
 
-        # Callback function for statistics and logging
         def on_step(tensordict: TensorDict, info: dict):
-            """Callback called at each step during rollout collection."""
             reward = tensordict["reward"]
             done = tensordict["done"]
-            # Episode stats
             self.curr_ret += reward
             self.curr_len += 1
             done_idx = torch.nonzero(done, as_tuple=False).squeeze(-1)
@@ -226,7 +229,6 @@ def on_step(tensordict: TensorDict, info: dict):
         )
         self.buffer.add(rollout)
 
-        # Sync global_step and episode stats across ranks in distributed mode
         if self.distributed:
             if not torch.distributed.is_available():
                 raise RuntimeError(
@@ -419,6 +421,15 @@ def save_checkpoint(self, path: str | None = None) -> str | None:
         optimizer = getattr(self.algorithm, "optimizer", None)
         if optimizer is not None:
             checkpoint["optimizer"] = optimizer.state_dict()
+        lr_scheduler = getattr(self.algorithm, "lr_scheduler", None)
+        if lr_scheduler is not None:
+            checkpoint["lr_scheduler"] = lr_scheduler.state_dict()
+            sched_cfg = getattr(self.algorithm, "_lr_scheduler_cfg", None)
+            if sched_cfg is not None:
+                checkpoint["lr_scheduler_cfg"] = {
+                    "name": sched_cfg.name,
+                    "kwargs": dict(sched_cfg.kwargs),
+                }
         torch.save(checkpoint, path)
         self.latest_checkpoint_path = path
         print(f"Checkpoint saved: {path}")
diff --git a/embodichain_tasks/configs/agents/rl/basic/cart_pole/train_config.json b/embodichain_tasks/configs/agents/rl/basic/cart_pole/train_config.json
index a3fcabdc9..7a2deedb2 100644
--- a/embodichain_tasks/configs/agents/rl/basic/cart_pole/train_config.json
+++ b/embodichain_tasks/configs/agents/rl/basic/cart_pole/train_config.json
@@ -79,7 +79,10 @@
     "algorithm": {
         "name": "ppo",
         "cfg": {
-            "learning_rate": 0.0001,
+            "optimizer": {
+                "name": "adam",
+                "learning_rate": 0.0001
+            },
             "n_epochs": 10,
             "batch_size": 8192,
             "gamma": 0.99,
diff --git a/embodichain_tasks/configs/agents/rl/basic/cart_pole/train_config.yaml b/embodichain_tasks/configs/agents/rl/basic/cart_pole/train_config.yaml
index b729186e1..c160305e1 100644
--- a/embodichain_tasks/configs/agents/rl/basic/cart_pole/train_config.yaml
+++ b/embodichain_tasks/configs/agents/rl/basic/cart_pole/train_config.yaml
@@ -60,7 +60,9 @@ policy:
 algorithm:
   name: ppo
   cfg:
-    learning_rate: 0.0001
+    optimizer:
+      name: adam
+      learning_rate: 0.0001
     n_epochs: 10
     batch_size: 8192
     gamma: 0.99
diff --git a/embodichain_tasks/configs/agents/rl/basic/cart_pole/train_config_grpo.json b/embodichain_tasks/configs/agents/rl/basic/cart_pole/train_config_grpo.json
index 57ac46543..80bedf8a5 100644
--- a/embodichain_tasks/configs/agents/rl/basic/cart_pole/train_config_grpo.json
+++ b/embodichain_tasks/configs/agents/rl/basic/cart_pole/train_config_grpo.json
@@ -70,7 +70,10 @@
     "algorithm": {
         "name": "grpo",
         "cfg": {
-            "learning_rate": 0.0001,
+            "optimizer": {
+                "name": "adam",
+                "learning_rate": 0.0001
+            },
             "n_epochs": 10,
             "batch_size": 8192,
             "gamma": 0.99,
diff --git a/embodichain_tasks/configs/agents/rl/basic/cart_pole/train_config_grpo.yaml b/embodichain_tasks/configs/agents/rl/basic/cart_pole/train_config_grpo.yaml
index 7ec0ecf50..2895a7b52 100644
--- a/embodichain_tasks/configs/agents/rl/basic/cart_pole/train_config_grpo.yaml
+++ b/embodichain_tasks/configs/agents/rl/basic/cart_pole/train_config_grpo.yaml
@@ -54,7 +54,9 @@ policy:
 algorithm:
   name: grpo
   cfg:
-    learning_rate: 0.0001
+    optimizer:
+      name: adam
+      learning_rate: 0.0001
     n_epochs: 10
     batch_size: 8192
     gamma: 0.99
diff --git a/embodichain_tasks/configs/agents/rl/basic/point_mass/train_apg.yaml b/embodichain_tasks/configs/agents/rl/basic/point_mass/train_apg.yaml
index 38f4b6dc6..cf2180fe2 100644
--- a/embodichain_tasks/configs/agents/rl/basic/point_mass/train_apg.yaml
+++ b/embodichain_tasks/configs/agents/rl/basic/point_mass/train_apg.yaml
@@ -36,7 +36,9 @@ policy:
 algorithm:
   name: apg
   cfg:
-    learning_rate: 0.00025
+    optimizer:
+      name: adam
+      learning_rate: 0.00025
     gamma: 0.99
     max_grad_norm: 0.5
     ent_coef: 0.0
diff --git a/embodichain_tasks/configs/agents/rl/basic/point_mass/train_ppo.yaml b/embodichain_tasks/configs/agents/rl/basic/point_mass/train_ppo.yaml
index e18b8f8db..7ffa60103 100644
--- a/embodichain_tasks/configs/agents/rl/basic/point_mass/train_ppo.yaml
+++ b/embodichain_tasks/configs/agents/rl/basic/point_mass/train_ppo.yaml
@@ -41,7 +41,9 @@ policy:
 algorithm:
   name: ppo
   cfg:
-    learning_rate: 0.00025
+    optimizer:
+      name: adam
+      learning_rate: 0.00025
     n_epochs: 10
     batch_size: 4096
     gamma: 0.99
diff --git a/embodichain_tasks/configs/agents/rl/push_cube/train_config.json b/embodichain_tasks/configs/agents/rl/push_cube/train_config.json
index d8ea93d04..263498078 100644
--- a/embodichain_tasks/configs/agents/rl/push_cube/train_config.json
+++ b/embodichain_tasks/configs/agents/rl/push_cube/train_config.json
@@ -62,7 +62,10 @@
     "algorithm": {
         "name": "ppo",
         "cfg": {
-            "learning_rate": 0.0001,
+            "optimizer": {
+                "name": "adam",
+                "learning_rate": 0.0001
+            },
             "n_epochs": 10,
             "batch_size": 8192,
             "gamma": 0.99,
diff --git a/embodichain_tasks/configs/agents/rl/push_cube/train_config_grpo.json b/embodichain_tasks/configs/agents/rl/push_cube/train_config_grpo.json
index e87ddb628..4dea67c20 100644
--- a/embodichain_tasks/configs/agents/rl/push_cube/train_config_grpo.json
+++ b/embodichain_tasks/configs/agents/rl/push_cube/train_config_grpo.json
@@ -48,7 +48,10 @@
     "algorithm": {
         "name": "grpo",
         "cfg": {
-            "learning_rate": 0.0001,
+            "optimizer": {
+                "name": "adam",
+                "learning_rate": 0.0001
+            },
             "n_epochs": 10,
             "batch_size": 8192,
             "gamma": 0.99,
diff --git a/scripts/benchmark/rl/algorithms/apg.yaml b/scripts/benchmark/rl/algorithms/apg.yaml
index 8f0a579cf..7ddd1ec0e 100644
--- a/scripts/benchmark/rl/algorithms/apg.yaml
+++ b/scripts/benchmark/rl/algorithms/apg.yaml
@@ -13,7 +13,9 @@ config:
   algorithm:
     name: apg
     cfg:
-      learning_rate: 0.00025
+      optimizer:
+        name: adam
+        learning_rate: 0.00025
       gamma: 0.99
       max_grad_norm: 0.5
       ent_coef: 0.0
diff --git a/scripts/benchmark/rl/algorithms/grpo.yaml b/scripts/benchmark/rl/algorithms/grpo.yaml
index e33c673b0..64956992e 100644
--- a/scripts/benchmark/rl/algorithms/grpo.yaml
+++ b/scripts/benchmark/rl/algorithms/grpo.yaml
@@ -10,7 +10,9 @@ config:
   algorithm:
     name: grpo
     cfg:
-      learning_rate: 0.0001
+      optimizer:
+        name: adam
+        learning_rate: 0.0001
       n_epochs: 10
       batch_size: 8192
       gamma: 0.99
diff --git a/scripts/benchmark/rl/algorithms/ppo.yaml b/scripts/benchmark/rl/algorithms/ppo.yaml
index 361c93866..addf877f5 100644
--- a/scripts/benchmark/rl/algorithms/ppo.yaml
+++ b/scripts/benchmark/rl/algorithms/ppo.yaml
@@ -15,7 +15,9 @@ config:
   algorithm:
     name: ppo
     cfg:
-      learning_rate: 0.0001
+      optimizer:
+        name: adam
+        learning_rate: 0.0001
       n_epochs: 10
       batch_size: 8192
       gamma: 0.99
diff --git a/tests/learning/test_apg.py b/tests/learning/test_apg.py
index 3cd509853..32a7348fc 100644
--- a/tests/learning/test_apg.py
+++ b/tests/learning/test_apg.py
@@ -39,6 +39,7 @@
     DifferentiableTransition,
 )
 from embodichain.learning.rl.models import ActorOnly
+from embodichain.learning.rl.utils import LRSchedulerCfg, OptimizerCfg
 
 
 class _LinearDifferentiableEnv:
@@ -203,7 +204,12 @@ def test_apg_update_improves_deterministic_objective() -> None:
     env = _LinearDifferentiableEnv()
     policy = _make_policy()
     algorithm = APG(
-        APGCfg(device="cpu", learning_rate=0.05, gamma=0.9, max_grad_norm=100.0),
+        APGCfg(
+            device="cpu",
+            optimizer=OptimizerCfg(learning_rate=0.05),
+            gamma=0.9,
+            max_grad_norm=100.0,
+        ),
         policy,
     )
     collector = DifferentiableCollector(env, policy, env.device)
@@ -232,3 +238,65 @@ def test_apg_skips_nonfinite_update_without_changing_policy() -> None:
 
     assert metrics["skipped_update"] == 1.0
     assert torch.equal(policy.actor.weight, initial_weight)
+
+
+def test_apg_optimizer_is_configurable() -> None:
+    policy = _make_policy()
+    algorithm = APG(
+        APGCfg(
+            device="cpu",
+            optimizer=OptimizerCfg(
+                name="adamw",
+                learning_rate=3e-4,
+                kwargs={"weight_decay": 1e-4},
+            ),
+        ),
+        policy,
+    )
+
+    assert isinstance(algorithm.optimizer, torch.optim.AdamW)
+    assert algorithm.optimizer.defaults["weight_decay"] == pytest.approx(1e-4)
+
+
+def test_apg_rejects_unknown_optimizer() -> None:
+    with pytest.raises(ValueError, match="Unsupported optimizer"):
+        APG(
+            APGCfg(
+                device="cpu",
+                optimizer=OptimizerCfg(name="not-an-optimizer"),
+            ),
+            _make_policy(),
+        )
+
+
+def test_apg_linear_scheduler_decays_on_valid_updates_only() -> None:
+    env = _LinearDifferentiableEnv()
+    policy = _make_policy()
+    algorithm = APG(
+        APGCfg(
+            device="cpu",
+            optimizer=OptimizerCfg(learning_rate=0.1),
+            lr_scheduler=LRSchedulerCfg(name="linear"),
+            gamma=0.9,
+            max_grad_norm=100.0,
+        ),
+        policy,
+    )
+    algorithm.bind_schedule(total_updates=4)
+    collector = DifferentiableCollector(env, policy, env.device)
+    initial_lr = algorithm.current_learning_rate()
+
+    algorithm.update(collector.collect(num_steps=2, deterministic=True))
+    after_valid = algorithm.current_learning_rate()
+    assert after_valid < initial_lr
+
+    skip_env = _NonFiniteRewardEnv()
+    lr_before_skip = algorithm.current_learning_rate()
+    skip_metrics = algorithm.update(
+        DifferentiableCollector(skip_env, policy, env.device).collect(
+            num_steps=1,
+            deterministic=True,
+        )
+    )
+    assert skip_metrics["skipped_update"] == 1.0
+    assert algorithm.current_learning_rate() == pytest.approx(lr_before_skip)
diff --git a/tests/learning/test_differentiable_trainer.py b/tests/learning/test_differentiable_trainer.py
index ea601d851..0fc48f9ed 100644
--- a/tests/learning/test_differentiable_trainer.py
+++ b/tests/learning/test_differentiable_trainer.py
@@ -32,6 +32,7 @@
 )
 from embodichain.learning.rl.algo import APG, APGCfg
 from embodichain.learning.rl.models import ActorOnly
+from embodichain.learning.rl.utils import LRSchedulerCfg, OptimizerCfg
 
 
 class _QuadraticActionEnv:
@@ -83,7 +84,7 @@ def _make_components(
     algorithm = APG(
         APGCfg(
             device="cpu",
-            learning_rate=0.05,
+            optimizer=OptimizerCfg(learning_rate=0.05),
             max_grad_norm=10.0,
             ent_coef=ent_coef,
         ),
@@ -228,6 +229,63 @@ def test_checkpoint_restores_policy_optimizer_and_counters(tmp_path: Path) -> No
     assert restored_algorithm.optimizer.state_dict()["state"]
 
 
+def test_checkpoint_restores_lr_scheduler_state(tmp_path: Path) -> None:
+    env = _QuadraticActionEnv()
+    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.1),
+            lr_scheduler=LRSchedulerCfg(name="linear"),
+            max_grad_norm=10.0,
+        ),
+        policy,
+    )
+    trainer = DifferentiableTrainer(
+        DifferentiableTrainerCfg(
+            segment_length=2,
+            update_horizon=2,
+            deterministic_actions=True,
+        ),
+        env,
+        policy,
+        algorithm,
+    )
+    trainer.train(total_timesteps=8)
+    lr_before = algorithm.current_learning_rate()
+    assert algorithm.lr_scheduler is not None
+    scheduler_state = algorithm.lr_scheduler.state_dict()
+    checkpoint_path = trainer.save_checkpoint(tmp_path / "apg_sched.pt")
+
+    restored_env = _QuadraticActionEnv()
+    restored_actor = nn.Linear(1, 1, bias=False)
+    restored_policy = ActorOnly(1, 1, restored_env.device, actor=restored_actor)
+    restored_algorithm = APG(
+        APGCfg(
+            device="cpu",
+            optimizer=OptimizerCfg(learning_rate=0.1),
+            lr_scheduler=LRSchedulerCfg(name="linear"),
+            max_grad_norm=10.0,
+        ),
+        restored_policy,
+    )
+    restored = DifferentiableTrainer(
+        trainer.cfg,
+        restored_env,
+        restored_policy,
+        restored_algorithm,
+    )
+    restored.load_checkpoint(checkpoint_path)
+
+    assert restored_algorithm.lr_scheduler is not None
+    assert restored_algorithm.current_learning_rate() == pytest.approx(lr_before)
+    assert restored_algorithm.lr_scheduler.state_dict()["last_epoch"] == scheduler_state[
+        "last_epoch"
+    ]
+
+
 def test_checkpoint_load_falls_back_for_older_torch(
     tmp_path: Path,
     monkeypatch: pytest.MonkeyPatch,
diff --git a/tests/learning/test_newton_planar_reach.py b/tests/learning/test_newton_planar_reach.py
index cc10c92c1..d9067778b 100644
--- a/tests/learning/test_newton_planar_reach.py
+++ b/tests/learning/test_newton_planar_reach.py
@@ -38,6 +38,7 @@
     train_planar_reach,
 )
 from embodichain.learning.rl.models import ActorOnly
+from embodichain.learning.rl.utils import OptimizerCfg
 
 
 def _make_env() -> NewtonPlanarReachEnv:
@@ -307,7 +308,7 @@ def test_newton_rollout_drives_apg_policy_update() -> None:
     algorithm = APG(
         APGCfg(
             device="cpu",
-            learning_rate=0.01,
+            optimizer=OptimizerCfg(learning_rate=0.01),
             max_grad_norm=100.0,
         ),
         policy,
diff --git a/tests/learning/test_optimizer.py b/tests/learning/test_optimizer.py
new file mode 100644
index 000000000..297a0bd88
--- /dev/null
+++ b/tests/learning/test_optimizer.py
@@ -0,0 +1,109 @@
+# ----------------------------------------------------------------------------
+# 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 shared optimizer and LR-scheduler builders."""
+
+from __future__ import annotations
+
+import pytest
+import torch
+import torch.nn as nn
+
+from embodichain.learning.rl.utils import (
+    LRSchedulerCfg,
+    OptimizerCfg,
+    build_lr_scheduler,
+    build_optimizer,
+    scheduler_needs_horizon,
+)
+
+
+def test_build_optimizer_adamw_kwargs() -> None:
+    model = nn.Linear(2, 1)
+    optimizer = build_optimizer(
+        model.parameters(),
+        OptimizerCfg(
+            name="adamw",
+            learning_rate=1e-3,
+            kwargs={"weight_decay": 0.01},
+        ),
+    )
+
+    assert isinstance(optimizer, torch.optim.AdamW)
+    assert optimizer.defaults["lr"] == pytest.approx(1e-3)
+    assert optimizer.defaults["weight_decay"] == pytest.approx(0.01)
+
+
+def test_build_optimizer_rejects_lr_in_kwargs() -> None:
+    model = nn.Linear(2, 1)
+    with pytest.raises(ValueError, match="OptimizerCfg.learning_rate"):
+        build_optimizer(
+            model.parameters(),
+            OptimizerCfg(kwargs={"lr": 1e-2}),
+        )
+
+
+def test_build_optimizer_accepts_mapping_cfg() -> None:
+    model = nn.Linear(2, 1)
+    optimizer = build_optimizer(
+        model.parameters(),
+        {"name": "sgd", "learning_rate": 0.05, "kwargs": {"momentum": 0.9}},
+    )
+    assert isinstance(optimizer, torch.optim.SGD)
+    assert optimizer.defaults["momentum"] == pytest.approx(0.9)
+
+
+def test_linear_and_cosine_need_horizon_until_bound() -> None:
+    assert scheduler_needs_horizon(LRSchedulerCfg(name="linear"))
+    assert scheduler_needs_horizon(LRSchedulerCfg(name="cosine"))
+    assert not scheduler_needs_horizon(
+        LRSchedulerCfg(name="linear", kwargs={"total_iters": 10})
+    )
+    assert not scheduler_needs_horizon(LRSchedulerCfg(name=None))
+
+
+def test_build_lr_scheduler_linear_decays() -> None:
+    model = nn.Linear(2, 1)
+    optimizer = build_optimizer(
+        model.parameters(),
+        OptimizerCfg(learning_rate=1.0),
+    )
+    scheduler = build_lr_scheduler(
+        optimizer,
+        LRSchedulerCfg(
+            name="linear",
+            kwargs={"total_iters": 4, "start_factor": 1.0, "end_factor": 0.0},
+        ),
+    )
+    assert scheduler is not None
+    rates = [optimizer.param_groups[0]["lr"]]
+    for _ in range(4):
+        optimizer.zero_grad(set_to_none=True)
+        loss = model(torch.ones(1, 2)).sum()
+        loss.backward()
+        optimizer.step()
+        scheduler.step()
+        rates.append(optimizer.param_groups[0]["lr"])
+    assert rates[0] == pytest.approx(1.0)
+    assert rates[-1] == pytest.approx(0.0)
+    assert rates[1] < rates[0]
+
+
+def test_build_lr_scheduler_cosine_requires_t_max() -> None:
+    model = nn.Linear(2, 1)
+    optimizer = build_optimizer(model.parameters(), OptimizerCfg())
+    with pytest.raises(ValueError, match="T_max"):
+        build_lr_scheduler(optimizer, LRSchedulerCfg(name="cosine"))
diff --git a/tests/learning/test_point_mass.py b/tests/learning/test_point_mass.py
index 7850b6f3f..853e76353 100644
--- a/tests/learning/test_point_mass.py
+++ b/tests/learning/test_point_mass.py
@@ -32,6 +32,7 @@
 )
 from embodichain.learning.rl.models import ActorCritic
 from embodichain.learning.rl.train import train_from_config
+from embodichain.learning.rl.utils import OptimizerCfg
 from embodichain.learning.rl.utils.trainer import Trainer
 from embodichain_tasks.rl.basic.point_mass import PointMassEnv
 
@@ -112,7 +113,7 @@ def test_unified_train_entry_runs_apg_and_ppo(
         },
     }
     algorithm_cfg = {
-        "learning_rate": 1e-3,
+        "optimizer": {"name": "adam", "learning_rate": 1e-3},
         "batch_size": 8,
         "gamma": 0.99,
         "max_grad_norm": 1.0,
@@ -188,7 +189,7 @@ def test_trainer_saves_best_checkpoint_from_eval(tmp_path: Path) -> None:
     algorithm = PPO(
         PPOCfg(
             device="cpu",
-            learning_rate=1e-3,
+            optimizer=OptimizerCfg(learning_rate=1e-3),
             n_epochs=1,
             batch_size=8,
             gamma=0.99,

From 8f61faeacd59ad8ecc547caf14a6f02e4322a1b4 Mon Sep 17 00:00:00 2001
From: yangchen73 <122090643@link.cuhk.edu.cn>
Date: Sun, 2 Aug 2026 03:54:15 +0800
Subject: [PATCH 3/7] WIP

---
 .../agents/rl/basic/point_mass/train_ppo.yaml  |  7 ++++++-
 .../embodichain_tasks/rl/basic/point_mass.py   | 18 +++++++++++-------
 tests/learning/test_point_mass.py              | 14 ++++++++++++++
 3 files changed, 31 insertions(+), 8 deletions(-)

diff --git a/embodichain_tasks/configs/agents/rl/basic/point_mass/train_ppo.yaml b/embodichain_tasks/configs/agents/rl/basic/point_mass/train_ppo.yaml
index 7ffa60103..e2d2aa604 100644
--- a/embodichain_tasks/configs/agents/rl/basic/point_mass/train_ppo.yaml
+++ b/embodichain_tasks/configs/agents/rl/basic/point_mass/train_ppo.yaml
@@ -5,10 +5,11 @@ trainer:
     cfg:
       max_episode_steps: 100
       success_threshold: 0.03
+      success_bonus: 5.0
   seed: 42
   device: cuda:0
   num_envs: 64
-  iterations: 500
+  iterations: 550
   buffer_size: 256
   enable_eval: true
   eval_freq: 32768
@@ -44,6 +45,10 @@ algorithm:
     optimizer:
       name: adam
       learning_rate: 0.00025
+    lr_scheduler:
+      name: linear
+      kwargs:
+        end_factor: 0.4
     n_epochs: 10
     batch_size: 4096
     gamma: 0.99
diff --git a/embodichain_tasks/embodichain_tasks/rl/basic/point_mass.py b/embodichain_tasks/embodichain_tasks/rl/basic/point_mass.py
index 0e018fe64..b40a1c429 100644
--- a/embodichain_tasks/embodichain_tasks/rl/basic/point_mass.py
+++ b/embodichain_tasks/embodichain_tasks/rl/basic/point_mass.py
@@ -34,10 +34,9 @@
 class PointMassEnv:
     """Navigate a damped point mass to a goal while avoiding two obstacles.
 
-    The implementation is always differentiable. Standard collectors execute
-    it under ``torch.no_grad()``, while APG retains the path from policy action
-    through dynamics and reward.
-    """
+    The imphementation is ale implementation is.aStandardfcollectofs execete
+e   it ntiable. Standard collectorwhile s exrctaine it urat `fram,lA icyrictpm action
+    roruugh gynnancr.d
 
     observation_dim = 14
     action_dim = 2
@@ -54,6 +53,7 @@ def __init__(
         linear_damping: float = 5.0,
         dt: float = 1.0 / 60.0,
         success_threshold: float = 0.03,
+        success_bonus: float = 0.0,
         obstacle_radius_range: Sequence[float] = (0.08, 0.15),
         obstacle_min_goal_distance: float = 0.25,
     ) -> None:
@@ -66,6 +66,8 @@ def __init__(
         radius_min, radius_max = map(float, obstacle_radius_range)
         if radius_min <= 0 or radius_max < radius_min:
             raise ValueError("Invalid obstacle_radius_range.")
+        if success_bonus < 0:
+            raise ValueError("success_bonus cannot be negative.")
 
         self.num_envs = int(num_envs)
         self.device = torch.device(device)
@@ -76,6 +78,7 @@ def __init__(
         self.linear_damping = float(linear_damping)
         self.dt = float(dt)
         self.success_threshold = float(success_threshold)
+        self.success_bonus = float(success_bonus)
         self.obstacle_radius_range = (radius_min, radius_max)
         self.obstacle_min_goal_distance = float(obstacle_min_goal_distance)
 
@@ -164,16 +167,17 @@ def step(self, action: torch.Tensor) -> tuple[
         self.collision_count = self.collision_count + collided.long()
 
         distance = (self.position - self.goal_position).norm(dim=-1)
+        terminated = distance < self.success_threshold
+        truncated = self.episode_step >= self.max_episode_steps
+        done = terminated | truncated
         reward = (
             -distance
             + 0.5 * torch.exp(-(distance.square()) / (2.0 * 0.05**2))
             - 2.0 * penetration.square().sum(dim=-1)
             - 0.01 * self.velocity.square().sum(dim=-1)
             - 0.001 * (action - self.last_action).square().sum(dim=-1)
+            + self.success_bonus * terminated.float()
         )
-        terminated = distance < self.success_threshold
-        truncated = self.episode_step >= self.max_episode_steps
-        done = terminated | truncated
 
         terminal_distance = distance.detach()
         terminal_path_length = self.path_length.detach().clone()
diff --git a/tests/learning/test_point_mass.py b/tests/learning/test_point_mass.py
index 853e76353..55fb06185 100644
--- a/tests/learning/test_point_mass.py
+++ b/tests/learning/test_point_mass.py
@@ -85,6 +85,20 @@ def test_point_mass_auto_reset_keeps_terminal_metrics() -> None:
     assert next_observation.shape == (2, 14)
 
 
+def test_point_mass_success_bonus_applies_only_on_success() -> None:
+    env = PointMassEnv(num_envs=2, success_bonus=5.0)
+    env.reset(seed=5)
+    env.velocity.zero_()
+    env.goal_position[0] = env.position[0]
+    env.goal_position[1] = env.position[1] + 0.5
+
+    _, reward, terminated, _, _ = env.step(torch.zeros(2, 2))
+
+    assert bool(terminated[0])
+    assert not bool(terminated[1])
+    assert float(reward[0]) > float(reward[1]) + 4.0
+
+
 def test_point_mass_detach_state_keeps_values() -> None:
     env = PointMassEnv(num_envs=3)
     env.reset(seed=9)

From 9f693837fd78b93ffadabab198cb68bf9972e19c Mon Sep 17 00:00:00 2001
From: yangchen73 <122090643@link.cuhk.edu.cn>
Date: Sun, 2 Aug 2026 04:07:48 +0800
Subject: [PATCH 4/7] Visualize PointMass policy

---
 docs/source/tutorial/rl.rst                   |  12 +
 .../embodichain_tasks/rl/basic/point_mass.py  |   8 +-
 scripts/tutorials/rl/visualize_point_mass.py  | 322 ++++++++++++++++++
 3 files changed, 339 insertions(+), 3 deletions(-)
 create mode 100644 scripts/tutorials/rl/visualize_point_mass.py

diff --git a/docs/source/tutorial/rl.rst b/docs/source/tutorial/rl.rst
index 54ce4590f..0cb6093fd 100644
--- a/docs/source/tutorial/rl.rst
+++ b/docs/source/tutorial/rl.rst
@@ -81,6 +81,18 @@ APG and PPO. Launch either config with the same CLI:
    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
 ---------------------
 
diff --git a/embodichain_tasks/embodichain_tasks/rl/basic/point_mass.py b/embodichain_tasks/embodichain_tasks/rl/basic/point_mass.py
index b40a1c429..a19f21d6f 100644
--- a/embodichain_tasks/embodichain_tasks/rl/basic/point_mass.py
+++ b/embodichain_tasks/embodichain_tasks/rl/basic/point_mass.py
@@ -34,9 +34,11 @@
 class PointMassEnv:
     """Navigate a damped point mass to a goal while avoiding two obstacles.
 
-    The imphementation is ale implementation is.aStandardfcollectofs execete
-e   it ntiable. Standard collectorwhile s exrctaine it urat `fram,lA icyrictpm action
-    roruugh gynnancr.d
+    Always differentiable: PPO/GRPO run under ``torch.no_grad()``, APG keeps
+    the graph. ``success_bonus`` is optional reward shaping for on-policy
+    methods that optimize dense distance rewards but are evaluated on strict
+    success; leave it at ``0`` for APG.
+    """
 
     observation_dim = 14
     action_dim = 2
diff --git a/scripts/tutorials/rl/visualize_point_mass.py b/scripts/tutorials/rl/visualize_point_mass.py
new file mode 100644
index 000000000..d3dcc6a9a
--- /dev/null
+++ b/scripts/tutorials/rl/visualize_point_mass.py
@@ -0,0 +1,322 @@
+# ----------------------------------------------------------------------------
+# 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.
+# ----------------------------------------------------------------------------
+
+"""Visualize a trained PointMass policy as 2D trajectory plots."""
+
+from __future__ import annotations
+
+import argparse
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+import matplotlib
+
+matplotlib.use("Agg")
+
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+from matplotlib.patches import Circle
+from tensordict import TensorDict
+
+from embodichain.learning.rl.models import Policy, build_mlp_from_cfg, build_policy
+from embodichain.utils.logger import log_info
+from embodichain.utils.utility import load_config
+from embodichain_tasks.rl.basic.point_mass import PointMassEnv
+
+
+@dataclass(frozen=True)
+class EpisodeTrajectory:
+    """One deterministic PointMass rollout used for plotting."""
+
+    positions: np.ndarray
+    goal: np.ndarray
+    obstacles: np.ndarray
+    radii: np.ndarray
+    success: bool
+    final_distance: float
+
+    @property
+    def num_steps(self) -> int:
+        return max(int(self.positions.shape[0]) - 1, 0)
+
+
+def build_policy_from_config(
+    policy_block: dict[str, Any],
+    env: PointMassEnv,
+    device: torch.device,
+) -> Policy:
+    """Build a PointMass policy matching the training config."""
+    obs_dim = int(env.single_observation_space.shape[-1])
+    action_dim = int(env.single_action_space.shape[-1])
+    actor_cfg = policy_block.get("actor")
+    critic_cfg = policy_block.get("critic")
+    actor = (
+        build_mlp_from_cfg(actor_cfg, obs_dim, action_dim)
+        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
+    )
+    return build_policy(
+        policy_block,
+        env.single_observation_space,
+        env.single_action_space,
+        device,
+        actor=actor,
+        critic=critic,
+    )
+
+
+def load_policy_checkpoint(
+    policy: Policy,
+    checkpoint_path: str | Path,
+    device: torch.device,
+) -> None:
+    """Load policy weights from a Trainer / DifferentiableTrainer checkpoint."""
+    path = Path(checkpoint_path)
+    try:
+        checkpoint = torch.load(path, map_location=device, weights_only=True)
+    except TypeError:
+        checkpoint = torch.load(path, map_location=device)
+    if "policy" not in checkpoint:
+        raise KeyError(f"Checkpoint {path} does not contain a 'policy' state dict.")
+    policy.load_state_dict(checkpoint["policy"])
+    policy.eval()
+
+
+@torch.no_grad()
+def rollout_episode(
+    env: PointMassEnv,
+    policy: Policy,
+    *,
+    seed: int,
+    env_index: int = 0,
+) -> EpisodeTrajectory:
+    """Roll out one deterministic episode and record the 2D trajectory."""
+    if not 0 <= env_index < env.num_envs:
+        raise ValueError(f"env_index out of range for num_envs={env.num_envs}.")
+
+    observation, _ = env.reset(seed=seed)
+    positions = [env.position[env_index].detach().cpu().numpy().copy()]
+    goal = env.goal_position[env_index].detach().cpu().numpy().copy()
+    obstacles = env.obstacle_position[env_index].detach().cpu().numpy().copy()
+    radii = env.obstacle_radius[env_index].detach().cpu().numpy().copy()
+
+    success = False
+    final_distance = float("nan")
+    for _ in range(env.max_episode_steps):
+        policy_input = TensorDict(
+            {"obs": observation},
+            batch_size=[env.num_envs],
+            device=env.device,
+        )
+        action = policy.get_action(policy_input, deterministic=True)["action"]
+        observation, _, terminated, truncated, info = env.step(action)
+        positions.append(env.position[env_index].detach().cpu().numpy().copy())
+        if bool(terminated[env_index] or truncated[env_index]):
+            success = bool(info["success"][env_index])
+            final_distance = float(info["metrics"]["final_distance"][env_index])
+            # Auto-reset already replaced state; drop the next-episode sample.
+            positions.pop()
+            break
+
+    return EpisodeTrajectory(
+        positions=np.asarray(positions, dtype=np.float32),
+        goal=np.asarray(goal, dtype=np.float32),
+        obstacles=np.asarray(obstacles, dtype=np.float32),
+        radii=np.asarray(radii, dtype=np.float32),
+        success=success,
+        final_distance=final_distance,
+    )
+
+
+def plot_episode(
+    episode: EpisodeTrajectory,
+    *,
+    title: str,
+    output_path: Path,
+    workspace_half: float,
+) -> None:
+    """Save a top-down trajectory plot for one episode."""
+    fig, ax = plt.subplots(figsize=(5.5, 5.5), dpi=160)
+    ax.set_aspect("equal")
+    ax.set_xlim(-workspace_half, workspace_half)
+    ax.set_ylim(-workspace_half, workspace_half)
+    ax.grid(True, alpha=0.25)
+    ax.set_xlabel("x")
+    ax.set_ylabel("y")
+
+    for center, radius in zip(episode.obstacles, episode.radii):
+        ax.add_patch(
+            Circle(
+                (float(center[0]), float(center[1])),
+                float(radius),
+                facecolor="#c44e52",
+                edgecolor="#8c2f32",
+                alpha=0.35,
+                zorder=1,
+            )
+        )
+
+    ax.plot(
+        episode.positions[:, 0],
+        episode.positions[:, 1],
+        color="#4c72b0",
+        linewidth=2.0,
+        zorder=2,
+        label="trajectory",
+    )
+    ax.scatter(
+        episode.positions[0, 0],
+        episode.positions[0, 1],
+        c="#4c72b0",
+        s=60,
+        marker="o",
+        zorder=3,
+        label="start",
+    )
+    ax.scatter(
+        episode.positions[-1, 0],
+        episode.positions[-1, 1],
+        c="#55a868",
+        s=70,
+        marker="*",
+        zorder=3,
+        label="end",
+    )
+    ax.scatter(
+        episode.goal[0],
+        episode.goal[1],
+        c="#dd8452",
+        s=80,
+        marker="x",
+        linewidths=2.0,
+        zorder=3,
+        label="goal",
+    )
+
+    status = "success" if episode.success else "fail"
+    ax.set_title(
+        f"{title}\n{status}, steps={episode.num_steps}, "
+        f"final_distance={episode.final_distance:.3f}"
+    )
+    ax.legend(loc="upper right", fontsize=8)
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    fig.tight_layout()
+    fig.savefig(output_path)
+    plt.close(fig)
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="Visualize a trained PointMass policy as 2D trajectory plots."
+    )
+    parser.add_argument(
+        "--config",
+        type=str,
+        required=True,
+        help="Training config used to build the env/policy (.yaml/.yml/.json).",
+    )
+    parser.add_argument(
+        "--checkpoint",
+        type=str,
+        required=True,
+        help="Checkpoint containing a 'policy' state dict.",
+    )
+    parser.add_argument(
+        "--output-dir",
+        type=str,
+        default="outputs/point_mass_viz",
+        help="Directory for saved PNG plots.",
+    )
+    parser.add_argument(
+        "--num-episodes",
+        type=int,
+        default=4,
+        help="Number of deterministic episodes to visualize.",
+    )
+    parser.add_argument(
+        "--seed",
+        type=int,
+        default=10042,
+        help="Base seed; episode i uses seed + i.",
+    )
+    parser.add_argument(
+        "--device",
+        type=str,
+        default="cpu",
+        help="Torch device for rollout (cpu or cuda:0).",
+    )
+    parser.add_argument(
+        "--tag",
+        type=str,
+        default="policy",
+        help="Filename prefix for saved plots.",
+    )
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    if args.num_episodes <= 0:
+        raise ValueError("--num-episodes must be positive.")
+
+    cfg = load_config(args.config)
+    if "policy" not in cfg:
+        raise KeyError(f"Config {args.config} is missing a 'policy' block.")
+    trainer_cfg = cfg.get("trainer", {})
+    env_block = trainer_cfg.get("learning_env", {})
+    env_cfg = {} if isinstance(env_block, str) else dict(env_block.get("cfg", {}))
+
+    device = torch.device(args.device)
+    env = PointMassEnv(num_envs=1, device=device, **env_cfg)
+    policy = build_policy_from_config(cfg["policy"], env, device)
+    load_policy_checkpoint(policy, args.checkpoint, device)
+
+    output_dir = Path(args.output_dir)
+    successes = 0
+    try:
+        for episode_idx in range(args.num_episodes):
+            episode = rollout_episode(
+                env,
+                policy,
+                seed=args.seed + episode_idx,
+            )
+            path = output_dir / f"{args.tag}_ep{episode_idx}.png"
+            plot_episode(
+                episode,
+                title=f"{args.tag} episode {episode_idx}",
+                output_path=path,
+                workspace_half=env.workspace_half,
+            )
+            successes += int(episode.success)
+            log_info(
+                f"saved {path} success={episode.success} "
+                f"distance={episode.final_distance:.4f} steps={episode.num_steps}"
+            )
+    finally:
+        env.close()
+
+    log_info(
+        f"finished {successes}/{args.num_episodes} successful episodes -> {output_dir}"
+    )
+
+
+if __name__ == "__main__":
+    main()

From c4e611d938244f988c39ce5f428fff473561e202 Mon Sep 17 00:00:00 2001
From: yangchen73 <122090643@link.cuhk.edu.cn>
Date: Sun, 2 Aug 2026 04:26:10 +0800
Subject: [PATCH 5/7] WIP

---
 embodichain/learning/rl/differentiable_trainer.py     |  3 +++
 embodichain/learning/rl/train.py                      |  2 +-
 embodichain/learning/rl/utils/config.py               |  2 +-
 .../embodichain_tasks/rl/basic/point_mass.py          |  5 ++++-
 scripts/benchmark/rl/runtime.py                       |  9 ++++++++-
 scripts/tutorials/rl/visualize_point_mass.py          |  8 +++++---
 tests/learning/test_point_mass.py                     | 11 +++++++++++
 7 files changed, 33 insertions(+), 7 deletions(-)

diff --git a/embodichain/learning/rl/differentiable_trainer.py b/embodichain/learning/rl/differentiable_trainer.py
index fe2472f85..a69cee762 100644
--- a/embodichain/learning/rl/differentiable_trainer.py
+++ b/embodichain/learning/rl/differentiable_trainer.py
@@ -263,9 +263,12 @@ def load_checkpoint(self, path: str | Path) -> None:
 
     def get_summary(self) -> dict[str, Any]:
         """Return the current in-memory training summary."""
+        elapsed = max(1e-6, time.time() - self.start_time)
         return {
             "global_step": self.global_step,
             "num_updates": self.num_updates,
+            "elapsed_time_sec": float(elapsed),
+            "training_fps": float(self.global_step / elapsed),
             "last_train_metrics": (
                 dict(self.train_history[-1]) if self.train_history else {}
             ),
diff --git a/embodichain/learning/rl/train.py b/embodichain/learning/rl/train.py
index 47677ca68..8ec255051 100644
--- a/embodichain/learning/rl/train.py
+++ b/embodichain/learning/rl/train.py
@@ -193,7 +193,7 @@ def _train_learning_env(
     use_wandb = bool(trainer_cfg.get("use_wandb", False))
     if use_wandb:
         wandb.init(
-            project=trainer_cfg.get("wandb_project_name", "embodichain-point-mass"),
+            project=trainer_cfg.get("wandb_project_name", "embodichain-generic"),
             name=exp_name,
             config=cfg_data,
         )
diff --git a/embodichain/learning/rl/utils/config.py b/embodichain/learning/rl/utils/config.py
index d1d491c32..d5c838dc7 100644
--- a/embodichain/learning/rl/utils/config.py
+++ b/embodichain/learning/rl/utils/config.py
@@ -48,7 +48,7 @@ class LRSchedulerCfg:
 class AlgorithmCfg:
     """Shared fields for RL algorithm configs."""
 
-    device: str = "cuda"
+    device: str = "cpu"
     optimizer: OptimizerCfg = OptimizerCfg()
     lr_scheduler: LRSchedulerCfg = LRSchedulerCfg()
     batch_size: int = 64
diff --git a/embodichain_tasks/embodichain_tasks/rl/basic/point_mass.py b/embodichain_tasks/embodichain_tasks/rl/basic/point_mass.py
index a19f21d6f..9fc5925f5 100644
--- a/embodichain_tasks/embodichain_tasks/rl/basic/point_mass.py
+++ b/embodichain_tasks/embodichain_tasks/rl/basic/point_mass.py
@@ -100,7 +100,7 @@ def __init__(
         self.action_space = self.single_action_space
 
         self._generator = torch.Generator(device=self.device)
-        self._generator.manual_seed(torch.seed())
+        self._generator.manual_seed(torch.initial_seed())
         self.position = torch.zeros(self.num_envs, 2, device=self.device)
         self.velocity = torch.zeros_like(self.position)
         self.goal_position = torch.zeros_like(self.position)
@@ -181,12 +181,15 @@ def step(self, action: torch.Tensor) -> tuple[
             + self.success_bonus * terminated.float()
         )
 
+        # Capture terminal state before auto-reset replaces done rows.
+        terminal_position = self.position.detach().clone()
         terminal_distance = distance.detach()
         terminal_path_length = self.path_length.detach().clone()
         terminal_collisions = self.collision_count.detach().clone()
         info = {
             "success": terminated.detach(),
             "metrics": {
+                "final_position": terminal_position,
                 "final_distance": terminal_distance,
                 "path_length": terminal_path_length,
                 "collision_count": terminal_collisions,
diff --git a/scripts/benchmark/rl/runtime.py b/scripts/benchmark/rl/runtime.py
index 351bdb819..216d98d7c 100644
--- a/scripts/benchmark/rl/runtime.py
+++ b/scripts/benchmark/rl/runtime.py
@@ -489,7 +489,14 @@ def evaluate_checkpoint(
             eval_rollout_buffer = _allocate_eval_rollout_buffer(env, policy, device)
             env.set_rollout_buffer(eval_rollout_buffer)
 
-        checkpoint = torch.load(checkpoint_path, map_location=device)
+        try:
+            checkpoint = torch.load(
+                checkpoint_path,
+                map_location=device,
+                weights_only=True,
+            )
+        except TypeError:
+            checkpoint = torch.load(checkpoint_path, map_location=device)
         policy.load_state_dict(checkpoint["policy"])
         policy.eval()
 
diff --git a/scripts/tutorials/rl/visualize_point_mass.py b/scripts/tutorials/rl/visualize_point_mass.py
index d3dcc6a9a..1cfb0fd8d 100644
--- a/scripts/tutorials/rl/visualize_point_mass.py
+++ b/scripts/tutorials/rl/visualize_point_mass.py
@@ -128,13 +128,15 @@ def rollout_episode(
         )
         action = policy.get_action(policy_input, deterministic=True)["action"]
         observation, _, terminated, truncated, info = env.step(action)
-        positions.append(env.position[env_index].detach().cpu().numpy().copy())
         if bool(terminated[env_index] or truncated[env_index]):
+            # Auto-reset replaces env.position; use the terminal snapshot in info.
+            positions.append(
+                info["metrics"]["final_position"][env_index].detach().cpu().numpy().copy()
+            )
             success = bool(info["success"][env_index])
             final_distance = float(info["metrics"]["final_distance"][env_index])
-            # Auto-reset already replaced state; drop the next-episode sample.
-            positions.pop()
             break
+        positions.append(env.position[env_index].detach().cpu().numpy().copy())
 
     return EpisodeTrajectory(
         positions=np.asarray(positions, dtype=np.float32),
diff --git a/tests/learning/test_point_mass.py b/tests/learning/test_point_mass.py
index 55fb06185..fd3fdab48 100644
--- a/tests/learning/test_point_mass.py
+++ b/tests/learning/test_point_mass.py
@@ -75,12 +75,23 @@ def test_point_mass_auto_reset_keeps_terminal_metrics() -> None:
     env.reset(seed=5)
     env.velocity.zero_()
     env.goal_position[0] = env.position[0]
+    terminal_goal = env.goal_position[0].clone()
 
     next_observation, _, terminated, truncated, info = env.step(torch.zeros(2, 2))
 
     assert bool(terminated[0])
     assert not bool(truncated[0])
     assert float(info["metrics"]["final_distance"][0]) < env.success_threshold
+    torch.testing.assert_close(
+        info["metrics"]["final_position"][0],
+        terminal_goal,
+        atol=env.success_threshold,
+        rtol=0.0,
+    )
+    # Auto-reset replaced the live state; terminal metrics stay pre-reset.
+    assert not torch.allclose(
+        env.position[0], info["metrics"]["final_position"][0], atol=1e-6
+    )
     assert env.episode_step[0] == 0
     assert next_observation.shape == (2, 14)
 

From bbc15753b88f75c5267b67e17e297dd4addd7147 Mon Sep 17 00:00:00 2001
From: yangchen73 <122090643@link.cuhk.edu.cn>
Date: Sun, 2 Aug 2026 04:51:38 +0800
Subject: [PATCH 6/7] WIP

---
 embodichain/learning/rl/evaluation.py         | 24 ++++++++++++++++++-
 scripts/tutorials/rl/visualize_point_mass.py  |  6 ++++-
 tests/learning/test_differentiable_trainer.py |  7 +++---
 tests/learning/test_evaluation.py             | 10 +++++++-
 4 files changed, 41 insertions(+), 6 deletions(-)

diff --git a/embodichain/learning/rl/evaluation.py b/embodichain/learning/rl/evaluation.py
index d64ebcfe3..974a2ab7c 100644
--- a/embodichain/learning/rl/evaluation.py
+++ b/embodichain/learning/rl/evaluation.py
@@ -58,6 +58,26 @@ def _selected_values(value: Any, indices: torch.Tensor) -> list[float]:
     return [float(array[index]) for index in indices.cpu().tolist()]
 
 
+def _is_per_env_scalar_metric(value: Any, num_envs: int) -> bool:
+    """Return True for metrics with one scalar per env (shape ``[N]`` / ``[N, 1]``)."""
+    if isinstance(value, torch.Tensor):
+        if value.numel() == 0:
+            return False
+        if value.ndim == 0:
+            return num_envs == 1
+        if value.shape[0] != num_envs:
+            return False
+        return value.ndim == 1 or value.shape[1:] in {(1,)}
+    array = np.asarray(value)
+    if array.size == 0:
+        return False
+    if array.ndim == 0:
+        return num_envs == 1
+    if array.shape[0] != num_envs:
+        return False
+    return array.ndim == 1 or array.shape[1:] in {(1,)}
+
+
 @torch.no_grad()
 def evaluate_episodes(
     *,
@@ -121,11 +141,13 @@ def evaluate_episodes(
                 metrics = info.get("metrics", {})
                 if isinstance(metrics, Mapping):
                     for name, value in metrics.items():
+                        # Skip vector/matrix metrics: reshape(-1) would break env indexing.
+                        if not _is_per_env_scalar_metric(value, num_envs):
+                            continue
                         metric_values.setdefault(str(name), []).extend(
                             _selected_values(value, selected)
                         )
 
-
             current_return[done_indices] = 0.0
             current_length[done_indices] = 0
     finally:
diff --git a/scripts/tutorials/rl/visualize_point_mass.py b/scripts/tutorials/rl/visualize_point_mass.py
index 1cfb0fd8d..c94287f69 100644
--- a/scripts/tutorials/rl/visualize_point_mass.py
+++ b/scripts/tutorials/rl/visualize_point_mass.py
@@ -131,7 +131,11 @@ def rollout_episode(
         if bool(terminated[env_index] or truncated[env_index]):
             # Auto-reset replaces env.position; use the terminal snapshot in info.
             positions.append(
-                info["metrics"]["final_position"][env_index].detach().cpu().numpy().copy()
+                info["metrics"]["final_position"][env_index]
+                .detach()
+                .cpu()
+                .numpy()
+                .copy()
             )
             success = bool(info["success"][env_index])
             final_distance = float(info["metrics"]["final_distance"][env_index])
diff --git a/tests/learning/test_differentiable_trainer.py b/tests/learning/test_differentiable_trainer.py
index 0fc48f9ed..a36023eed 100644
--- a/tests/learning/test_differentiable_trainer.py
+++ b/tests/learning/test_differentiable_trainer.py
@@ -281,9 +281,10 @@ def test_checkpoint_restores_lr_scheduler_state(tmp_path: Path) -> None:
 
     assert restored_algorithm.lr_scheduler is not None
     assert restored_algorithm.current_learning_rate() == pytest.approx(lr_before)
-    assert restored_algorithm.lr_scheduler.state_dict()["last_epoch"] == scheduler_state[
-        "last_epoch"
-    ]
+    assert (
+        restored_algorithm.lr_scheduler.state_dict()["last_epoch"]
+        == scheduler_state["last_epoch"]
+    )
 
 
 def test_checkpoint_load_falls_back_for_older_torch(
diff --git a/tests/learning/test_evaluation.py b/tests/learning/test_evaluation.py
index 9afa8560b..991a139e8 100644
--- a/tests/learning/test_evaluation.py
+++ b/tests/learning/test_evaluation.py
@@ -57,7 +57,14 @@ def step(self, action):
         terminal_steps = self.steps.clone()
         info = {
             "success": done.clone(),
-            "metrics": {"terminal_step": terminal_steps.float()},
+            "metrics": {
+                "terminal_step": terminal_steps.float(),
+                # Non-scalar per-env metric must not be flattened into eval logs.
+                "final_position": torch.stack(
+                    (terminal_steps.float(), terminal_steps.float() + 10.0),
+                    dim=-1,
+                ),
+            },
         }
         self.steps[done] = 0
         return (
@@ -88,3 +95,4 @@ def test_evaluate_episodes_counts_actual_completions_and_restores_mode() -> None
     assert result["eval/avg_length"] == 1.25
     assert result["eval/success_rate"] == 1.0
     assert result["eval/metrics/terminal_step"] == 1.25
+    assert "eval/metrics/final_position" not in result

From dcfc9f3859f37d19ce88c0fc024ef5f03ae99de5 Mon Sep 17 00:00:00 2001
From: yangchen73 <122090643@link.cuhk.edu.cn>
Date: Sun, 2 Aug 2026 06:30:57 +0800
Subject: [PATCH 7/7] Fix Newton and PointMass environments

---
 embodichain/learning/rl/env.py                | 27 +++++++++++++++++--
 .../rl/experimental/newton/planar_reach.py    |  3 +++
 2 files changed, 28 insertions(+), 2 deletions(-)

diff --git a/embodichain/learning/rl/env.py b/embodichain/learning/rl/env.py
index 5bf3716a4..d3c51299c 100644
--- a/embodichain/learning/rl/env.py
+++ b/embodichain/learning/rl/env.py
@@ -88,6 +88,23 @@ def detach_state(self) -> DifferentiableObservation:
         ...
 
 
+def _is_nested_package_shadow(existing: Any, candidate: Any) -> bool:
+    """Return True when ``candidate`` is an editable-install nested duplicate.
+
+    Editable installs can expose both ``embodichain_tasks.rl...`` and
+    ``embodichain_tasks.embodichain_tasks.rl...``. Prefer the shorter canonical
+    module path already registered.
+    """
+    existing_module = getattr(existing, "__module__", "") or ""
+    candidate_module = getattr(candidate, "__module__", "") or ""
+    if not existing_module or not candidate_module:
+        return False
+    nested_prefix = existing_module.partition(".")[0] + "." + existing_module
+    return candidate_module == nested_prefix or candidate_module.startswith(
+        nested_prefix + "."
+    )
+
+
 def register_learning_env(
     name: str,
     factory: LearningEnvFactory | None = None,
@@ -102,8 +119,14 @@ def register_learning_env(
 
     def decorator(env_factory: LearningEnvFactory) -> LearningEnvFactory:
         key = name.lower()
-        if key in _LEARNING_ENV_REGISTRY and not override:
-            raise ValueError(f"Learning environment '{name}' is already registered.")
+        if key in _LEARNING_ENV_REGISTRY:
+            existing = _LEARNING_ENV_REGISTRY[key]
+            if _is_nested_package_shadow(existing, env_factory):
+                return env_factory
+            if not override:
+                raise ValueError(
+                    f"Learning environment '{name}' is already registered."
+                )
         _LEARNING_ENV_REGISTRY[key] = env_factory
         return env_factory
 
diff --git a/embodichain/learning/rl/experimental/newton/planar_reach.py b/embodichain/learning/rl/experimental/newton/planar_reach.py
index 09f437f5f..41c3e5adf 100644
--- a/embodichain/learning/rl/experimental/newton/planar_reach.py
+++ b/embodichain/learning/rl/experimental/newton/planar_reach.py
@@ -347,6 +347,9 @@ def detach_state(self) -> torch.Tensor:
         self._last_action = self._last_action.detach()
         return self._get_observation().detach()
 
+    def close(self) -> None:
+        """No-op; Warp/Newton resources are released with the Python object."""
+
     def _build_model(self) -> tuple[Any, wp.array]:
         template = newton.ModelBuilder(gravity=0.0, up_axis=newton.Axis.Y)
         first_link = template.add_link(label="planar_first_link")