Reinforcement Learning Training

This tutorial shows you how to train reinforcement learning agents using EmbodiChain’s RL framework. You’ll learn how to configure training via JSON or YAML, set up environments, policies, and algorithms, and launch training sessions.

Overview

The RL framework provides a modular, extensible stack for robotics tasks:

• Trainer: Orchestrates the training loop (calls algorithm for data collection and updates, handles logging/eval/save)

• Algorithm: Controls data collection process (interacts with environment, fills buffer, computes advantages/returns) and updates the policy (e.g., PPO)

• Policy: Neural network models implementing a unified interface (get_action/get_value/evaluate_actions)

• Buffer: On-policy rollout storage and minibatch iterator (managed by algorithm)

• Env Factory: Build environments from a JSON or YAML config via registry

Architecture

The framework follows a clean separation of concerns:

• Trainer: Orchestrates the training loop (calls algorithm for data collection and updates, handles logging/eval/save)

• Algorithm: Controls data collection process (interacts with environment, fills buffer, computes advantages/returns) and updates the policy (e.g., PPO)

• Policy: Neural network models implementing a unified interface

• Buffer: On-policy rollout storage and minibatch iterator (managed by algorithm)

• Env Factory: Build environments from a JSON or YAML config via registry

The core components and their relationships:

• Trainer → Policy, Env, Algorithm (via callbacks for statistics)

• Algorithm → Policy, RolloutBuffer (algorithm manages its own buffer)

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(). 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

Training configurations are task-local: PushCube uses embodichain_tasks/configs/tasks/manipulation/push_cube/agents and CartPole uses embodichain_tasks/configs/tasks/classic_control/cart_pole/agents. PointMass APG/PPO configs live under embodichain_tasks/configs/tasks/classic_control/point_mass/agents:

Example: train_config.json

Example: train_config.yaml (CartPole)

Example: train_apg.yaml (PointMass)

PointMass uses one differentiable PyTorch dynamics implementation for both APG and PPO. Launch either config with the same CLI:

embodichain train-rl --config embodichain_tasks/configs/tasks/classic_control/point_mass/agents/apg.yaml
embodichain train-rl --config embodichain_tasks/configs/tasks/classic_control/point_mass/agents/ppo.yaml

Configuration Sections

Runtime Settings

The trainer section controls experiment setup:

• exp_name: Experiment name (used for output directories)

• seed: Random seed for reproducibility

• device: Runtime device string, e.g. "cpu" or "cuda:0"

• headless: Whether to run simulation in headless mode

• iterations: Number of training iterations

• buffer_size: Steps collected per rollout (e.g., 1024)

• eval_freq: Frequency of evaluation (in steps)

• save_freq: Frequency of checkpoint saving (in steps)

• use_wandb: Whether to enable Weights & Biases logging (set in the config file)

• wandb_project_name: Weights & Biases project name

Environment Configuration

The env section defines the task environment:

• id: Environment registry ID (e.g., “PushCubeRL”)

• cfg: Environment-specific configuration parameters

For RL environments, use the actions field for action preprocessing and extensions for task-specific parameters:

• actions: Ordered flat Action Manager config (for example, RelativeJointPositionAction with a selected control part and scale)

• extensions: Task-specific parameters (e.g., success_threshold)

Example:

"env": {
  "id": "PushCubeRL",
  "cfg": {
    "num_envs": 4,
    "actions": {
      "delta_qpos": {
        "func": "RelativeJointPositionAction",
        "params": { "part_name": "arm", "scale": 0.1 }
      }
    },
    "extensions": {
      "success_threshold": 0.1
    }
  }
}

Policy Configuration

The policy section defines the neural network policy:

• name: Policy name (e.g., “actor_critic”, “vla”)

• action_dim: Optional policy output action dimension. If omitted, it is inferred from env.action_space.

• actor: Actor network configuration (required for actor_critic)

• critic: Critic network configuration (required for actor_critic)

Example:

"policy": {
  "name": "actor_critic",
  "actor": {
    "type": "mlp",
    "network_cfg": {
      "hidden_sizes": [256, 256],
      "activation": "relu"
    }
  },
  "critic": {
    "type": "mlp",
    "network_cfg": {
      "hidden_sizes": [256, 256],
      "activation": "relu"
    }
  }
}

Algorithm Configuration

The algorithm section defines the RL algorithm:

• name: Algorithm name (e.g., “ppo”, “grpo”)

• cfg: Algorithm-specific hyperparameters

PPO example:

"algorithm": {
  "name": "ppo",
  "cfg": {
    "learning_rate": 0.0001,
    "n_epochs": 10,
    "batch_size": 64,
    "gamma": 0.99,
    "gae_lambda": 0.95,
    "clip_coef": 0.2,
    "ent_coef": 0.01,
    "vf_coef": 0.5,
    "max_grad_norm": 0.5
  }
}

