Upload 11 files

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+replay.mp4 filter=lfs diff=lfs merge=lfs -text

amp_callback.py

@@ -0,0 +1,191 @@
+import os
+from collections import deque
+
+import numpy as np
+import torch
+from stable_baselines3.common.callbacks import BaseCallback
+
+
+class AMPDiscriminatorCallback(BaseCallback):
+    """Trains the AMP discriminator and syncs weights into the envs.
+
+    FIXES vs original:
+    - Least-squares (MSE) loss to +1/-1 targets, per the AMP paper
+      (was smooth_l1, which goes linear and weakens gradients exactly when
+      the discriminator is wrong).
+    - Robust trigger (`num_timesteps - last >= train_freq` instead of a
+      modulo that silently never fires if train_freq isn't a multiple of
+      n_envs).
+    - Pushes updated weights to all envs via env_method("load_disc_state")
+      -> required for SubprocVecEnv workers.
+    - Sensible default budget: updates_per_call=8 x batch 512 per rollout
+      (was 1 x 256 per 16384 steps at lr 3e-5 - the discriminator barely
+      moved over an entire run).
+    """
+
+    def __init__(
+        self,
+        motion_lib,
+        discriminator,
+        optimizer,
+        batch_size=512,
+        updates_per_call=8,
+        train_freq=16384,
+        save_freq=500_000,
+        save_path="./models/checkpoints",
+        device="cpu",
+        amp_mean=None,
+        amp_std=None,
+        fake_replay_size=100_000,
+        gradient_penalty_weight=5.0,
+        score_reg_weight=1e-4,
+        max_grad_norm=1.0,
+        verbose=1,
+    ):
+        super().__init__(verbose)
+        self.motion_lib = motion_lib
+        self.disc = discriminator
+        self.optimizer = optimizer
+        self.batch_size = batch_size
+        self.updates_per_call = updates_per_call
+        self.train_freq = train_freq
+        self.save_freq = save_freq
+        self.save_path = save_path
+        self.device = device
+        self.amp_mean = amp_mean
+        self.amp_std = amp_std
+        self.fake_replay = deque(maxlen=fake_replay_size)
+        self.gradient_penalty_weight = gradient_penalty_weight
+        self.score_reg_weight = score_reg_weight
+        self.max_grad_norm = max_grad_norm
+        self.last_train_step = 0
+        self.last_save_step = 0
+        os.makedirs(self.save_path, exist_ok=True)
+
+    # ------------------------------------------------------------------ #
+
+    def _on_training_start(self):
+        # Make sure every worker starts from the same weights as the master.
+        self.push_disc_weights()
+
+    def _on_step(self):
+        if self.num_timesteps - self.last_train_step >= self.train_freq:
+            self.train_discriminator()
+            self.last_train_step = self.num_timesteps
+
+        if self.num_timesteps - self.last_save_step >= self.save_freq:
+            self.save_discriminator()
+            self.last_save_step = self.num_timesteps
+        return True
+
+    def normalize_amp(self, x_np):
+        if self.amp_mean is not None and self.amp_std is not None:
+            return (x_np - self.amp_mean) / self.amp_std
+        return x_np
+
+    def push_disc_weights(self):
+        state = {k: v.detach().cpu() for k, v in self.disc.state_dict().items()}
+        self.training_env.env_method("load_disc_state", state)
+
+    # ------------------------------------------------------------------ #
+
+    def train_discriminator(self):
+        fake_lists = self.training_env.env_method("pop_fake_transitions")
+        fake_transitions = [t for lst in fake_lists for t in lst]
+        if fake_transitions:
+            self.fake_replay.extend(fake_transitions)
+
+        if len(self.fake_replay) < self.batch_size:
+            return
+
+        self.disc.train()
+        last_loss = None
+
+        for _ in range(self.updates_per_call):
+            idx = np.random.randint(0, len(self.fake_replay), size=self.batch_size)
+            fake_np = self.normalize_amp(
+                np.array([self.fake_replay[i] for i in idx], dtype=np.float32)
+            )
+            real_np = self.normalize_amp(
+                np.array(
+                    [self.motion_lib.sample_amp_transition()
+                     for _ in range(self.batch_size)],
+                    dtype=np.float32,
+                )
+            )
+
+            real_x = torch.as_tensor(real_np, device=self.device)
+            fake_x = torch.as_tensor(fake_np, device=self.device)
+
+            real_scores = self.disc(real_x)
+            fake_scores = self.disc(fake_x)
+
+            # Least-squares GAN loss to targets +1 / -1 (AMP paper).
+            real_loss = torch.square(real_scores - 1.0).mean()
+            fake_loss = torch.square(fake_scores + 1.0).mean()
+            gp_loss = self.gradient_penalty(real_x)
+            score_reg = self.score_reg_weight * (
+                real_scores.pow(2).mean() + fake_scores.pow(2).mean()
+            )
+
+            loss = (
+                0.5 * (real_loss + fake_loss)
+                + self.gradient_penalty_weight * gp_loss
+                + score_reg
+            )
+
+            self.optimizer.zero_grad()
+            loss.backward()
+            torch.nn.utils.clip_grad_norm_(self.disc.parameters(), self.max_grad_norm)
+            self.optimizer.step()
+            last_loss = loss.item()
+
+        self.disc.eval()
+        self.push_disc_weights()
+
+        # ---- diagnostics on the last minibatch ---- #
+        with torch.no_grad():
+            real_scores = self.disc(real_x)
+            fake_scores = self.disc(fake_x)
+            real_score = real_scores.mean().item()
+            fake_score = fake_scores.mean().item()
+            real_reward = self.disc.amp_reward(real_x).mean().item()
+            fake_rewards = self.disc.amp_reward(fake_x)
+            fake_reward = fake_rewards.mean().item()
+            disc_acc = 0.5 * (
+                (real_scores > 0).float().mean()
+                + (fake_scores < 0).float().mean()
+            ).item()
+
+        if self.verbose:
+            print(
+                f"AMP Disc | loss={last_loss:.4f} acc={disc_acc:.2f} "
+                f"real_score={real_score:.3f} fake_score={fake_score:.3f} "
+                f"real_reward={real_reward:.3f} fake_reward={fake_reward:.3f}"
+            )
+
+        self.logger.record("amp/loss", last_loss)
+        self.logger.record("amp/disc_acc", disc_acc)
+        self.logger.record("amp/real_score", real_score)
+        self.logger.record("amp/fake_score", fake_score)
+        self.logger.record("amp/real_reward", real_reward)
+        self.logger.record("amp/fake_reward", fake_reward)
+
+    def gradient_penalty(self, real):
+        real = real.clone().detach().requires_grad_(True)
+        scores = self.disc(real)
+        gradients = torch.autograd.grad(
+            outputs=scores.sum(),
+            inputs=real,
+            create_graph=True,
+            retain_graph=True,
+            only_inputs=True,
+        )[0]
+        return gradients.pow(2).sum(dim=1).mean()
+
+    def save_discriminator(self):
+        path = os.path.join(
+            self.save_path, f"amp_discriminator_{self.num_timesteps}.pt"
+        )
+        torch.save(self.disc.state_dict(), path)
+        print("Saved discriminator checkpoint:", path)

