train_ppo_dagger.py

import swanlab
swanlab.sync_wandb()

import functools
import json
import time
from dataclasses import dataclass, field
from pathlib import Path

import jax
import numpy as np
import tyro
from absl import logging

import cat_ppo
from cat_ppo.constant import get_latest_ckpt
from cat_ppo.learning.policy.ppo import train as ppo
from train_ppo import (
    Args,
    _apply_args_to_config,
    _init_wandb,
    _log_checkpoint_path,
    _prepare_exp_name,
    _prepare_training_params,
    _progress,
    _report_training_time,
    _setup_paths,
    _validate_exp_name_format,
)


@dataclass
class DaggerArgs(Args):
    num_timesteps: int = 5_000_000_000
    num_evals: int = 16
    teacher_restore_names: list[str] = field(default_factory=list)
    dagger_timesteps: int = 0
    dagger_actor_loss_scale: float = 1.0
    dagger_value_loss_scale: float = 1.0
    pf_sampling_weights: list[float] = field(default_factory=list)
    pf_sampling_alpha: float = 1.0
    pf_sampling_ema_decay: float = 0.95
    network_kind: str = "mlp"
    policy_hidden_layer_sizes: list[int] = field(default_factory=list)
    value_hidden_layer_sizes: list[int] = field(default_factory=list)
    transformer_embed_dim: int = 256
    transformer_num_heads: int = 4
    transformer_ff_dim: int = 512
    transformer_num_layers: int = 4


def _dagger_task_name(task: str) -> str:
    if task == "G1Cat":
        return "G1CatDagger"
    return task


def _checkpoint_config_path(ckpt_path: Path) -> Path:
    if ckpt_path.name.isdigit():
        return ckpt_path.parent / "config.json"
    if ckpt_path.name == "checkpoints":
        return ckpt_path / "config.json"
    return ckpt_path / "checkpoints" / "config.json"


def _network_factory_key(network_factory: dict) -> str:
    return json.dumps(network_factory, sort_keys=True)