GRPO example (for Embodied AI / from-scratch training, e.g. CartPole):

"algorithm": {
  "name": "grpo",
  "cfg": {
    "learning_rate": 0.0001,
    "n_epochs": 10,
    "batch_size": 8192,
    "gamma": 0.99,
    "clip_coef": 0.2,
    "ent_coef": 0.001,
    "kl_coef": 0,
    "group_size": 4,
    "eps": 1e-8,
    "reset_every_rollout": true,
    "max_grad_norm": 0.5,
    "truncate_at_first_done": true
  }
}

For GRPO: use actor_only policy. Set kl_coef=0 for from-scratch training; kl_coef=0.02 for VLA/LLM fine-tuning.

Training Script

The training script (train.py) is located in embodichain/learning/rl/:

Code for train.py

# ----------------------------------------------------------------------------
# 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 argparse
import json
import os
import random
import time
from collections.abc import Sequence
from copy import deepcopy
from pathlib import Path

import numpy as np
import torch
import wandb
import yaml
from torch.utils.tensorboard import SummaryWriter

from embodichain.learning.rl.checkpoint import (
    _atomic_write,
    _resolve_resume,
    _sync_errors,
    _validate_saved_files,
)
from embodichain.learning.rl.policies import get_registered_policy_names
from embodichain.learning.rl.policy_evaluation.manifest import (
    write_run_manifest,
)
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.runtime import (
    _close_environment,
    _build_learning_environment,
    _build_learning_policy,
    build_gym_policy_runtime,
    build_learning_policy_runtime,
)
from embodichain.learning.rl.routing import get_trainer_class
from embodichain.learning.rl.utils.trainer import Trainer
from embodichain.utils import logger, set_seed
from embodichain.lab.gym.utils.registration import (
    build_env,
    discover_task_packages,
    execute_init_hooks,
)
from embodichain.lab.gym.utils.profiler import EnvProfilerCfg
from embodichain.utils.utility import load_config
from embodichain.utils.module_utils import find_function_from_modules
from embodichain.lab.gym.envs.managers.cfg import EventCfg

_CAMERA_RECORDERS = {"record_camera_data", "record_camera_data_async"}


def _seed_training_rng(seed: int, device: torch.device) -> None:
    """Seed policy/trainer RNGs without changing deterministic-kernel settings."""
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if device.type == "cuda":
        torch.cuda.manual_seed_all(seed)


def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
    """Parse command-line arguments.

    Args:
        argv: Arguments excluding the command name. Uses ``sys.argv`` when
            omitted.

    Returns:
        Parsed training arguments.
    """
    parser = argparse.ArgumentParser(
        prog="embodichain train-rl",
        description="Train an RL agent from a JSON or YAML config.",
    )
    parser.add_argument(
        "--config",
        type=str,
        required=True,
        help="Path to training config file (.json, .yaml, or .yml).",
    )
    parser.add_argument(
        "--distributed",
        action=argparse.BooleanOptionalAction,
        default=None,
        help="Enable or disable multi-GPU distributed training",
    )
    parser.add_argument(
        "--profile",
        action="store_true",
        default=False,
        help=(
            "Enable per-section time profiling of gym env reset/step "
            "(report on env.close()). Requires trainer.gym_config."
        ),
    )
    parser.add_argument(
        "--profile_output",
        type=str,
        default=None,
        help="Dump the profiling report as JSON on env.close() (requires --profile).",
    )
    parser.add_argument(
        "--resume",
        type=str,
        default=None,
        help="Resume a checkpoint file, or latest/best selected from --run.",
    )
    parser.add_argument(
        "--run",
        type=str,
        default=None,
        help="Source run directory for --resume latest/best.",
    )
    parser.add_argument(
        "--resume-scheduler",
        choices=("restore", "restart"),
        default="restore",
        help="Restore the saved LR curve, or restart it with the current config.",
    )
    args = parser.parse_args(argv)
    if args.resume in {"latest", "best"} and args.run is None:
        parser.error("--resume latest/best requires --run")
    if args.run is not None and args.resume not in {"latest", "best"}:
        parser.error("--run requires --resume latest/best")
    if args.resume is None and args.resume_scheduler != "restore":
        parser.error("--resume-scheduler requires --resume")
    return args


def _resolve_profile_output(
    path: str | None,
    *,
    rank: int,
    world_size: int,
) -> str | None:
    if path is None or world_size <= 1:
        return path
    output = Path(path)
    return str(output.with_name(f"{output.stem}_rank{rank}{output.suffix}"))


def _event_params(
    event_info: dict,
    *,
    run_base: str | Path,
    phase: str,
) -> dict:
    """Place default camera recordings under the current training run."""
    params = dict(event_info.get("params", {}))
    function_name = str(event_info.get("func", "")).rsplit(".", 1)[-1]
    if function_name in _CAMERA_RECORDERS:
        params.setdefault("save_path", str(Path(run_base) / "videos" / phase))
    return params