amp_discriminator.py

@@ -0,0 +1,48 @@
+import torch
+import torch.nn as nn
+
+
+class AMPDiscriminator(nn.Module):
+    """Least-squares AMP discriminator (predicts +1 real / -1 fake).
+
+    FIX vs original: `amp_reward` previously used exp(-0.05*(d-1)^2), which is
+    almost flat: a confidently-fake score of -1 still earned reward 0.82 vs 1.0
+    for perfectly-real. The style signal was ~constant, so PPO had nothing to
+    optimize. This restores the paper reward
+
+        r = clamp(1 - 0.25 * (d - 1)^2, 0, 1)
+
+    which gives r=1 at d=+1 (real) and r=0 at d<=-1 (fake) - a full-range,
+    informative reward.
+    """
+
+    def __init__(self, input_dim=90, hidden_dim=512):
+        super().__init__()
+        self.input_dim = input_dim
+
+        self.net = nn.Sequential(
+            nn.Linear(input_dim, hidden_dim),
+            nn.ReLU(),
+            nn.Linear(hidden_dim, hidden_dim),
+            nn.ReLU(),
+            nn.Linear(hidden_dim, 1),
+        )
+
+    def forward(self, x):
+        if x.ndim == 1:
+            x = x.unsqueeze(0)
+
+        if x.shape[-1] != self.input_dim:
+            raise ValueError(
+                f"Expected AMP input dim {self.input_dim}, got {x.shape[-1]}"
+            )
+
+        return self.net(x).view(-1, 1)
+
+    def predict_score(self, x):
+        return self.forward(x)
+
+    def amp_reward(self, x):
+        d = self.forward(x)
+        r = 1.0 - 0.25 * torch.square(d - 1.0)
+        return torch.clamp(r, min=0.0, max=1.0)

amp_env.py

