guguji_simulation/guguji_rl/scripts/run_policy.py

from __future__ import annotations

import argparse
import copy
import sys
from pathlib import Path

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

from guguji_rl.config import load_config


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description='Run a trained PPO policy in Gazebo or MuJoCo.')
    parser.add_argument('--config', default='configs/walk_ppo.yaml', help='配置文件路径')
    parser.add_argument('--model', required=True, help='训练好的模型路径，例如 outputs/.../final_model.zip')
    parser.add_argument(
        '--deterministic',
        action='store_true',
        help='是否强制使用确定性动作，适合策略回放和部署调试',
    )
    parser.add_argument(
        '--max-episodes',
        type=int,
        default=0,
        help='最多运行多少个 episode，0 表示一直循环运行',
    )
    parser.add_argument(
        '--render-human',
        action='store_true',
        help='如果当前后端是 MuJoCo，则以交互窗口方式显示策略回放',
    )
    parser.add_argument('--command-forward', type=float, default=None, help='可选固定策略回放时的目标前进速度')
    parser.add_argument('--command-yaw', type=float, default=None, help='可选固定策略回放时的目标偏航角速度')
    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 main() -> int:
    args = parse_args()

    try:
        from stable_baselines3 import PPO
    except ImportError:
        print('缺少 stable-baselines3，请先安装 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.render_human:
        # 训练通常不渲染，但策略回放时可以临时把 MuJoCo 切到 human 模式。
        config = copy.deepcopy(config)
        config.setdefault('mujoco', {})
        config['mujoco']['render_mode'] = 'human'
    if args.command_forward is not None or args.command_yaw is not None:
        config = copy.deepcopy(config)
        config.setdefault('commands', {})
        config['commands']['enabled'] = True
        fixed_command = dict(config['commands'].get('fixed_command', {}))
        fixed_command['enabled'] = True
        if args.command_forward is not None:
            fixed_command['forward_velocity'] = float(args.command_forward)
        if args.command_yaw is not None:
            fixed_command['yaw_rate'] = float(args.command_yaw)
        config['commands']['fixed_command'] = fixed_command

    env = build_env_from_config(config)
    model = PPO.load(resolve_input_path(args.model))

    episode_index = 0
    try:
        while args.max_episodes <= 0 or episode_index < args.max_episodes:
            # 这里直接复用训练环境，方便策略回放和后续在线部署共用同一套观测/动作逻辑。
            observation, _ = env.reset()
            terminated = False
            truncated = False
            total_reward = 0.0

            while not terminated and not truncated:
                action, _ = model.predict(
                    observation,
                    deterministic=args.deterministic or bool(config['evaluation']['deterministic']),
                )
                observation, reward, terminated, truncated, _ = env.step(action)
                total_reward += reward

            print(
                f'episode={episode_index} '
                f'total_reward={total_reward:.3f} '
                f'terminated={terminated} truncated={truncated}'
            )
            episode_index += 1
    except KeyboardInterrupt:
        print('收到 Ctrl+C，停止在线策略运行。')
    finally:
        env.close()

    return 0


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