def _train_learning_env(
    cfg_data: dict,
    *,
    config_path: str | Path,
    distributed: bool | None,
    profile: bool = False,
    resume: str | Path | None = None,
    run: str | Path | None = None,
    resume_scheduler: str = "restore",
) -> dict[str, object]:
    """Train a lightweight registered environment through the unified CLI."""
    if profile:
        raise ValueError(
            "--profile requires trainer.gym_config; learning_env is unsupported."
        )
    trainer_cfg = cfg_data["trainer"]
    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)
    set_seed(
        seed,
        deterministic=bool(trainer_cfg.get("torch_deterministic", False)),
    )

    num_envs = int(trainer_cfg.get("num_envs", 64))
    runtime = build_learning_policy_runtime(
        cfg_data,
        num_envs=num_envs,
        device=device,
    )
    env = runtime.env
    policy = runtime.policy
    env_name = runtime.env_id

    enable_eval = bool(trainer_cfg.get("enable_eval", False))
    eval_env = None
    if enable_eval:
        _eval_name, eval_env = _build_learning_environment(
            cfg_data,
            num_envs=int(trainer_cfg.get("num_eval_envs", 16)),
            device=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 = f"{time.strftime('%Y%m%d_%H%M%S')}_{time.time_ns() % 1_000_000_000:09d}"
    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-generic"),
            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,
                rollout_mode=str(trainer_cfg.get("rollout_mode", "segmented")),
                gradient_accumulation_steps=int(
                    trainer_cfg.get("gradient_accumulation_steps", 1)
                ),
                deterministic_actions=bool(
                    trainer_cfg.get("deterministic_actions", False)
                ),
                clip_actions_to_space=bool(
                    trainer_cfg.get("clip_actions_to_space", False)
                ),
                action_adjoint_max_norm=float(
                    trainer_cfg.get("action_adjoint_max_norm", 0.0)
                ),
                normalize_observations=bool(
                    trainer_cfg.get("normalize_observations", False)
                ),
                rollout_seed=seed,
                checkpoint_dir=str(checkpoint_dir),
                experiment_name=exp_name,
                save_frequency_updates=int(
                    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
                * diff_cfg.gradient_accumulation_steps
            )
        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"),
                save_frequency_updates=int(
                    trainer_cfg.get("save_frequency_updates", 0)
                ),
            )
            default_steps = iterations * buffer_size * num_envs
        _configure_checkpoints(trainer, cfg_data, config_path, run_base)
        _resume_training(trainer, resume, run, resume_scheduler)
        if (
            trainer_class is DifferentiableTrainer
            and diff_cfg.rollout_mode == "complete"
            and "total_timesteps" not in trainer_cfg
        ):
            trainer.train(total_updates=iterations)
        else:
            total_timesteps = int(trainer_cfg.get("total_timesteps", default_steps))
            trainer.train(total_timesteps)
        _save_final_checkpoint(trainer)
        summary = trainer.get_summary()
    finally:
        writer.close()
        if use_wandb:
            wandb.finish()
        env.close()
        if eval_env is not None:
            eval_env.close()
    return summary