@@ -0,0 +1,130 @@
+import numpy as np
+import torch
+import gymnasium as gym
+
+import register_env  # noqa: F401  (registers HumanoidDirection-v0)
+from amp_obs import build_amp_obs
+
+
+class AMPHumanoidEnv(gym.Wrapper):
+    """Wraps HumanoidDirection-v0 with an AMP style reward.
+
+    FIXES vs original:
+    - AMP features come from the shared, heading-invariant `build_amp_obs`
+      (was: raw world-frame qpos[2:]+qvel, letting the discriminator cheat on
+      global heading).
+    - Reward mixing is paper-style: total = w_task * r_task + w_style * r_style
+      with both rewards in [0, 1] and default weights 0.5 / 0.5
+      (was: r_task in ~[0, 7] plus 0.2 * near-constant style reward).
+    - `load_disc_state` lets the training callback push fresh discriminator
+      weights into each env. This is REQUIRED under SubprocVecEnv, where every
+      worker process holds its own copy of the discriminator; without syncing,
+      the style reward would be computed with the frozen initial weights
+      forever.
+    """
+
+    def __init__(
+        self,
+        discriminator,
+        motion_lib=None,
+        env_id="HumanoidDirection-v0",
+        render_mode=None,
+        task_weight=0.5,
+        amp_weight=0.5,
+        device="cpu",
+        amp_mean=None,
+        amp_std=None,
+        reference_state_init_prob=0.5,
+    ):
+        env = gym.make(env_id, render_mode=render_mode)
+        super().__init__(env)
+
+        self.disc = discriminator
+        self.disc.eval()
+        self.motion_lib = motion_lib
+        self.task_weight = task_weight
+        self.amp_weight = amp_weight
+        self.device = device
+        self.prev_amp_obs = None
+        self.fake_amp_transitions = []
+        self.amp_mean = amp_mean
+        self.amp_std = amp_std
+        self.reference_state_init_prob = reference_state_init_prob
+
+    # ------------------------------------------------------------------ #
+
+    def get_amp_obs(self):
+        data = self.env.unwrapped.data
+        return build_amp_obs(data.qpos.copy(), data.qvel.copy())
+
+    def load_disc_state(self, state_dict):
+        """Called via VecEnv.env_method by the AMP callback after each
+        discriminator update."""
+        self.disc.load_state_dict(state_dict)
+        self.disc.eval()
+        return True
+
+    def _current_policy_obs(self):
+        unwrapped = self.env.unwrapped
+        humanoid_obs = unwrapped._get_obs()
+        if hasattr(unwrapped, "_get_obs_with_direction"):
+            return unwrapped._get_obs_with_direction(humanoid_obs)
+        return humanoid_obs
+
+    def maybe_reference_state_init(self):
+        if self.motion_lib is None or self.reference_state_init_prob <= 0.0:
+            return False
+
+        rng = self.env.unwrapped.np_random
+        if rng.random() >= self.reference_state_init_prob:
+            return False
+
+        qpos, qvel = self.motion_lib.sample_reference_state()
+        self.env.unwrapped.set_state(qpos, qvel)
+        return True
+
+    def reset(self, **kwargs):
+        obs, info = self.env.reset(**kwargs)
+
+        used_reference_state = self.maybe_reference_state_init()
+        if used_reference_state:
+            obs = self._current_policy_obs()
+
+        self.prev_amp_obs = self.get_amp_obs()
+        info["reference_state_init"] = used_reference_state
+        return obs, info
+
+    def step(self, action):
+        obs, task_reward, terminated, truncated, info = self.env.step(action)
+
+        current_amp_obs = self.get_amp_obs()
+        amp_transition = np.concatenate(
+            [self.prev_amp_obs, current_amp_obs]
+        ).astype(np.float32)
+        self.fake_amp_transitions.append(amp_transition)
+
+        x_np = amp_transition
+        if self.amp_mean is not None and self.amp_std is not None:
+            x_np = (x_np - self.amp_mean) / self.amp_std
+
+        with torch.no_grad():
+            x = torch.tensor(
+                x_np, dtype=torch.float32, device=self.device
+            ).unsqueeze(0)
+            amp_reward = self.disc.amp_reward(x).item()
+
+        total_reward = (
+            self.task_weight * task_reward + self.amp_weight * amp_reward
+        )
+
+        info["task_reward"] = task_reward
+        info["amp_reward"] = amp_reward
+        info["total_reward"] = total_reward
+
+        self.prev_amp_obs = current_amp_obs
+        return obs, total_reward, terminated, truncated, info
+
+    def pop_fake_transitions(self):
+        transitions = self.fake_amp_transitions
+        self.fake_amp_transitions = []
+        return transitions

amp_obs.py

