initial commit

examples/train_sb3.py

@@ -0,0 +1,54 @@
+"""Example: train an agent with Stable-Baselines3 using action masking.
+
+Requires: pip install blokus-gym[train]
+"""
+from stable_baselines3 import PPO
+
+from blokus_gym import ActionMaskWrapper, BlokusEnv
+
+
+def main():
+    # Create environment with ActionMaskWrapper
+    env = BlokusEnv(num_players=2, board_size=7, bot_type=None)
+    env = ActionMaskWrapper(env)
+
+    # Train PPO agent
+    model = PPO(
+        "MultiInputPolicy",
+        env,
+        verbose=1,
+        learning_rate=3e-4,
+        n_steps=2048,
+        batch_size=64,
+        n_epochs=10,
+        gamma=0.99,
+        seed=42,
+    )
+
+    print("Training for 10,000 steps...")
+    model.learn(total_timesteps=10000)
+
+    # Save the model
+    model.save("blokus_ppo_agent")
+
+    # Evaluate
+    obs, info = env.reset(seed=42)
+    total_reward = 0
+    steps = 0
+    terminated = False
+    truncated = False
+
+    while not (terminated or truncated):
+        action, _ = model.predict(obs, deterministic=True)
+        obs, reward, terminated, truncated, info = env.step(action)
+        total_reward += reward
+        steps += 1
+
+    print(f"Evaluation: {steps} steps, reward: {total_reward}")
+
+    # Cleanup
+    env.close()
+
+
+if __name__ == "__main__":
+    main()