123 lines
4.8 KiB
Python
123 lines
4.8 KiB
Python
#!/usr/bin/env python3
|
|
"""DreamWaQ training via the faithful rsl_rl-1.0.2 port (upstream-aligned).
|
|
|
|
Uses upstream's ActorCritic_DWAQ + PPO (joint VAE training) + OnPolicyRunner,
|
|
with MotrixLab's DreamWaQ env. Config matches upstream Go1RoughCfgPPO.
|
|
|
|
Usage:
|
|
uv run scripts/train_dreamwaq_rsl.py --num-envs 2048 --iterations 3000
|
|
"""
|
|
import argparse
|
|
import os
|
|
import sys
|
|
import torch
|
|
|
|
import motrix_envs.locomotion.go1.dreamwaq # noqa: F401 (registers env)
|
|
from motrix_envs import registry as env_registry
|
|
from motrix_rl.dwaq_rsl import OnPolicyRunner, DwaqVecEnv
|
|
|
|
|
|
# Upstream Go1RoughCfgPPO (legged_robot_config.py LeggedRobotCfgPPO + Go1 overrides)
|
|
TRAIN_CFG = {
|
|
"runner": {
|
|
"policy_class_name": "ActorCritic_DWAQ",
|
|
"algorithm_class_name": "PPO",
|
|
"num_steps_per_env": 24,
|
|
"save_interval": 50,
|
|
},
|
|
"algorithm": {
|
|
"value_loss_coef": 1.0,
|
|
"use_clipped_value_loss": True,
|
|
"clip_param": 0.2,
|
|
"entropy_coef": 0.01,
|
|
"num_learning_epochs": 5,
|
|
"num_mini_batches": 4,
|
|
"learning_rate": 1.0e-3,
|
|
"schedule": "adaptive",
|
|
"gamma": 0.99,
|
|
"lam": 0.95,
|
|
"desired_kl": 0.01,
|
|
"max_grad_norm": 1.0,
|
|
},
|
|
"policy": {
|
|
"init_noise_std": 1.0,
|
|
},
|
|
}
|
|
|
|
|
|
def main():
|
|
p = argparse.ArgumentParser()
|
|
p.add_argument("--num-envs", type=int, default=2048)
|
|
p.add_argument("--iterations", type=int, default=3000)
|
|
p.add_argument("--seed", type=int, default=1)
|
|
p.add_argument("--init-noise-std", type=float, default=1.0)
|
|
p.add_argument("--force-std", action="store_true",
|
|
help="force reset action std to --init-noise-std even when resuming")
|
|
p.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
|
|
p.add_argument("--resume", default=None,
|
|
help="checkpoint .pt to warm-start from (e.g. the flat policy before "
|
|
"the pyramid-terrain curriculum). Loads model + optimizer; the "
|
|
"action std comes from the checkpoint, not --init-noise-std.")
|
|
p.add_argument("--level", type=int, default=None,
|
|
help="force ALL envs to this terrain level (0-9), skip curriculum")
|
|
args = p.parse_args()
|
|
|
|
torch.manual_seed(args.seed)
|
|
import numpy as np
|
|
np.random.seed(args.seed)
|
|
|
|
TRAIN_CFG["policy"]["init_noise_std"] = args.init_noise_std
|
|
|
|
raw_env = env_registry.make("go1-dreamwaq-walk", num_envs=args.num_envs)
|
|
if args.level is not None:
|
|
raw_env._force_level = args.level
|
|
print(f"[DreamWaQ-rsl] forcing ALL envs at terrain level {args.level}")
|
|
env = DwaqVecEnv(raw_env, device=args.device)
|
|
print(f"[DreamWaQ-rsl] {args.num_envs} envs | obs={env.num_obs} priv={env.num_privileged_obs} "
|
|
f"hist={env.num_obs_hist} act={env.num_actions} | device={args.device}")
|
|
print(f"[DreamWaQ-rsl] actor_in={env.num_obs + 19} critic_in={env.num_privileged_obs} "
|
|
f"cenet_in={env.num_obs_hist * env.num_obs}")
|
|
|
|
log_dir = os.path.join("runs", "go1-dreamwaq-walk", "rsl_dwaq",
|
|
__import__("datetime").datetime.now().strftime("%m-%d_%H-%M-%S"))
|
|
os.makedirs(log_dir, exist_ok=True)
|
|
|
|
# Save config snapshot for later reference
|
|
import json
|
|
cfg_snapshot = {
|
|
"command_line": sys.argv,
|
|
"kp": raw_env.cfg.control_config.stiffness,
|
|
"kd": raw_env.cfg.control_config.damping,
|
|
"action_scale": raw_env.cfg.control_config.action_scale,
|
|
"rewards": dict(raw_env.cfg.reward_config.scales),
|
|
"sigma": raw_env.cfg.reward_config.tracking_sigma,
|
|
"only_positive": raw_env.cfg.reward_config.only_positive_rewards,
|
|
"force_level": args.level,
|
|
"init_noise_std": TRAIN_CFG["policy"]["init_noise_std"],
|
|
"entropy_coef": TRAIN_CFG["algorithm"]["entropy_coef"],
|
|
"terrain_rows": raw_env.cfg.terrain_rows,
|
|
"terrain_cols": raw_env.cfg.terrain_cols,
|
|
"scene": raw_env.cfg.model_file,
|
|
}
|
|
with open(os.path.join(log_dir, "config.json"), "w") as f:
|
|
json.dump(cfg_snapshot, f, indent=2, default=str)
|
|
|
|
runner = OnPolicyRunner(env, TRAIN_CFG, log_dir=log_dir, device=args.device)
|
|
if args.resume:
|
|
runner.load(args.resume)
|
|
runner.current_learning_iteration = 0
|
|
if args.force_std:
|
|
runner.alg.actor_critic.std.data.fill_(args.init_noise_std)
|
|
print(f"[DreamWaQ-rsl] warm-start from {args.resume} "
|
|
f"(std FORCED to {args.init_noise_std})")
|
|
else:
|
|
print(f"[DreamWaQ-rsl] warm-start from {args.resume} "
|
|
f"(model+optimizer; std from checkpoint)")
|
|
print(f"[DreamWaQ-rsl] log_dir={log_dir} | training {args.iterations} iterations...")
|
|
runner.learn(args.iterations, init_at_random_ep_len=True)
|
|
print("[DreamWaQ-rsl] done.")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|