guguji_simulation/guguji_rl/scripts/train.py

from __future__ import annotations

import argparse
import copy
import sys
from datetime import datetime
from pathlib import Path
from typing import Any

SCRIPT_ROOT = Path(__file__).resolve().parents[1]
if str(SCRIPT_ROOT) not in sys.path:
    sys.path.insert(0, str(SCRIPT_ROOT))

import torch

from guguji_rl.config import load_config, resolve_project_path, save_yaml
from guguji_rl.evaluation import evaluate_forward_progress, print_forward_progress_summary


def resolve_device(device_name: str) -> str:
    if device_name == 'auto':
        return 'cuda' if torch.cuda.is_available() else 'cpu'
    if device_name == 'cuda' and not torch.cuda.is_available():
        raise RuntimeError('配置要求使用 CUDA，但当前 torch 检测不到可用 GPU。')
    return device_name


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description='Train PPO policy for guguji biped robot.')
    parser.add_argument(
        '--config',
        default='configs/balance_ppo.yaml',
        help='训练配置文件路径，默认使用 balance_ppo.yaml',
    )
    parser.add_argument(
        '--device',
        default=None,
        help='可选覆盖配置文件中的设备设置，例如 cpu / cuda / auto',
    )
    parser.add_argument(
        '--total-timesteps',
        type=int,
        default=None,
        help='可选覆盖配置文件中的 total_timesteps',
    )
    parser.add_argument(
        '--init-model',
        default=None,
        help='可选指定一个已有 PPO 模型，用于继续训练或做课程学习初始化',
    )
    parser.add_argument(
        '--skip-auto-eval',
        action='store_true',
        help='训练完成后跳过自动前进评估',
    )
    parser.add_argument(
        '--render-human',
        action='store_true',
        help='如果当前后端是 MuJoCo，则在训练时同步打开 GUI 画面',
    )
    return parser.parse_args()


def resolve_input_path(path_str: str) -> Path:
    path = Path(path_str)
    if path.is_absolute() or path.exists():
        return path
    return SCRIPT_ROOT / path


def maybe_override_policy_log_std(model: object, initial_log_std: float | None) -> None:
    """可选地缩小 PPO 的初始探索方差，适合课程学习后的精修阶段。"""
    if initial_log_std is None:
        return
    policy = getattr(model, 'policy', None)
    if policy is None or not hasattr(policy, 'log_std'):
        raise RuntimeError('当前策略对象不支持直接设置 log_std。')

    # 这里直接把每个动作维度的对数标准差统一改成同一个值，
    # 方便在“已有步态基础上继续训练”时降低探索噪声，减少无意义的乱踢。
    policy.log_std.data.fill_(float(initial_log_std))
    print(f'已将策略初始 log_std 设为: {float(initial_log_std):.3f}')


def sanitize_stage_name(stage_name: str) -> str:
    sanitized = ''.join(
        character if character.isalnum() or character in {'-', '_'} else '_'
        for character in stage_name.strip()
    )
    return sanitized.strip('_') or 'stage'


def deep_merge_stage_override(base: dict[str, Any], override: dict[str, Any]) -> None:
    for key, value in override.items():
        if isinstance(value, dict) and isinstance(base.get(key), dict):
            deep_merge_stage_override(base[key], value)
        else:
            base[key] = copy.deepcopy(value)


def format_stage_target_summary(stage_config: dict[str, Any]) -> str:
    commands_config = stage_config.get('commands', {})
    if bool(commands_config.get('enabled', False)):
        forward_velocity_range = commands_config.get('forward_velocity_range', [0.0, 0.0])
        yaw_rate_range = commands_config.get('yaw_rate_range', [0.0, 0.0])
        return (
            'command_conditioned '
            f'vx=[{float(forward_velocity_range[0]):.2f}, {float(forward_velocity_range[1]):.2f}] '
            f'yaw=[{float(yaw_rate_range[0]):.2f}, {float(yaw_rate_range[1]):.2f}]'
        )
    return f'target_forward_velocity={float(stage_config["task"]["target_forward_velocity"]):.2f}'