@@ -0,0 +1,85 @@
+"""Shared, heading-invariant AMP feature construction.
+
+WHY THIS FILE EXISTS (the main bug in the original code):
+The discriminator was fed `concat[qpos[2:], qvel]`, which contains the root
+quaternion in the WORLD frame and the root linear velocity in the WORLD frame.
+Your task randomizes the walking direction every episode, but the mocap clips
+walk in one fixed world direction. So the discriminator learns to separate
+real/fake by *global heading*, not by gait quality -> the style reward
+actively punishes the policy whenever it follows a target direction that
+differs from the dataset's heading. Task reward and style reward fight each
+other and neither wins.
+
+Fix: express all root quantities in a local "heading frame" (yaw removed),
+exactly like the AMP paper. Both the policy transitions (amp_env.py) and the
+mocap transitions (motion_lib.py) MUST use this same function.
+
+Feature layout (45 dims -> transition pair is 90 = discriminator input_dim):
+    root height                                   1
+    root orientation, yaw removed (quat, w>=0)    4
+    root linear velocity in heading frame         3
+    root angular velocity (MuJoCo body frame)     3
+    joint angles                                 17
+    joint velocities                             17
+"""
+
+import numpy as np
+
+AMP_OBS_DIM = 45
+AMP_TRANSITION_DIM = 2 * AMP_OBS_DIM
+
+
+def quat_mul(a, b):
+    """Hamilton product, MuJoCo [w, x, y, z] convention."""
+    w1, x1, y1, z1 = a
+    w2, x2, y2, z2 = b
+    return np.array([
+        w1 * w2 - x1 * x2 - y1 * y2 - z1 * z2,
+        w1 * x2 + x1 * w2 + y1 * z2 - z1 * y2,
+        w1 * y2 - x1 * z2 + y1 * w2 + z1 * x2,
+        w1 * z2 + x1 * y2 - y1 * x2 + z1 * w2,
+    ])
+
+
+def yaw_from_quat(q):
+    """ZYX-convention yaw (heading) of a [w, x, y, z] quaternion."""
+    w, x, y, z = q
+    return np.arctan2(2.0 * (w * z + x * y), 1.0 - 2.0 * (y * y + z * z))
+
+
+def build_amp_obs(qpos, qvel):
+    """Humanoid-v5 layout:
+        qpos = [x, y, z, qw, qx, qy, qz, 17 joint angles]        (24,)
+        qvel = [vx, vy, vz (world), wx, wy, wz (body), 17 vels]  (23,)
+    Returns float32 (45,) heading- and position-invariant features.
+    """
+    quat = np.asarray(qpos[3:7], dtype=np.float64)
+    yaw = yaw_from_quat(quat)
+
+    c, s = np.cos(yaw), np.sin(yaw)
+    # Rotates a world-frame vector into the heading frame (inverse yaw).
+    R_inv = np.array([[c, s, 0.0],
+                      [-s, c, 0.0],
+                      [0.0, 0.0, 1.0]])
+
+    # Remove yaw from the root orientation: q_local = q_yaw^-1 (x) q
+    half = -0.5 * yaw
+    q_yaw_inv = np.array([np.cos(half), 0.0, 0.0, np.sin(half)])
+    quat_local = quat_mul(q_yaw_inv, quat)
+    if quat_local[0] < 0.0:
+        quat_local = -quat_local  # canonical sign (quats double-cover)
+
+    lin_vel_local = R_inv @ np.asarray(qvel[0:3], dtype=np.float64)
+    # MuJoCo free-joint angular velocity (qvel[3:6]) is already expressed in
+    # the body-local frame, i.e. heading-invariant. Left untouched. Either
+    # way, expert and policy use the identical convention/transform.
+    ang_vel = qvel[3:6]
+
+    return np.concatenate([
+        qpos[2:3],       # root height
+        quat_local,      # roll/pitch information only
+        lin_vel_local,   # forward/lateral/vertical speed relative to facing
+        ang_vel,
+        qpos[7:],        # joint angles (already local)
+        qvel[6:],        # joint velocities (already local)
+    ]).astype(np.float32)

humanoid_direction_env.py

@@ -0,0 +1,78 @@
+import numpy as np
+from gymnasium.envs.mujoco.humanoid_v5 import HumanoidEnv
+from gymnasium.utils import EzPickle
+from gymnasium import spaces
+
+
+class HumanoidDirectionEnv(HumanoidEnv, EzPickle):
+    """Humanoid-v5 with a random target heading each episode.
+
+    FIX vs original: the task reward is now BOUNDED in [0, 1], AMP-paper style.
+    Old reward: 5.0 * dot(v, dir) + 0.5 * healthy - ctrl_cost
+      - Unbounded in speed -> PPO maximizes it by sprinting/lunging, which is
+        exactly the behavior the motion prior punishes, so the two rewards
+        fight each other.
+      - Its scale (~2.5 to 7+ per step) dwarfed the [0, 0.2] style reward, so
+        the discriminator was effectively ignored.
+    New reward: full credit once velocity along the target reaches
+    `target_speed` (default 1.4 m/s, a normal walking pace present in mocap),
+    no extra credit for going faster. Falling is handled by termination and by
+    the style reward (mocap contains no falling), so the constant alive bonus
+    is gone.
+    """
+
+    def __init__(self, direction=(1.0, 0.0), target_speed=1.4, **kwargs):
+        self.target_dir = np.array(direction, dtype=np.float32)
+        self.target_dir /= np.linalg.norm(self.target_dir)
+        self.target_speed = float(target_speed)
+
+        EzPickle.__init__(self, direction, target_speed, **kwargs)
+        super().__init__(**kwargs)
+
+        old_space = self.observation_space
+        self.observation_space = spaces.Box(
+            low=-np.inf,
+            high=np.inf,
+            shape=(old_space.shape[0] + 2,),
+            dtype=np.float64,
+        )
+
+    def reset_model(self):
+        obs = super().reset_model()
+        angle = self.np_random.uniform(-np.pi, np.pi)
+        self.target_dir = np.array(
+            [np.cos(angle), np.sin(angle)], dtype=np.float32
+        )
+        return self._get_obs_with_direction(obs)
+
+    def step(self, action):
+        xy_before = self.data.xpos[self.model.body("torso").id][:2].copy()
+        self.do_simulation(action, self.frame_skip)
+        xy_after = self.data.xpos[self.model.body("torso").id][:2].copy()
+
+        xy_velocity = (xy_after - xy_before) / self.dt
+        v_proj = float(np.dot(xy_velocity, self.target_dir))
+
+        # Only penalize the *deficit* below target speed; exceeding it earns
+        # nothing extra. Reward is in [0, 1].
+        speed_deficit = max(0.0, self.target_speed - v_proj)
+        direction_reward = float(np.exp(-2.0 * speed_deficit ** 2))
+
+        reward = direction_reward
+        terminated = not self.is_healthy
+        obs = self._get_obs_with_direction(self._get_obs())
+
+        info = {
+            "xy_velocity": xy_velocity,
+            "target_dir": self.target_dir,
+            "v_proj": v_proj,
+            "direction_reward": direction_reward,
+        }
+
+        if self.render_mode == "human":
+            self.render()
+
+        return obs, reward, terminated, False, info
+
+    def _get_obs_with_direction(self, humanoid_obs):
+        return np.concatenate([humanoid_obs, self.target_dir]).astype(np.float64)

