55 lines
1.2 KiB
Python
55 lines
1.2 KiB
Python
"""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()
|