def warm_start_policy_from_checkpoint(
    *,
    ppo_class: type,
    model: object,
    init_model_path: Path,
    device: str,
) -> None:
    try:
        model.set_parameters(
            str(init_model_path),
            exact_match=False,
            device=device,
        )
        print(f'已加载课程初始化模型: {init_model_path}')
        return
    except Exception as error:
        print(f'完整参数加载失败，将改用兼容 warm start: {error}')

    source_model = ppo_class.load(str(init_model_path), device=device)
    source_state_dict = source_model.policy.state_dict()
    target_state_dict = model.policy.state_dict()
    matched_keys: list[str] = []
    skipped_keys: list[str] = []

    for key, source_value in source_state_dict.items():
        target_value = target_state_dict.get(key)
        if target_value is None or tuple(target_value.shape) != tuple(source_value.shape):
            skipped_keys.append(key)
            continue
        target_state_dict[key] = source_value.detach().clone()
        matched_keys.append(key)

    model.policy.load_state_dict(target_state_dict, strict=False)
    print(
        '已完成兼容 warm start: '
        f'匹配 {len(matched_keys)} 个张量，'
        f'跳过 {len(skipped_keys)} 个形状不兼容张量。'
    )
    if skipped_keys:
        print(f'跳过的典型张量: {", ".join(skipped_keys[:4])}')


def build_curriculum_stage_configs(config: dict[str, Any]) -> list[tuple[str | None, dict[str, Any]]]:
    """把课程学习阶段展开成一组可直接训练的独立配置。"""
    raw_stages = config['training'].get('curriculum_stages') or []
    if not raw_stages:
        single_stage_config = copy.deepcopy(config)
        single_stage_config['training'].pop('curriculum_stages', None)
        return [(None, single_stage_config)]

    stage_configs: list[tuple[str | None, dict[str, Any]]] = []
    for stage_index, raw_stage in enumerate(raw_stages, start=1):
        if not isinstance(raw_stage, dict):
            raise RuntimeError('training.curriculum_stages 里的每个阶段都必须是字典。')

        stage_config = copy.deepcopy(config)
        stage_config['training'].pop('curriculum_stages', None)

        for section_name in ('task', 'commands', 'rewards', 'robot', 'sim', 'mujoco', 'evaluation', 'training'):
            section_override = raw_stage.get(section_name)
            if isinstance(section_override, dict):
                deep_merge_stage_override(stage_config[section_name], section_override)

        raw_name = str(raw_stage.get('name') or f'stage_{stage_index}')
        stage_name = f'{stage_index:02d}_{sanitize_stage_name(raw_name)}'

        # 课程阶段目前主要控制“目标前进速度 + 本阶段训练步数 + 探索方差”。
        # 这样 walking 阶段就能从慢到快逐段抬升，而不用一次把目标速度顶太高。
        if 'target_forward_velocity' in raw_stage:
            stage_config['task']['target_forward_velocity'] = float(raw_stage['target_forward_velocity'])
        if 'target_yaw_rate' in raw_stage:
            stage_config['task']['target_yaw_rate'] = float(raw_stage['target_yaw_rate'])
        if 'total_timesteps' in raw_stage:
            stage_config['training']['total_timesteps'] = int(raw_stage['total_timesteps'])
        if 'initial_log_std' in raw_stage:
            stage_config['training']['initial_log_std'] = float(raw_stage['initial_log_std'])

        stage_config['experiment']['name'] = f"{config['experiment']['name']}_{stage_name}"
        stage_configs.append((stage_name, stage_config))

    return stage_configs