humanoid_direction_evaluate.py

@@ -0,0 +1,79 @@
+import gymnasium as gym
+import register_env
+
+from stable_baselines3 import PPO
+from stable_baselines3.common.vec_env import DummyVecEnv, VecNormalize
+
+
+MODEL_PATH = (
+    "models_exp2/"
+    "ppo_humanoid_direction_amp_fixed.zip"
+)
+
+VECNORMALIZE_PATH = (
+    "models_exp2/"
+    "vecnormalize_amp_fixed.pkl"
+)
+
+
+def make_env():
+    return gym.make(
+        "HumanoidDirection-v0",
+        render_mode="human",
+    )
+
+
+# Create the same vectorized environment structure used during training
+env = DummyVecEnv([make_env])
+
+# Load the saved observation-normalization statistics
+env = VecNormalize.load(
+    VECNORMALIZE_PATH,
+    env,
+)
+
+# Evaluation settings
+env.training = False
+env.norm_reward = False
+
+# Load the matching PPO model and attach the environment
+model = PPO.load(
+    MODEL_PATH,
+    env=env,
+    device="cpu",
+)
+
+# DummyVecEnv.reset() returns only observations
+obs = env.reset()
+
+episode_reward = 0.0
+episode = 0
+num_episodes = 1000
+
+while episode < num_episodes:
+    action, _ = model.predict(
+        obs,
+        deterministic=True,
+    )
+
+    # VecEnv.step() returns four values, not five
+    obs, rewards, dones, infos = env.step(action)
+
+    # rewards and dones are arrays because this is a vectorized environment
+    episode_reward += float(rewards[0])
+
+    if dones[0]:
+        episode += 1
+
+        print(f"Episode: {episode}")
+        print(f"Episode reward: {episode_reward:.2f}")
+        print(f"Episode info: {infos[0]}")
+        print()
+
+        episode_reward = 0.0
+
+        # DummyVecEnv normally resets automatically after done.
+        # The returned obs is already the next episode's initial observation.
+
+env.close()
+

humanoid_direction_record.py

@@ -0,0 +1,47 @@
+import os
+os.environ["MUJOCO_GL"] = "glfw"
+
+import imageio
+import gymnasium as gym
+import register_env
+from stable_baselines3 import PPO
+
+model = PPO.load("models/ppo_humanoid_direction_20m.zip", device="cpu")
+
+env = gym.make("HumanoidDirection-v0", render_mode="rgb_array")
+obs, info = env.reset()
+
+frames = []
+num_episodes = 5
+episode_count = 0
+
+while episode_count < num_episodes:
+    action, _ = model.predict(obs, deterministic=True)
+    obs, reward, terminated, truncated, info = env.step(action)
+
+    frame = env.render()
+    frames.append(frame)
+    if terminated or truncated:
+        episode_count += 1
+        print(f"Episode: {episode_count}")
+        last_frame = frames[-1]
+
+        #pause using last frame
+        for _ in range(15):
+            frames.append(last_frame)
+        if episode_count < num_episodes:
+            obs, info = env.reset()
+            
+
+env.close()
+
+print(f"number of frames: {len(frames)}")
+print(f"frames shape: {frames[0].shape}")
+writer = imageio.get_writer("humanoid_direction_20m.mp4", fps=30)
+
+for frame in frames:
+    writer.append_data(frame)
+
+writer.close()
+
+print(f"saved humanoid_direction.mp4 with episodes: {num_episodes}")