def _prepare_dagger_config(policy_cfg, env_config, args: DaggerArgs):
    dagger_cfg = getattr(policy_cfg, "dagger_config", None)
    if dagger_cfg is None or not getattr(dagger_cfg, "enable", False):
        raise ValueError("train_ppo_dagger requires a task with enabled dagger_config")

    teacher_names = list(args.teacher_restore_names) or list(getattr(dagger_cfg, "teacher_restore_names", []))
    if not teacher_names:
        raise ValueError("Set --teacher_restore_names")
    teacher_checkpoint_paths = []
    for teacher_name in teacher_names:
        ckpt = get_latest_ckpt(teacher_name)
        if ckpt is None:
            raise ValueError(f"No checkpoint found for DAgger teacher run: {teacher_name}")
        teacher_checkpoint_paths.append(str(ckpt))

    scene_paths = []
    teacher_network_factory = None
    for teacher_name, ckpt in zip(teacher_names, teacher_checkpoint_paths):
        config_path = _checkpoint_config_path(Path(ckpt))
        if not config_path.exists():
            raise ValueError(f"Missing teacher config.json for DAgger checkpoint: {ckpt}")
        with config_path.open("r", encoding="utf-8") as f:
            teacher_config = json.load(f)
        current_teacher_network_factory = teacher_config["policy_config"]["network_factory"]
        if teacher_network_factory is None:
            teacher_network_factory = current_teacher_network_factory
        elif _network_factory_key(current_teacher_network_factory) != _network_factory_key(teacher_network_factory):
            raise ValueError(
                "All DAgger teachers must use the same policy_config.network_factory; "
                f"got mismatch at teacher {teacher_name}"
            )
        teacher_pf_config = teacher_config["env_config"]["pf_config"]
        teacher_dx = teacher_pf_config.get("dx", env_config.pf_config.dx)
        if float(teacher_dx) != float(env_config.pf_config.dx):
            raise ValueError(
                f"Teacher scene dx mismatch for {ckpt}: teacher dx={teacher_dx}, "
                f"student dx={env_config.pf_config.dx}"
            )
        scene_paths.append(teacher_pf_config["path"])

    env_config.pf_config.paths = scene_paths
    env_config.pf_config.path = scene_paths[0]
    env_config.pf_config.origin = [-0.5, -1.0, 0.0]
    if args.pf_sampling_weights:
        if len(args.pf_sampling_weights) != len(scene_paths):
            raise ValueError(
                f"--pf_sampling_weights length must match --teacher_restore_names: "
                f"{len(args.pf_sampling_weights)} != {len(scene_paths)}"
            )
        env_config.pf_config.sampling_weights = args.pf_sampling_weights
    else:
        env_config.pf_config.sampling_weights = [1.0] * len(scene_paths)
    env_config.pf_config.sampling_alpha = args.pf_sampling_alpha
    env_config.pf_config.sampling_ema_decay = args.pf_sampling_ema_decay
    dagger_cfg.teacher_restore_names = teacher_names
    dagger_cfg.teacher_checkpoint_paths = teacher_checkpoint_paths
    dagger_cfg.teacher_network_factory = teacher_network_factory
    dagger_cfg.dagger_timesteps = args.dagger_timesteps or (policy_cfg.num_timesteps // 2)
    if args.dagger_actor_loss_scale <= 0:
        raise ValueError("dagger_actor_loss_scale must be > 0 because DAgger phase should train the actor.")
    dagger_cfg.actor_loss_scale = args.dagger_actor_loss_scale
    dagger_cfg.value_loss_scale = args.dagger_value_loss_scale


def _apply_dagger_network_config(policy_cfg, args: DaggerArgs):
    if args.network_kind not in ("mlp", "humanoid_transformer"):
        raise ValueError(
            f"Unsupported network_kind={args.network_kind!r}. Supported values are "
            "'mlp' and 'humanoid_transformer'."
        )
    policy_cfg.network_factory.network_kind = args.network_kind
    if args.policy_hidden_layer_sizes:
        policy_cfg.network_factory.policy_hidden_layer_sizes = tuple(args.policy_hidden_layer_sizes)
    if args.value_hidden_layer_sizes:
        policy_cfg.network_factory.value_hidden_layer_sizes = tuple(args.value_hidden_layer_sizes)
    policy_cfg.network_factory.transformer_embed_dim = args.transformer_embed_dim
    policy_cfg.network_factory.transformer_num_heads = args.transformer_num_heads
    policy_cfg.network_factory.transformer_ff_dim = args.transformer_ff_dim
    policy_cfg.network_factory.transformer_num_layers = args.transformer_num_layers


def train(args: DaggerArgs):
    task_name = _dagger_task_name(args.task)
    env_class = cat_ppo.registry.get(task_name, "train_env_class")
    task_cfg = cat_ppo.registry.get(task_name, "config")
    env_cfg = task_cfg.env_config
    policy_cfg = task_cfg.policy_config

    base_args = {
        key: value
        for key, value in args.__dict__.items()
        if key not in (
            "teacher_restore_names",
            "dagger_timesteps",
            "dagger_actor_loss_scale",
            "dagger_value_loss_scale",
            "pf_sampling_weights",
            "pf_sampling_alpha",
            "pf_sampling_ema_decay",
            "network_kind",
            "policy_hidden_layer_sizes",
            "value_hidden_layer_sizes",
            "transformer_embed_dim",
            "transformer_num_heads",
            "transformer_ff_dim",
            "transformer_num_layers",
        )
    }
    train_args = Args(**{**base_args, "task": task_name})

    exp_name = _prepare_exp_name(task_name, train_args.generate_exp_name())
    debug_mode = "debug" in exp_name

    _validate_exp_name_format(exp_name, debug_mode)

    logdir, ckpt_path = _setup_paths(exp_name)
    _log_checkpoint_path(ckpt_path)

    _apply_args_to_config(train_args, policy_cfg, env_cfg, debug_mode)
    _apply_dagger_network_config(policy_cfg, args)
    _prepare_dagger_config(policy_cfg, env_cfg, args)
    task_cfg.env_config = env_cfg
    policy_params = _prepare_training_params(policy_cfg, ckpt_path)

    if not debug_mode:
        _init_wandb(train_args, exp_name, env_class, task_cfg, ckpt_path)

    train_fn = functools.partial(ppo.train, **policy_params)
    times = [time.monotonic()]

    env = env_class(task_type=env_cfg.task_type, config=env_cfg)
    eval_env = env_class(task_type=env_cfg.task_type, config=env_cfg)

    make_inference_fn, params, _ = train_fn(
        environment=env,
        progress_fn=lambda s, m: _progress(s, m, times, policy_cfg.num_timesteps, debug_mode, exp_name),
        eval_env=eval_env,
        policy_params_fn=lambda *args: None,
    )

    _report_training_time(times)
    inference_fn = jax.jit(make_inference_fn(params, deterministic=True))

    logging.info(f"Run {exp_name} DAgger train done.")

    if train_args.convert_onnx:
        try:
            from cat_ppo.eval.brax2onnx import convert_jax2onnx, get_latest_ckpt as get_latest_saved_ckpt

            ckpt_dir = get_latest_saved_ckpt(ckpt_path)
            obs_size = {
                "privileged_state": (env_cfg.num_pri,),
                "state": (env_cfg.num_obs,),
            }
            act_size = env_cfg.num_act
            policy_obs_key = policy_cfg.network_factory.policy_obs_key
            convert_jax2onnx(
                ckpt_dir=ckpt_dir,
                output_path=f"{ckpt_dir}/policy.onnx",
                inference_fn=inference_fn,
                hidden_layer_sizes=policy_cfg.network_factory.policy_hidden_layer_sizes,
                obs_size=obs_size,
                action_size=act_size,
                policy_obs_key=policy_obs_key,
                jax_params=params,
                network_kind=policy_cfg.network_factory.get("network_kind", "mlp"),
                transformer_embed_dim=policy_cfg.network_factory.get("transformer_embed_dim", 256),
                transformer_num_heads=policy_cfg.network_factory.get("transformer_num_heads", 4),
                transformer_ff_dim=policy_cfg.network_factory.get("transformer_ff_dim", 512),
                transformer_num_layers=policy_cfg.network_factory.get("transformer_num_layers", 4),
                activation="swish",
            )
        except ImportError:
            logging.warning(
                "TensorFlow is not installed. Please install TensorFlow to use ONNX conversion."
            )


if __name__ == "__main__":
    train(tyro.cli(DaggerArgs))