def train_from_config(
    config_path: str,
    distributed: bool | None = None,
    *,
    profile: bool = False,
    profile_output: str | None = None,
    resume: str | Path | None = None,
    run: str | Path | None = None,
    resume_scheduler: str = "restore",
) -> dict[str, object] | None:
    """Run training from a config file path.

    Args:
        config_path: Path to the training config file (.json, .yaml, or .yml).
        distributed: If True, run multi-GPU distributed training.
            If None, use trainer.distributed from config.
        profile: Enable gym ``EnvProfiler`` on the training environment.
        profile_output: Optional JSON dump path for the profiling report.
        resume: Checkpoint file, or latest/best selected from run.
        run: Source run directory when resume selects latest/best.
        resume_scheduler: Restore the saved LR curve or restart the configured
            curve with optimizer momentum retained.

    Returns:
        The lightweight trainer summary, or ``None`` for simulator training.
    """
    if profile_output is not None and not profile:
        raise ValueError("--profile_output requires --profile.")

    if resume_scheduler not in {"restore", "restart"}:
        raise ValueError("resume_scheduler must be 'restore' or 'restart'.")
    if resume is None and resume_scheduler != "restore":
        raise ValueError("resume_scheduler restart requires resume.")
    if resume is None and run is not None:
        raise ValueError("run requires resume latest/best.")
    cfg_data = load_config(config_path)

    trainer_cfg = cfg_data["trainer"]
    if "learning_env" in trainer_cfg:
        return _train_learning_env(
            cfg_data,
            config_path=config_path,
            distributed=distributed,
            profile=profile,
            resume=resume,
            run=run,
            resume_scheduler=resume_scheduler,
        )
    policy_block = cfg_data["policy"]
    algo_block = cfg_data["algorithm"]

    if distributed is None:
        distributed = bool(trainer_cfg.get("distributed", False))

    rank = 0
    world_size = 1
    local_rank = 0
    if distributed:
        if not torch.distributed.is_available():
            raise RuntimeError(
                "Distributed training requested but torch.distributed is not available."
            )
        if not torch.cuda.is_available():
            raise RuntimeError(
                "Distributed training with NCCL backend requires CUDA, "
                "but torch.cuda.is_available() is False."
            )
        local_rank = int(os.environ.get("LOCAL_RANK", 0))
        if local_rank < 0 or local_rank >= torch.cuda.device_count():
            raise ValueError(
                f"LOCAL_RANK {local_rank} is out of range "
                f"(available GPUs: {torch.cuda.device_count()})."
            )
        torch.cuda.set_device(local_rank)
        if not torch.distributed.is_initialized():
            torch.distributed.init_process_group(backend="nccl")
        rank = torch.distributed.get_rank()
        world_size = torch.distributed.get_world_size()

    exp_name = trainer_cfg.get("exp_name", "generic_exp")
    seed = int(trainer_cfg.get("seed", 1))
    device_str = trainer_cfg.get("device", "cpu")
    if distributed:
        device_str = f"cuda:{local_rank}"
    iterations = int(trainer_cfg.get("iterations", 250))
    buffer_size = int(
        trainer_cfg.get("buffer_size", trainer_cfg.get("rollout_steps", 2048))
    )
    enable_eval = bool(trainer_cfg.get("enable_eval", False))
    eval_freq = int(trainer_cfg.get("eval_freq", 10000))
    save_freq = int(
        trainer_cfg.get(
            "save_freq", 0 if "save_frequency_updates" in trainer_cfg else 50000
        )
    )
    num_eval_episodes = int(trainer_cfg.get("num_eval_episodes", 5))
    eval_seed = int(trainer_cfg.get("eval_seed", seed + 10_000))
    headless = bool(trainer_cfg.get("headless", True))
    renderer = trainer_cfg.get("renderer")
    gpu_id = int(trainer_cfg.get("gpu_id", 0))
    num_envs = trainer_cfg.get("num_envs", None)
    wandb_project_name = trainer_cfg.get("wandb_project_name", "embodichain-generic")

    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}"
        )
    try:
        device = torch.device(device_str)
    except RuntimeError as exc:
        raise ValueError(
            f"Failed to parse runtime.device='{device_str}': {exc}"
        ) from exc

    if device.type == "cuda":
        if not torch.cuda.is_available():
            raise ValueError(
                "CUDA device requested but torch.cuda.is_available() is False."
            )
        index = (
            device.index if device.index is not None else torch.cuda.current_device()
        )
        device_count = torch.cuda.device_count()
        if index < 0 or index >= device_count:
            raise ValueError(
                f"CUDA device index {index} is out of range (available devices: {device_count})."
            )
        torch.cuda.set_device(index)
        device = torch.device(f"cuda:{index}")
    elif device.type != "cpu":
        raise ValueError(f"Unsupported device type: {device}")
    if rank == 0:
        logger.log_info(f"Device: {device}")
    if distributed and rank == 0:
        logger.log_info(f"Distributed training: world_size={world_size}")

    # Seeds
    effective_seed = seed + rank
    _seed_training_rng(effective_seed, device)
    torch.backends.cudnn.deterministic = True

    # Outputs
    if distributed:
        run_stamp = (
            f"{time.strftime('%Y%m%d_%H%M%S')}_{time.time_ns() % 1_000_000_000:09d}"
            if rank == 0
            else None
        )
        run_stamp_list = [run_stamp]
        torch.distributed.broadcast_object_list(run_stamp_list, src=0)
        run_stamp = run_stamp_list[0]
    else:
        run_stamp = (
            f"{time.strftime('%Y%m%d_%H%M%S')}_{time.time_ns() % 1_000_000_000:09d}"
        )
    run_base = os.path.join("outputs", f"{exp_name}_{run_stamp}")
    log_dir = os.path.join(run_base, "logs")
    checkpoint_dir = os.path.join(run_base, "checkpoints")
    if rank == 0:
        os.makedirs(log_dir, exist_ok=True)
        os.makedirs(checkpoint_dir, exist_ok=True)
    writer = SummaryWriter(f"{log_dir}/{exp_name}") if rank == 0 else None

    # Initialize Weights & Biases (optional)
    use_wandb = trainer_cfg.get("use_wandb", False)
    if use_wandb and rank == 0:
        wandb.init(project=wandb_project_name, name=exp_name, config=cfg_data)

    if rank == 0:
        logger.log_info(f"Current working directory: {Path.cwd()}")

    profiler = (
        EnvProfilerCfg(
            enable_time=True,
            output_path=_resolve_profile_output(
                profile_output,
                rank=rank,
                world_size=world_size,
            ),
        )
        if profile
        else None
    )
    # Resolve active event callables before constructing the first World.
    event_modules = [
        "embodichain.lab.gym.envs.managers.randomization",
        "embodichain.lab.gym.envs.managers.record",
        "embodichain.lab.gym.envs.managers.events",
    ]
    events_dict = trainer_cfg.get("events", {})
    train_event_cfg = {}
    eval_event_cfg = {}
    # Parse train events
    for event_name, event_info in events_dict.get("train", {}).items():
        event_func_str = event_info.get("func")
        mode = event_info.get("mode", "interval")
        params = _event_params(event_info, run_base=run_base, phase="train")
        interval_step = event_info.get("interval_step", 1)
        event_func = find_function_from_modules(
            event_func_str, event_modules, raise_if_not_found=True
        )
        train_event_cfg[event_name] = EventCfg(
            func=event_func,
            mode=mode,
            params=params,
            interval_step=interval_step,
            is_global=event_info.get("is_global", False),
        )
    # Parse eval events (only if evaluation is enabled)
    if enable_eval and rank == 0:
        for event_name, event_info in events_dict.get("eval", {}).items():
            event_func_str = event_info.get("func")
            mode = event_info.get("mode", "interval")
            params = _event_params(event_info, run_base=run_base, phase="eval")
            interval_step = event_info.get("interval_step", 1)
            event_func = find_function_from_modules(
                event_func_str, event_modules, raise_if_not_found=True
            )
            eval_event_cfg[event_name] = EventCfg(
                func=event_func,
                mode=mode,
                params=params,
                interval_step=interval_step,
                is_global=event_info.get("is_global", False),
            )
    runtime = build_gym_policy_runtime(
        cfg_data,
        device=device,
        num_envs=num_envs,
        headless=headless,
        renderer=renderer,
        gpu_id=gpu_id,
        seed=effective_seed,
        config_dir=Path(config_path).expanduser().resolve().parent,
        profiler=profiler,
        trainer_events=(*train_event_cfg.values(), *eval_event_cfg.values()),
    )
    env = runtime.env
    policy = runtime.policy
    gym_config_path = runtime.gym_config_path
    gym_config_data = runtime.gym_config
    gym_env_cfg = runtime.env_cfg
    if gym_config_path is None or gym_config_data is None or gym_env_cfg is None:
        raise RuntimeError("Simulator Policy runtime is missing task configuration")
    if rank == 0:
        logger.log_info(
            f"Loaded gym_config from {gym_config_path} (env_id={gym_config_data['id']}, num_envs={gym_env_cfg.num_envs}, headless={gym_env_cfg.sim_cfg.headless}, renderer={gym_env_cfg.sim_cfg.render_cfg.renderer}, device={gym_env_cfg.sim_cfg.device})"
        )

    # Create evaluation environment only if enabled
    eval_env = None
    num_eval_envs = trainer_cfg.get("num_eval_envs", 4)
    if enable_eval and rank == 0:
        eval_gym_env_cfg = deepcopy(gym_env_cfg)
        eval_gym_env_cfg.num_envs = int(num_eval_envs)
        eval_gym_env_cfg.seed = eval_seed
        eval_gym_env_cfg.sim_cfg.headless = True
        eval_gym_env_cfg.profiler = None
        eval_env = build_env(gym_config_data["id"], base_env_cfg=eval_gym_env_cfg)
        logger.log_info(
            f"Evaluation environment created (num_envs={num_eval_envs}, headless=True)"
        )

    # Environment construction intentionally uses task/evaluation seeds. Reset
    # the trainer stream so policy initialization is independent of scene work.
    _seed_training_rng(effective_seed, device)

    policy_name = policy_block["name"]

    # Build Algorithm via factory
    algo_name = algo_block["name"].lower()
    algo_cfg = algo_block["cfg"]
    algo = build_algo(
        algo_name,
        algo_cfg,
        policy,
        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
    trainer = Trainer(
        policy=policy,
        env=env,
        algorithm=algo,
        buffer_size=buffer_size,
        batch_size=algo_cfg["batch_size"],
        writer=writer,
        eval_freq=eval_freq if enable_eval else 0,  # Disable eval if not enabled
        save_freq=save_freq,
        save_frequency_updates=int(trainer_cfg.get("save_frequency_updates", 0)),
        checkpoint_dir=checkpoint_dir,
        exp_name=exp_name,
        use_wandb=use_wandb,
        eval_env=eval_env,  # None if enable_eval=False
        event_cfg=train_event_cfg,
        eval_event_cfg=eval_event_cfg if (enable_eval and rank == 0) else {},
        num_eval_episodes=num_eval_episodes,
        distributed=distributed,
        rank=rank,
        world_size=world_size,
        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"),
    )

    if rank == 0:
        logger.log_info("Generic training initialized")
        logger.log_info(f"Task: {type(env).__name__}")
        logger.log_info(
            f"Policy: {policy_name} (available: {get_registered_policy_names()})"
        )
        logger.log_info(
            f"Algorithm: {algo_name} (available: {get_registered_algo_names()})"
        )

    total_steps = int(iterations * buffer_size * env.num_envs * world_size)
    if rank == 0:
        logger.log_info(
            f"Total steps: {total_steps} (iterations≈{iterations}, world_size={world_size})"
        )

    try:
        _configure_checkpoints(
            trainer,
            cfg_data,
            config_path,
            run_base,
            gym_config=gym_config_path,
            gym_data=gym_config_data,
        )
        _resume_training(trainer, resume, run, resume_scheduler)
        trainer.train(total_steps)
        _save_final_checkpoint(trainer)
    except KeyboardInterrupt:
        if rank == 0:
            logger.log_info("Training interrupted by user")
    finally:
        if writer is not None:
            writer.close()
        if use_wandb and rank == 0:
            try:
                wandb.finish()
            except Exception:
                pass

        # Clean up environments to prevent resource leaks
        try:
            if env is not None:
                runtime.close()
        except Exception as e:
            if rank == 0:
                logger.log_warning(f"Failed to close training environment: {e}")

        try:
            if eval_env is not None:
                _close_environment(eval_env)
        except Exception as e:
            if rank == 0:
                logger.log_warning(f"Failed to close evaluation environment: {e}")

        if distributed and torch.distributed.is_initialized():
            torch.distributed.destroy_process_group()

    if rank == 0:
        logger.log_info("Training finished")


def _checkpoint_context(cfg_data: dict, gym_data: dict | None = None) -> dict:
    policy = dict(cfg_data["policy"])
    policy.pop("initial_log_std", None)
    trainer = cfg_data["trainer"]
    result = {"policy": policy, "learning_env": trainer.get("learning_env")}
    if gym_data is not None:
        result["gym"] = gym_data
    result["best_eval_metric"] = trainer.get("best_eval_metric", "eval/avg_reward")
    result["best_eval_mode"] = trainer.get("best_eval_mode", "max")
    return result


def _configure_checkpoints(
    trainer,
    cfg_data: dict,
    config_path: str | Path,
    run_base: str | Path,
    *,
    gym_config: str | Path | None = None,
    gym_data: dict | None = None,
) -> None:
    trainer._checkpoint_context = deepcopy(_checkpoint_context(cfg_data, gym_data))
    trainer._checkpoint_config = deepcopy(cfg_data)
    if getattr(trainer, "rank", 0) != 0:
        return
    trainer._checkpoint_run = Path(run_base)
    config_dir = Path(run_base) / "configs"
    train_path = config_dir / f"train{Path(config_path).suffix.lower() or '.yaml'}"
    gym_path = (
        config_dir / f"gym{Path(gym_config or 'gym.yaml').suffix.lower() or '.yaml'}"
        if gym_data is not None
        else None
    )
    snapshots = {train_path: trainer._checkpoint_config}
    if gym_path is not None:
        snapshots[gym_path] = deepcopy(gym_data)

    def update_manifest(summary: dict) -> None:
        for path, data in snapshots.items():
            contents = (
                json.dumps(data, indent=2)
                if path.suffix == ".json"
                else yaml.safe_dump(data, sort_keys=False)
            ).encode("utf-8")
            _atomic_write(path, lambda stream: stream.write(contents))
        snapshots.clear()
        write_run_manifest(
            run_base,
            train_config=train_path,
            gym_config=gym_path,
            latest_checkpoint=summary["latest_checkpoint_path"],
            best_checkpoint=summary["best_checkpoint_path"],
        )

    trainer._checkpoint_callback = update_manifest


def _save_final_checkpoint(trainer) -> None:
    """Write the final update only if periodic/best saving has not written it."""
    if getattr(trainer, "_last_saved_update", None) != trainer.num_updates:
        trainer.save_checkpoint()
        return
    error = None
    if getattr(trainer, "rank", 0) == 0:
        try:
            _validate_saved_files(trainer)
        except Exception as exc:
            error = exc
    _sync_errors(error, getattr(trainer, "distributed", False))


def _resume_training(trainer, resume, run, scheduler: str) -> None:
    error = None
    path = None
    try:
        path = _resolve_resume(resume, run)
    except Exception as exc:
        error = exc
    _sync_errors(error, getattr(trainer, "distributed", False))
    if path is not None:
        trainer.load_checkpoint(path, resume_scheduler=scheduler)
        if getattr(trainer, "rank", 0) == 0:
            logger.log_info(
                f"Resumed step={trainer.global_step}, updates={trainer.num_updates}, "
                f"learning_rate={trainer.algorithm.current_learning_rate():.6g}, "
                f"scheduler={scheduler}"
            )


def cli(argv: Sequence[str] | None = None) -> None:
    """Command-line interface for RL training.

    Parses CLI arguments and launches training from a config file.

    Task packages are discovered (and init hooks executed) before training so
    that task environments registered in separate packages (e.g.
    ``embodichain_tasks``) are available to ``build_env``. This mirrors the
    ``run_task`` CLI.
    """
    args = parse_args(argv)

    # Discover all installed task packages and run init hooks (register custom
    # manager modules / asset resolvers) before building any environment.
    discover_task_packages()
    execute_init_hooks()

    try:
        train_from_config(
            args.config,
            distributed=args.distributed,
            profile=args.profile,
            profile_output=args.profile_output,
            resume=args.resume,
            run=args.run,
            resume_scheduler=args.resume_scheduler,
        )
    finally:
        from embodichain.lab.sim.sim_manager import SimulationManager

        SimulationManager.flush_cleanup_queue()


if __name__ == "__main__":
    cli()


__all__ = ["cli", "parse_args", "train_from_config"]

The Script Explained

The training script performs the following steps:

1. Parse Configuration: Loads the config file (.json, .yaml, or .yml) and extracts runtime/env/policy/algorithm blocks

2. Setup: Initializes device, seeds, output directories, TensorBoard, and Weights & Biases

3. Build Components: - Environment via build_env() factory - Policy via build_policy() registry - Algorithm via build_algo() factory

4. Create Trainer: Instantiates the Trainer with all components

5. Train: Runs the training loop until completion

Launching Training

To start training, run:

embodichain train-rl --config embodichain_tasks/configs/tasks/classic_control/cart_pole/agents/ppo.yaml

JSON configs are also supported:

embodichain train-rl --config embodichain_tasks/configs/tasks/manipulation/push_cube/agents/ppo.json

Outputs

All outputs are written to ./outputs/<exp_name>_<timestamp>/:

• logs/: TensorBoard logs

• checkpoints/: Model checkpoints

• configs/: Training config and referenced gym config snapshots

• evaluations/: Timestamped policy evaluation reports

• run-manifest.json: Training configs and best/latest checkpoint index used by eval-policy

A training run can be evaluated Headless or opened in its simulator task:

embodichain eval-policy outputs/<exp_name>_<timestamp>
embodichain eval-policy outputs/<exp_name>_<timestamp> --viewer

See Policy Evaluation for EmbodiChain .pt training runs and the external Motion Profile example.

Training Process

The training process follows this sequence:

1. Rollout Phase: SyncCollector interacts with the environment and writes policy-side fields into a shared rollout TensorDict with uniform [N, T + 1] layout. EmbodiedEnv writes environment-side step fields such as reward, done, terminated, and truncated into the same rollout via set_rollout_buffer(). The final slot of transition-only fields is reserved as padding, while obs[:, -1] and value[:, -1] remain valid bootstrap data.

2. Advantage/Return Computation: Algorithm computes advantages and returns from the collected rollout (e.g. GAE for PPO, step-wise group normalization for GRPO) and converts it to a transition-aligned view over the valid first T steps before minibatch optimization.

3. Update Phase: Algorithm updates the policy with update(rollout)

4. Logging: Trainer logs training losses and aggregated metrics to TensorBoard and Weights & Biases

5. Evaluation (periodic): Trainer evaluates the current policy

6. Checkpointing (periodic): Trainer saves model checkpoints

Policy Interface

All policies must inherit from the Policy abstract base class:

from abc import ABC, abstractmethod
import torch.nn as nn

class Policy(nn.Module, ABC):
    device: torch.device

    def get_action(self, tensordict, deterministic: bool = False):
        """Samples action, sample_log_prob, and value into the TensorDict."""
        ...

    @abstractmethod
    def forward(self, tensordict, deterministic: bool = False):
        """Writes action, sample_log_prob, and value into the TensorDict."""
        raise NotImplementedError

    @abstractmethod
    def get_value(self, tensordict):
        """Writes value estimate into the TensorDict."""
        raise NotImplementedError

    @abstractmethod
    def evaluate_actions(self, tensordict):
        """Returns a new TensorDict with log_prob, entropy, and value."""
        raise NotImplementedError

Available Policies

• ActorCritic: MLP-based Gaussian policy with learnable log_std. Requires external actor and critic modules to be provided (defined in the training config file). Used with PPO.

• ActorOnly: Actor-only policy without Critic. Used with GRPO (group-relative advantage estimation).

• VLAPlaceholderPolicy: Placeholder for Vision-Language-Action policies

Algorithms

Available Algorithms

• PPO: Proximal Policy Optimization with GAE

• GRPO: Group Relative Policy Optimization (no Critic, step-wise returns, masked group normalization). Use actor_only policy. Set kl_coef=0 for from-scratch training (CartPole, dense reward); kl_coef=0.02 for VLA/LLM fine-tuning.

Adding a New Algorithm

To add a new algorithm:

1. Create a new algorithm class in embodichain/learning/rl/algo/

2. Implement update(rollout) and consume the shared rollout TensorDict

3. Register in algo/__init__.py:

from tensordict import TensorDict
from embodichain.learning.rl.algo import BaseAlgorithm, register_algo

@register_algo("my_algo")
class MyAlgorithm(BaseAlgorithm):
    def __init__(self, cfg, policy):
        self.cfg = cfg
        self.policy = policy
        self.device = torch.device(cfg.device)

    def update(self, rollout: TensorDict):
        """Update the policy using a collected rollout."""
        # compute advantages / returns from rollout
        # optimize policy parameters
        return {"loss": 0.0}

Adding a New Policy

To add a new policy:

1. Create a new policy class inheriting from the Policy abstract base class

2. Register in models/__init__.py:

from embodichain.learning.rl.policies import register_policy, Policy

@register_policy("my_policy")
class MyPolicy(Policy):
    def __init__(self, obs_dim, action_dim, device, config):
        super().__init__()
        self.device = device
        # Initialize your networks here

    def get_action(self, tensordict, deterministic=False):
        ...
    def forward(self, tensordict, deterministic=False):
        ...
    def get_value(self, tensordict):
        ...
    def evaluate_actions(self, tensordict):
        ...

Current built-in MLP policies use flattened observations in the training path. If your policy requires structured or multi-modal inputs, keep the richer obs_space interface and define a matching rollout/collector schema.

Adding a New Environment

To add a new RL environment:

1. Create an environment class inheriting from EmbodiedEnv (with Action Manager configured for action preprocessing and standardized info structure):

from embodichain.lab.gym.envs import EmbodiedEnv, EmbodiedEnvCfg
from embodichain.lab.gym.utils.registration import register_env
import torch

@register_env("MyTaskRL", override=True)
class MyTaskEnv(EmbodiedEnv):
    def __init__(self, cfg: EmbodiedEnvCfg = None, **kwargs):
        super().__init__(cfg, **kwargs)

    def compute_task_state(self, **kwargs):
        """Compute success/failure conditions and metrics."""
        is_success = ...  # Define success condition
        is_fail = torch.zeros_like(is_success)
        metrics = {"distance": ..., "error": ...}
        return is_success, is_fail, metrics

2. Configure the environment in your config file with actions and extensions:

"env": {
  "id": "MyTaskRL",
  "cfg": {
    "num_envs": 4,
    "actions": {
      "delta_qpos": {
        "func": "RelativeJointPositionAction",
        "params": { "part_name": "arm", "scale": 0.1 }
      }
    },
    "extensions": {
      "success_threshold": 0.05
    }
  }
}

The EmbodiedEnv with Action Manager provides:

• Action Composition: Configure ordered joint, EEF, and gripper terms via actions; each term owns and directly applies its selected resources.

• Standardized Info: Implements get_info() using compute_task_state() template method

Best Practices

• Use EmbodiedEnv with Action Manager for RL Tasks: Inherit from EmbodiedEnv and configure actions in your config. The Action Manager handles action preprocessing (delta_qpos, qpos, qvel, qf, eef_pose) in a modular way.

• Action Configuration: Use the actions field in your config file. For example: "delta_qpos": {"func": "RelativeJointPositionAction", "params": {"part_name": "arm", "scale": 0.1}}.

• Device Management: Device is single-sourced from runtime.cuda. All components (trainer/algorithm/policy/env) share the same device.

• Observation Format: Environments should provide consistent observation shape/types (torch.float32) and a single done = terminated | truncated.

• Algorithm Interface: Algorithms implement update(rollout) and consume a shared rollout TensorDict. Collection is handled by SyncCollector plus environment-side rollout writes in EmbodiedEnv.

• Reward Configuration: Use the RewardManager in your environment config to define reward components. Organize reward components in info["rewards"] dictionary and metrics in info["metrics"] dictionary. The trainer performs dense per-step logging directly from environment info.

• Template Methods: Override compute_task_state() to define success/failure conditions and metrics. Override check_truncated() for custom truncation logic.

• Configuration: Use JSON for all hyperparameters. This makes experiments reproducible and easy to track.

• Logging: Metrics are automatically logged to TensorBoard and Weights & Biases. Check outputs/<exp_name>_<timestamp>/logs/ for TensorBoard logs.

• Checkpoints: Regular checkpoints are saved to outputs/<exp_name>_<timestamp>/checkpoints/. Use these to resume training or evaluate policies.

See Also

• Reinforcement Learning — RL module architecture and component reference

• Embodied Environments — EmbodiedEnv configuration and Action Manager

• Creating a Basic Environment — Creating basic Gymnasium environments

• Creating a Modular Environment — Advanced modular environments with managers

• Supported Tasks — List of available RL task environments

• Policy Evaluation — Headless and Viewer evaluation of EmbodiChain .pt checkpoints and external Motion Profiles