Upload train_humanoid_gpu.py with huggingface_hub

train_humanoid_gpu.py

@@ -24,6 +24,7 @@ import numpy as np
 import gymnasium as gym
 from stable_baselines3 import SAC
 from stable_baselines3.common.callbacks import BaseCallback
+from stable_baselines3.common.vec_env import SubprocVecEnv
 
 # ── Config ───────────────────────────────────────────────────────────────
 TIMESTEPS = 10_000_000
@@ -198,19 +199,27 @@ def main():
     os.makedirs(SAVE_DIR, exist_ok=True)
     os.makedirs(RESULTS_DIR, exist_ok=True)
 
-    env = UprightHumanoidWrapper(gym.make("Humanoid-v5"))
+    N_ENVS = 16
+
+    def make_env(rank):
+        def _init():
+            env = UprightHumanoidWrapper(gym.make("Humanoid-v5"))
+            return env
+        return _init
+
+    env = SubprocVecEnv([make_env(i) for i in range(N_ENVS)])
 
     model = SAC(
         policy="MlpPolicy",
         env=env,
         learning_rate=3e-4,
         buffer_size=1_000_000,
-        batch_size=256,
+        batch_size=512,
         tau=0.005,
         gamma=0.99,
         learning_starts=25_000,
-        train_freq=4,
-        gradient_steps=4,
+        train_freq=1,
+        gradient_steps=2,
         verbose=0,
         device=DEVICE,
         seed=42,