def main() -> int:
    args = parse_args()

    try:
        from stable_baselines3 import PPO
        from stable_baselines3.common.callbacks import CheckpointCallback
        from stable_baselines3.common.monitor import Monitor
    except ImportError:
        print(
            '缺少 stable-baselines3，请先进入 guguji_rl 目录安装依赖: '
            'pip install -r requirements.txt',
            file=sys.stderr,
        )
        return 1

    from guguji_rl.envs import build_env_from_config

    config = load_config(resolve_input_path(args.config))
    if args.device is not None:
        config['training']['device'] = args.device
    if args.total_timesteps is not None:
        config['training']['total_timesteps'] = args.total_timesteps
    if args.init_model is not None:
        config['training']['init_model_path'] = str(resolve_input_path(args.init_model))
    if args.render_human:
        # MuJoCo 训练默认关闭渲染以保证速度。
        # 当你想一边训练一边看画面时，可以通过命令行临时打开 human 渲染。
        config = copy.deepcopy(config)
        config.setdefault('mujoco', {})
        config['mujoco']['render_mode'] = 'human'

    # 这里统一解析训练设备，方便你只改 YAML 就切换 CPU / GPU。
    config['training']['device'] = resolve_device(config['training']['device'])

    output_root = resolve_project_path(config, config['training']['output_root'])
    timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
    run_dir = output_root / f"{config['experiment']['name']}_{timestamp}"
    run_dir.mkdir(parents=True, exist_ok=True)
    # 保存一份展开后的配置，便于后面复现实验。
    save_yaml(config, run_dir / 'resolved_config.yaml')
    stage_configs = build_curriculum_stage_configs(config)

    print(f"训练设备: {config['training']['device']}")
    print(f"输出目录: {run_dir}")
    model = None
    final_model_path = run_dir / 'final_model'
    final_stage_config = config

    for stage_index, (stage_name, stage_config) in enumerate(stage_configs, start=1):
        stage_dir = run_dir if stage_name is None else run_dir / stage_name
        stage_dir.mkdir(parents=True, exist_ok=True)
        # 每个阶段都单独保存一份实际生效的配置，后面你回看实验会很方便。
        save_yaml(stage_config, stage_dir / 'resolved_config.yaml')

        if stage_name is not None:
            print(
                f'开始课程阶段 {stage_index}/{len(stage_configs)}: {stage_name} '
                f'({format_stage_target_summary(stage_config)}, '
                f'timesteps={int(stage_config["training"]["total_timesteps"])})'
            )

        env = Monitor(build_env_from_config(stage_config))
        checkpoint_callback = CheckpointCallback(
            save_freq=max(int(stage_config['training']['checkpoint_freq']), 1),
            save_path=str(stage_dir / 'checkpoints'),
            name_prefix='guguji_ppo',
        )

        try:
            if model is None:
                policy_kwargs = {
                    'net_arch': list(stage_config['training']['policy_net_arch']),
                }
                # 先用 MLP + PPO 跑通训练闭环，后面你可以再逐步增大网络规模。
                model = PPO(
                    policy='MlpPolicy',
                    env=env,
                    verbose=1,
                    seed=int(stage_config['training']['seed']),
                    learning_rate=float(stage_config['training']['learning_rate']),
                    n_steps=int(stage_config['training']['n_steps']),
                    batch_size=int(stage_config['training']['batch_size']),
                    gamma=float(stage_config['training']['gamma']),
                    gae_lambda=float(stage_config['training']['gae_lambda']),
                    clip_range=float(stage_config['training']['clip_range']),
                    ent_coef=float(stage_config['training']['ent_coef']),
                    vf_coef=float(stage_config['training']['vf_coef']),
                    device=stage_config['training']['device'],
                    tensorboard_log=str(run_dir / 'tensorboard'),
                    policy_kwargs=policy_kwargs,
                )

                init_model_path = stage_config['training'].get('init_model_path')
                if init_model_path:
                    resolved_init_model_path = resolve_input_path(str(init_model_path))
                    # 如果 observation 维度还没变，这里会完整复用旧权重。
                    # 如果我们给新任务增加了命令维度，这里会自动退化成“兼容 warm start”，
                    # 尽量保留平衡模型里已经学到的站立/稳定控制能力。
                    warm_start_policy_from_checkpoint(
                        ppo_class=PPO,
                        model=model,
                        init_model_path=resolved_init_model_path,
                        device=stage_config['training']['device'],
                    )
            else:
                model.set_env(env)

            maybe_override_policy_log_std(model, stage_config['training'].get('initial_log_std'))
            model.learn(
                total_timesteps=int(stage_config['training']['total_timesteps']),
                callback=checkpoint_callback,
                progress_bar=True,
                reset_num_timesteps=(stage_index == 1),
            )

            stage_model_path = stage_dir / 'final_model'
            model.save(stage_model_path)
            final_model_path = stage_model_path
            final_stage_config = stage_config

            if stage_name is not None:
                print(f'课程阶段完成，模型已保存到: {stage_model_path.with_suffix(".zip")}')
        finally:
            env.close()

    if final_model_path != run_dir / 'final_model' and model is not None:
        # 在课程学习模式下，额外在 run 根目录保存一份最终模型，方便统一引用。
        model.save(run_dir / 'final_model')
        final_model_path = run_dir / 'final_model'

    print(f'训练完成，模型已保存到: {run_dir / "final_model.zip"}')

    evaluation_config = final_stage_config['evaluation']
    if bool(evaluation_config.get('auto_forward_progress', True)) and not args.skip_auto_eval:
        try:
            # 每轮训练结束后自动做一次前进评估，方便你快速看 delta_x / mean_vx。
            summary = evaluate_forward_progress(
                config=final_stage_config,
                model_path=final_model_path,
                episodes=int(evaluation_config['forward_progress_episodes']),
                max_steps=int(evaluation_config['forward_progress_max_steps']),
                deterministic=bool(evaluation_config['forward_progress_deterministic']),
            )
            print_forward_progress_summary(summary)
        except Exception as error:
            print(f'自动前进评估失败: {error}', file=sys.stderr)

    return 0


if __name__ == '__main__':
    raise SystemExit(main())