fix: clamp std before distribution, lower init_noise to 0.5, NaN guard

scripts/train_dreamwaq_rsl.py

@@ -0,0 +1,122 @@
+#!/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()