register_env.py

@@ -0,0 +1,9 @@
+from gymnasium.envs.registration import register
+
+print("registering environment")
+
+register(
+    id="HumanoidDirection-v0",
+    entry_point="humanoid_direction_env:HumanoidDirectionEnv",
+    max_episode_steps=1000,
+)

replay.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:827d7f4bc9fd4654913f2c2c53dcfd2cd6254620e7fd81cd5c6ccc4f3eddc240
+size 8157631

train_sb3_real_amp.py

@@ -0,0 +1,164 @@
+from pathlib import Path
+
+import gymnasium as gym
+import torch
+
+import register_env  # noqa: F401
+from stable_baselines3 import PPO
+from stable_baselines3.common.callbacks import CallbackList, CheckpointCallback
+from stable_baselines3.common.monitor import Monitor
+from stable_baselines3.common.vec_env import SubprocVecEnv, VecNormalize
+
+from amp_callback import AMPDiscriminatorCallback
+from amp_discriminator import AMPDiscriminator
+from amp_env import AMPHumanoidEnv
+from motion_lib import MotionLib
+
+# --------------------------------------------------------------------- #
+TOTAL_TIMESTEPS = 50_000_000
+N_ENVS = 8                      # = --cpus-per-task in your slurm file
+N_STEPS = 2048
+TRAIN_FREQ = N_ENVS * N_STEPS   # train the discriminator once per rollout
+
+TASK_WEIGHT = 0.5               # both rewards are in [0, 1] now, so
+AMP_WEIGHT = 0.5                # 0.5 / 0.5 mixing as in the AMP paper
+REFERENCE_STATE_INIT_PROB = 0.5
+
+DISCRIMINATOR_LR = 1e-4
+DISCRIMINATOR_HIDDEN_DIM = 512
+DISC_UPDATES_PER_ROLLOUT = 8
+DISC_BATCH_SIZE = 512
+GRADIENT_PENALTY_WEIGHT = 5.0
+
+BASE_DIR = Path(__file__).resolve().parent
+MODEL_DIR = BASE_DIR / "models"
+CHECKPOINT_DIR = MODEL_DIR / "checkpoints"
+TENSORBOARD_DIR = BASE_DIR / "tensorboard_real_amp_fixed"
+CHECKPOINT_EVERY = 500_000
+# --------------------------------------------------------------------- #
+
+
+def make_env(motion_lib, amp_mean, amp_std):
+    def _init():
+        # Runs inside each SubprocVecEnv worker: keep torch single-threaded so
+        # 8 workers don't oversubscribe the 8-CPU slurm allocation.
+        torch.set_num_threads(1)
+        # Each worker holds its OWN cpu copy of the discriminator for reward
+        # evaluation; the callback pushes fresh weights after every update.
+        disc_local = AMPDiscriminator(
+            input_dim=90, hidden_dim=DISCRIMINATOR_HIDDEN_DIM
+        )
+        env = AMPHumanoidEnv(
+            discriminator=disc_local,
+            motion_lib=motion_lib,
+            task_weight=TASK_WEIGHT,
+            amp_weight=AMP_WEIGHT,
+            device="cpu",
+            amp_mean=amp_mean,
+            amp_std=amp_std,
+            reference_state_init_prob=REFERENCE_STATE_INIT_PROB,
+        )
+        return Monitor(env)
+
+    return _init
+
+
+def main():
+    MODEL_DIR.mkdir(exist_ok=True)
+    CHECKPOINT_DIR.mkdir(exist_ok=True)
+    TENSORBOARD_DIR.mkdir(exist_ok=True)
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    print("Using device:", device)
+
+    # Expert transition spacing must match the env control timestep exactly.
+    _tmp = gym.make("HumanoidDirection-v0")
+    env_dt = float(_tmp.unwrapped.dt)
+    _tmp.close()
+    print("Environment dt:", env_dt)
+
+    motion_lib = MotionLib(
+        str(BASE_DIR / "retargeted_pkl"),
+        transition_dt=env_dt,
+    )
+    amp_mean, amp_std = motion_lib.compute_amp_stats(num_samples=10000)
+
+    # Master discriminator (the one actually trained, on GPU if available).
+    disc = AMPDiscriminator(
+        input_dim=90, hidden_dim=DISCRIMINATOR_HIDDEN_DIM
+    ).to(device)
+    disc_optimizer = torch.optim.Adam(
+        disc.parameters(), lr=DISCRIMINATOR_LR, weight_decay=1e-4
+    )
+
+    # SubprocVecEnv: envs step in parallel processes. With DummyVecEnv all 8
+    # envs ran serially in one process - roughly an 8x throughput loss.
+    env = SubprocVecEnv(
+        [make_env(motion_lib, amp_mean, amp_std) for _ in range(N_ENVS)]
+    )
+    # Observation normalization matters a lot for Humanoid + PPO.
+    env = VecNormalize(env, norm_obs=True, norm_reward=False, clip_obs=10.0)
+
+    amp_callback = AMPDiscriminatorCallback(
+        motion_lib=motion_lib,
+        discriminator=disc,
+        optimizer=disc_optimizer,
+        batch_size=DISC_BATCH_SIZE,
+        updates_per_call=DISC_UPDATES_PER_ROLLOUT,
+        train_freq=TRAIN_FREQ,
+        save_freq=CHECKPOINT_EVERY,
+        save_path=str(CHECKPOINT_DIR),
+        device=device,
+        amp_mean=amp_mean,
+        amp_std=amp_std,
+        fake_replay_size=100_000,
+        gradient_penalty_weight=GRADIENT_PENALTY_WEIGHT,
+        score_reg_weight=1e-4,
+        max_grad_norm=1.0,
+    )
+
+    checkpoint_callback = CheckpointCallback(
+        save_freq=max(CHECKPOINT_EVERY // N_ENVS, 1),
+        save_path=str(CHECKPOINT_DIR),
+        name_prefix="ppo_amp_fixed",
+        save_replay_buffer=False,
+        save_vecnormalize=True,   # needed to evaluate the model later!
+    )
+
+    policy_kwargs = dict(
+        activation_fn=torch.nn.ReLU,
+        net_arch=dict(pi=[1024, 512], vf=[1024, 512]),
+    )
+
+    model = PPO(
+        "MlpPolicy",
+        env,
+        device=device,
+        learning_rate=1e-4,      # 5e-5 was very slow for 20M steps
+        n_steps=N_STEPS,
+        batch_size=512,
+        n_epochs=5,
+        target_kl=0.02,
+        gamma=0.99,
+        gae_lambda=0.95,
+        clip_range=0.2,
+        ent_coef=0.0,            # 0.01 keeps the gaussian noisy -> jittery gait
+        policy_kwargs=policy_kwargs,
+        verbose=1,
+        tensorboard_log=str(TENSORBOARD_DIR),
+    )
+
+    model.learn(
+        total_timesteps=TOTAL_TIMESTEPS,
+        callback=CallbackList([amp_callback, checkpoint_callback]),
+    )
+
+    model.save(str(MODEL_DIR / "ppo_humanoid_direction_amp_fixed"))
+    env.save(str(MODEL_DIR / "vecnormalize_amp_fixed.pkl"))
+    torch.save(disc.state_dict(), MODEL_DIR / "amp_discriminator_fixed.pt")
+    env.close()
+    print("Saved final PPO model, VecNormalize stats, and discriminator.")
+
+
+if __name__ == "__main__":
+    main()

train_sb3_real_amp_exp2.py

@@ -0,0 +1,164 @@
+from pathlib import Path
+
+import gymnasium as gym
+import torch
+
+import register_env  # noqa: F401
+from stable_baselines3 import PPO
+from stable_baselines3.common.callbacks import CallbackList, CheckpointCallback
+from stable_baselines3.common.monitor import Monitor
+from stable_baselines3.common.vec_env import SubprocVecEnv, VecNormalize
+
+from amp_callback import AMPDiscriminatorCallback
+from amp_discriminator import AMPDiscriminator
+from amp_env import AMPHumanoidEnv
+from motion_lib import MotionLib
+
+# --------------------------------------------------------------------- #
+TOTAL_TIMESTEPS = 50_000_000
+N_ENVS = 8                      # = --cpus-per-task in your slurm file
+N_STEPS = 2048
+TRAIN_FREQ = N_ENVS * N_STEPS   # train the discriminator once per rollout
+
+TASK_WEIGHT = 0.5               # both rewards are in [0, 1] now, so
+AMP_WEIGHT = 0.5                # 0.5 / 0.5 mixing as in the AMP paper
+REFERENCE_STATE_INIT_PROB = 0.3
+
+DISCRIMINATOR_LR = 3e-5
+DISCRIMINATOR_HIDDEN_DIM = 512
+DISC_UPDATES_PER_ROLLOUT = 4
+DISC_BATCH_SIZE = 512
+GRADIENT_PENALTY_WEIGHT = 10.0
+
+BASE_DIR = Path(__file__).resolve().parent
+MODEL_DIR = BASE_DIR / "models_exp2"
+CHECKPOINT_DIR = MODEL_DIR / "checkpoints"
+TENSORBOARD_DIR = BASE_DIR / "tensorboard_real_amp_fixed"
+CHECKPOINT_EVERY = 500_000
+# --------------------------------------------------------------------- #
+
+
+def make_env(motion_lib, amp_mean, amp_std):
+    def _init():
+        # Runs inside each SubprocVecEnv worker: keep torch single-threaded so
+        # 8 workers don't oversubscribe the 8-CPU slurm allocation.
+        torch.set_num_threads(1)
+        # Each worker holds its OWN cpu copy of the discriminator for reward
+        # evaluation; the callback pushes fresh weights after every update.
+        disc_local = AMPDiscriminator(
+            input_dim=90, hidden_dim=DISCRIMINATOR_HIDDEN_DIM
+        )
+        env = AMPHumanoidEnv(
+            discriminator=disc_local,
+            motion_lib=motion_lib,
+            task_weight=TASK_WEIGHT,
+            amp_weight=AMP_WEIGHT,
+            device="cpu",
+            amp_mean=amp_mean,
+            amp_std=amp_std,
+            reference_state_init_prob=REFERENCE_STATE_INIT_PROB,
+        )
+        return Monitor(env)
+
+    return _init
+
+
+def main():
+    MODEL_DIR.mkdir(exist_ok=True)
+    CHECKPOINT_DIR.mkdir(exist_ok=True)
+    TENSORBOARD_DIR.mkdir(exist_ok=True)
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    print("Using device:", device)
+
+    # Expert transition spacing must match the env control timestep exactly.
+    _tmp = gym.make("HumanoidDirection-v0")
+    env_dt = float(_tmp.unwrapped.dt)
+    _tmp.close()
+    print("Environment dt:", env_dt)
+
+    motion_lib = MotionLib(
+        str(BASE_DIR / "retargeted_pkl"),
+        transition_dt=env_dt,
+    )
+    amp_mean, amp_std = motion_lib.compute_amp_stats(num_samples=10000)
+
+    # Master discriminator (the one actually trained, on GPU if available).
+    disc = AMPDiscriminator(
+        input_dim=90, hidden_dim=DISCRIMINATOR_HIDDEN_DIM
+    ).to(device)
+    disc_optimizer = torch.optim.Adam(
+        disc.parameters(), lr=DISCRIMINATOR_LR, weight_decay=1e-4
+    )
+
+    # SubprocVecEnv: envs step in parallel processes. With DummyVecEnv all 8
+    # envs ran serially in one process - roughly an 8x throughput loss.
+    env = SubprocVecEnv(
+        [make_env(motion_lib, amp_mean, amp_std) for _ in range(N_ENVS)]
+    )
+    # Observation normalization matters a lot for Humanoid + PPO.
+    env = VecNormalize(env, norm_obs=True, norm_reward=False, clip_obs=10.0)
+
+    amp_callback = AMPDiscriminatorCallback(
+        motion_lib=motion_lib,
+        discriminator=disc,
+        optimizer=disc_optimizer,
+        batch_size=DISC_BATCH_SIZE,
+        updates_per_call=DISC_UPDATES_PER_ROLLOUT,
+        train_freq=TRAIN_FREQ,
+        save_freq=CHECKPOINT_EVERY,
+        save_path=str(CHECKPOINT_DIR),
+        device=device,
+        amp_mean=amp_mean,
+        amp_std=amp_std,
+        fake_replay_size=100_000,
+        gradient_penalty_weight=GRADIENT_PENALTY_WEIGHT,
+        score_reg_weight=1e-4,
+        max_grad_norm=1.0,
+    )
+
+    checkpoint_callback = CheckpointCallback(
+        save_freq=max(CHECKPOINT_EVERY // N_ENVS, 1),
+        save_path=str(CHECKPOINT_DIR),
+        name_prefix="ppo_amp_fixed",
+        save_replay_buffer=False,
+        save_vecnormalize=True,   # needed to evaluate the model later!
+    )
+
+    policy_kwargs = dict(
+        activation_fn=torch.nn.ReLU,
+        net_arch=dict(pi=[1024, 512], vf=[1024, 512]),
+    )
+
+    model = PPO(
+        "MlpPolicy",
+        env,
+        device=device,
+        learning_rate=5e-5,      # 5e-5 was very slow for 20M steps
+        n_steps=N_STEPS,
+        batch_size=512,
+        n_epochs=5,
+        target_kl=0.02,
+        gamma=0.99,
+        gae_lambda=0.95,
+        clip_range=0.2,
+        ent_coef=0.0,            # 0.01 keeps the gaussian noisy -> jittery gait
+        policy_kwargs=policy_kwargs,
+        verbose=1,
+        tensorboard_log=str(TENSORBOARD_DIR),
+    )
+
+    model.learn(
+        total_timesteps=TOTAL_TIMESTEPS,
+        callback=CallbackList([amp_callback, checkpoint_callback]),
+    )
+
+    model.save(str(MODEL_DIR / "ppo_humanoid_direction_amp_fixed"))
+    env.save(str(MODEL_DIR / "vecnormalize_amp_fixed.pkl"))
+    torch.save(disc.state_dict(), MODEL_DIR / "amp_discriminator_fixed.pt")
+    env.close()
+    print("Saved final PPO model, VecNormalize stats, and discriminator.")
+
+
+if __name__ == "__main__":
+    main()