Self-Play Connect4 with DQN + Curriculum Learning¶
Agent trained to play Connect Four through self-play¶
This tutorial shows how to train a DQN agent on the connect four classic environment.
This tutorial focuses on two techniques used in reinforcement learning - curriculum learning and self-play. Curriculum learning refers to training an agent on tasks of increasing difficulty in separate ‘lessons’. Imagine you were trying to become a chess world champion. You would not decide to learn to play chess by immediately taking on a grand master - it would be too difficult. Instead, you would practice against people of the same ability as you, improve slowly, and increasingly play against harder opponents until you were ready to compete with the best. The same concept applies to reinforcement learning models. Sometimes, tasks are too difficult to learn in one go, and so we must create a curriculum to guide an agent and teach it to solve our ultimate hard environment.
This tutorial also uses self-play. Self-play is a technique used in competitive reinforcement learning environments. An agent trains by playing against a copy of itself - the opponent - and learns to beat this opponent. The opponent is then updated to a copy of this better version of the agent, and the agent must then learn to beat itself again. This is done repeatedly, and the agent iteratively improves by exploiting its own weaknesses and discovering new strategies.
In this tutorial, self-play is treated as the final lesson in the curriculum. However, these two techniques can be used independently of each other, and with unlimited resources, self-play can beat agents trained with human-crafted lessons through curriculum learning. The Bitter Lesson by Richard Sutton provides an interesting take on curriculum learning and is definitely worth consideration from any engineer undertaking such a task. However, unlike Sutton, we do not all have the resources available to us that Deepmind and top institutions provide, and so one must be pragmatic when deciding how they will solve their own reinforcement learning problem. If you would like to discuss this exciting area of research further, please join the AgileRL Discord server and let us know what you think!
What is DQN?¶
DQN (Deep Q-Network) is an extension of Q-learning that makes use of a replay buffer and target network to improve learning stability. For further information on DQN, check out the AgileRL documentation.
Compatible Action Spaces¶
|
|
|
|
|---|---|---|---|
✔️ |
❌ |
❌ |
❌ |
Curriculum learning and Self-play Using DQN on Connect Four¶
The following code should run without any issues. The comments are designed to help you understand how to use PettingZoo with AgileRL. If you have any questions, please feel free to ask in the Discord server.
This is a complicated tutorial, and so we will go through it in stages. The full code can be found at the end of this section. Although much of this tutorial contains content specific to the Connect Four environment, it serves to demonstrate how techniques can be applied more generally to other problems.
Imports¶
Importing the following packages, functions and classes will enable us to run the tutorial.
Imports
import copy
import os
import random
from collections import deque
from datetime import datetime
import numpy as np
import torch
import wandb
import yaml
from tqdm import tqdm
from pettingzoo.classic import connect_four_v3
from agilerl.algorithms import DQN
from agilerl.components.replay_buffer import ReplayBuffer
from agilerl.hpo.mutation import Mutations
from agilerl.hpo.tournament import TournamentSelection
Curriculum Learning¶
First, we need to set up and modify our environment to enable curriculum learning. Curriculum learning is enabled by changing the environment that the agent trains in. This can be implemented by changing what happens when certain actions are taken - altering the next observation returned by the environment, or more simply by altering the reward. First, we will change the reward. By default, Connect Four uses the following rewards:
Win = +1
Lose = -1
Play continues = 0
To help guide our agent, we can introduce rewards for other outcomes in the environment, such as a small reward for placing 3 pieces in a row, or a small negative reward when the opponent manages the same feat. We can also use reward shaping to encourage our agent to explore more. In Connect Four, if playing against a random opponent, an easy way to win is to always play in the same column. An agent may find success doing this, and therefore not learn other, more sophisticated strategies that can help it win against better opponents. We may therefore elect to reward vertical wins slightly less than horizontal or diagonal wins, to encourage the agent to try winning in different ways. An example reward system could be defined as follows:
Win (horizontal or diagonal) = +1
Win (vertical) = +0.8
Three in a row = +0.05
Opponent three in a row = -0.05
Lose = -1
Play continues = 0
Config files¶
It is best to use YAML config files to define the lessons in our curriculum and easily change and keep track of our settings. The first three lessons in our curriculum can be defined as follows:
Lesson 1
---
# Connect Four Lesson 1
# Train against random agent: 'random', weak opponent: 'weak', strong opponent: 'strong', or use self-play: 'self'
opponent: random
opponent_pool_size: # Size of opponent pool for self-play
opponent_upgrade: # Epoch frequency to update opponent pool
eval_opponent: # 'random', 'weak' or 'strong'
pretrained_path: # Path to pretrained model weights
save_path: models/DQN/lesson1_trained_agent.pt # Path to save trained model
max_train_episodes: 0 # Maximum number of training episodes in environment
## Game specific:
buffer_warm_up: true # Fill replay buffer with random experiences
warm_up_opponent: random # Difficulty level of warm up experiences
agent_warm_up: 300 # Number of epochs to warm up agent by training on random experiences
block_vert_coef: 4 # How many times more likely to block vertically
rewards: # Rewards for different outcomes
win: 1
vertical_win: 0.7
three_in_row: 0.05
opp_three_in_row: -0.05
lose: -1
play_continues: 0
Lesson 2
---
# Connect Four Lesson 2
# Train against random agent: 'random', weak opponent: 'weak', strong opponent: 'strong', or use self-play: 'self'
opponent: weak
opponent_pool_size: # Size of opponent pool for self-play
opponent_upgrade: # Epoch frequency to update opponent pool
eval_opponent: weak # 'random', 'weak' or 'strong'
pretrained_path: models/DQN/lesson1_trained_agent.pt # Path to pretrained model weights
save_path: models/DQN/lesson2_trained_agent.pt # Path to save trained model
max_train_episodes: 100000 # Maximum number of training episodes in environment
## Game specific:
buffer_warm_up: false # Fill replay buffer with random experiences
warm_up_opponent: # Difficulty level of warm up experiences
agent_warm_up: 0 # Number of epochs to warm up agent by training on random experiences
block_vert_coef: 1 # How many times more likely to block vertically
rewards: # Rewards for different outcomes
win: 1
vertical_win: 1
three_in_row: 0.02
opp_three_in_row: -0.02
lose: -1
play_continues: 0
Lesson 3
---
# Connect Four Lesson 3
# Train against random agent: 'random', weak opponent: 'weak', strong opponent: 'strong', or use self-play: 'self'
opponent: strong
opponent_pool_size: # Size of opponent pool for self-play
opponent_upgrade: # Epoch frequency to update opponent pool
eval_opponent: strong # 'random', 'weak' or 'strong'
pretrained_path: models/DQN/lesson2_trained_agent.pt # Path to pretrained model weights
save_path: models/DQN/lesson3_trained_agent.pt # Path to save trained model
max_train_episodes: 200000 # Maximum number of training episodes in environment
## Game specific:
buffer_warm_up: false # Fill replay buffer with random experiences
warm_up_opponent: # Difficulty level of warm up experiences
agent_warm_up: 0 # Number of epochs to warm up agent by training on random experiences
block_vert_coef: 1 # How many times more likely to block vertically
rewards: # Rewards for different outcomes
win: 1
vertical_win: 1
three_in_row: 0.02
opp_three_in_row: -0.02
lose: -1
play_continues: 0
To implement our curriculum, we create a CurriculumEnv class that acts as a wrapper on top of our Connect Four environment and enables us
to alter the reward to guide the training of our agent. This uses the configs that we set up to define the lesson.
CurriculumEnv
import random
from pettingzoo.parallel import ParallelEnv
from agilerl.components.data import Transition
from agilerl.components.replay_buffer import ReplayBuffer
class CurriculumEnv:
"""Wrapper around environment to modify reward for curriculum learning.
:param env: Environment to learn in
:type env: PettingZoo-style environment
:param lesson: Lesson settings for curriculum learning
:type lesson: dict
"""
def __init__(self, env: ParallelEnv, lesson: dict):
self.env = env
self.lesson = lesson
def fill_replay_buffer(
self, memory: ReplayBuffer, opponent: "Opponent"
) -> ReplayBuffer:
"""Fill the replay buffer with experiences collected by taking random actions in the environment.
:param memory: Experience replay buffer
:type memory: AgileRL experience replay buffer
:param opponent: Opponent to train against
:type opponent: Opponent
:return: Filled replay buffer
:rtype: ReplayBuffer
"""
print("Filling replay buffer ...")
pbar = tqdm(total=memory.max_size)
while len(memory) < memory.max_size:
# Randomly decide whether random player will go first or second
opponent_first = random.random() > 0.5
mem_full = len(memory)
self.reset() # Reset environment at start of episode
observation, reward, done, truncation, _ = self.last()
(
p1_state,
p1_state_flipped,
p1_action,
p1_next_state,
p1_next_state_flipped,
) = (None, None, None, None, None)
done, truncation = False, False
while not (done or truncation):
# Player 0's turn
p0_action_mask = observation["action_mask"]
p0_state, p0_state_flipped = transform_and_flip(observation, player=0)
if opponent_first:
p0_action = self.env.action_space("player_0").sample(p0_action_mask)
else:
if self.lesson["warm_up_opponent"] == "random":
p0_action = opponent.get_action(
p0_action_mask, p1_action, self.lesson["block_vert_coef"]
)
else:
p0_action = opponent.get_action(player=0)
self.step(p0_action) # Act in environment
observation, env_reward, done, truncation, _ = self.last()
p0_next_state, p0_next_state_flipped = transform_and_flip(
observation, player=0
)
if done or truncation:
reward = self.reward(done=True, player=0)
transition = Transition(
obs=np.concatenate(
(p0_state, p1_state, p0_state_flipped, p1_state_flipped)
),
action=np.array(
[p0_action, p1_action, 6 - p0_action, 6 - p1_action]
),
reward=np.array(
[
reward,
LESSON["rewards"]["lose"],
reward,
LESSON["rewards"]["lose"],
]
),
next_obs=np.concatenate(
(
p0_next_state,
p1_next_state,
p0_next_state_flipped,
p1_next_state_flipped,
)
),
done=np.array([done, done, done, done]),
batch_size=[4],
)
memory.add(transition.to_tensordict(), is_vectorised=True)
else: # Play continues
if p1_state is not None:
reward = self.reward(done=False, player=1)
transition = Transition(
obs=np.concatenate((p1_state, p1_state_flipped)),
action=np.array([p1_action, 6 - p1_action]),
reward=np.array([reward, reward]),
next_obs=np.concatenate(
(p1_next_state, p1_next_state_flipped)
),
done=np.array([done, done]),
batch_size=[2],
)
memory.add(transition.to_tensordict(), is_vectorised=True)
# Player 1's turn
p1_action_mask = observation["action_mask"]
p1_state, p1_state_flipped = transform_and_flip(
observation, player=1
)
if not opponent_first:
p1_action = self.env.action_space("player_1").sample(
p1_action_mask
)
else:
if self.lesson["warm_up_opponent"] == "random":
p1_action = opponent.get_action(
p1_action_mask, p0_action, LESSON["block_vert_coef"]
)
else:
p1_action = opponent.get_action(player=1)
self.step(p1_action) # Act in environment
observation, env_reward, done, truncation, _ = self.last()
p1_next_state, p1_next_state_flipped = transform_and_flip(
observation, player=1
)
if done or truncation:
reward = self.reward(done=True, player=1)
transition = Transition(
obs=np.concatenate(
(p0_state, p1_state, p0_state_flipped, p1_state_flipped)
),
action=np.array(
[p0_action, p1_action, 6 - p0_action, 6 - p1_action]
),
reward=np.array(
[
LESSON["rewards"]["lose"],
reward,
LESSON["rewards"]["lose"],
reward,
]
),
next_obs=np.concatenate(
(
p0_next_state,
p1_next_state,
p0_next_state_flipped,
p1_next_state_flipped,
)
),
done=np.array([done, done, done, done]),
batch_size=[4],
)
memory.add(transition.to_tensordict(), is_vectorised=True)
else: # Play continues
reward = self.reward(done=False, player=0)
transition = Transition(
obs=np.concatenate((p0_state, p0_state_flipped)),
action=np.array([p0_action, 6 - p0_action]),
reward=np.array([reward, reward]),
next_obs=np.concatenate(
(p0_next_state, p0_next_state_flipped)
),
done=np.array([done, done]),
batch_size=[2],
)
memory.add(transition.to_tensordict(), is_vectorised=True)
pbar.update(len(memory) - mem_full)
pbar.close()
print("Replay buffer warmed up.")
return memory
def check_winnable(self, lst: list[int], piece: int) -> bool:
"""Checks if four pieces in a row represent a winnable opportunity, e.g. [1, 1, 1, 0] or [2, 0, 2, 2].
:param lst: List of pieces in row
:type lst: list
:param piece: Player piece we are checking (1 or 2)
:type piece: int
"""
return lst.count(piece) == 3 and lst.count(0) == 1
def check_vertical_win(self, player: int) -> bool:
"""Checks if a win is vertical.
:param player: Player who we are checking, 0 or 1
:type player: int
"""
board = np.array(self.env.env.board).reshape(6, 7)
piece = player + 1
column_count = 7
row_count = 6
# Check vertical locations for win
for c in range(column_count):
for r in range(row_count - 3):
if (
board[r][c] == piece
and board[r + 1][c] == piece
and board[r + 2][c] == piece
and board[r + 3][c] == piece
):
return True
return False
def check_three_in_row(self, player: int) -> int:
"""Checks if there are three pieces in a row and a blank space next, or two pieces - blank - piece.
:param player: Player who we are checking, 0 or 1
:type player: int
"""
board = np.array(self.env.env.board).reshape(6, 7)
piece = player + 1
# Check horizontal locations
column_count = 7
row_count = 6
three_in_row_count = 0
# Check vertical locations
for c in range(column_count):
for r in range(row_count - 3):
if self.check_winnable(board[r : r + 4, c].tolist(), piece):
three_in_row_count += 1
# Check horizontal locations
for r in range(row_count):
for c in range(column_count - 3):
if self.check_winnable(board[r, c : c + 4].tolist(), piece):
three_in_row_count += 1
# Check positively sloped diagonals
for c in range(column_count - 3):
for r in range(row_count - 3):
if self.check_winnable(
[
board[r, c],
board[r + 1, c + 1],
board[r + 2, c + 2],
board[r + 3, c + 3],
],
piece,
):
three_in_row_count += 1
# Check negatively sloped diagonals
for c in range(column_count - 3):
for r in range(3, row_count):
if self.check_winnable(
[
board[r, c],
board[r - 1, c + 1],
board[r - 2, c + 2],
board[r - 3, c + 3],
],
piece,
):
three_in_row_count += 1
return three_in_row_count
def reward(self, done: bool, player: int) -> float:
"""Processes and returns reward from environment according to lesson criteria.
:param done: Environment has terminated
:type done: bool
:param player: Player who we are checking, 0 or 1
:type player: int
"""
if done:
reward = (
self.lesson["rewards"]["vertical_win"]
if self.check_vertical_win(player)
else self.lesson["rewards"]["win"]
)
else:
agent_three_count = self.check_three_in_row(1 - player)
opp_three_count = self.check_three_in_row(player)
if (agent_three_count + opp_three_count) == 0:
reward = self.lesson["rewards"]["play_continues"]
else:
reward = (
self.lesson["rewards"]["three_in_row"] * agent_three_count
+ self.lesson["rewards"]["opp_three_in_row"] * opp_three_count
)
return reward
def last(self) -> tuple[dict, float, bool, bool, dict]:
"""Wrapper around PettingZoo env last method."""
return self.env.last()
def step(self, action: int) -> None:
"""Wrapper around PettingZoo env step method."""
self.env.step(action)
def reset(self) -> None:
"""Wrapper around PettingZoo env reset method."""
self.env.reset()
When defining the different lessons in our curriculum, we can increase the difficulty of our task by modifying environment observations for
our agent - in Connect Four, we can increase the skill level of our opponent. By progressively doing this, we can help our agent improve. We
can change our rewards between lessons too; for example, we may wish to reward wins in all directions equally once we have learned to beat a
random agent and now wish to train against a harder opponent. In this tutorial, an Opponent class is implemented to provide different
levels of difficulty for training our agent.
Opponent
class Opponent:
"""Connect 4 opponent to train and/or evaluate against.
:param env: Environment to learn in
:type env: PettingZoo-style environment
:param difficulty: Difficulty level of opponent, 'random', 'weak' or 'strong'
:type difficulty: str
"""
def __init__(self, env: ParallelEnv, difficulty: str):
self.env = env.env
self.difficulty = difficulty
if self.difficulty == "random":
self.get_action = self.random_opponent
elif self.difficulty == "weak":
self.get_action = self.weak_rule_based_opponent
else:
self.get_action = self.strong_rule_based_opponent
self.num_cols = 7
self.num_rows = 6
self.length = 4
self.top = [0] * self.num_cols
def update_top(self) -> None:
"""Updates self.top, a list which tracks the row on top of the highest piece in each column."""
board = np.array(self.env.env.board).reshape(self.num_rows, self.num_cols)
non_zeros = np.where(board != 0)
rows, cols = non_zeros
top = np.zeros(board.shape[1], dtype=int)
for col in range(board.shape[1]):
column_pieces = rows[cols == col]
if len(column_pieces) > 0:
top[col] = np.min(column_pieces) - 1
else:
top[col] = 5
full_columns = np.all(board != 0, axis=0)
top[full_columns] = 6
self.top = top
def random_opponent(
self,
action_mask: list[int],
last_opp_move: int | None = None,
block_vert_coef: float = 1,
) -> int:
"""Takes move for random opponent. If the lesson aims to randomly block vertical
wins with a higher probability, this is done here too.
:param action_mask: Mask of legal actions: 1=legal, 0=illegal
:type action_mask: list
:param last_opp_move: Most recent action taken by agent against this opponent
:type last_opp_move: int
:param block_vert_coef: How many times more likely to block vertically
:type block_vert_coef: float
"""
if last_opp_move is not None:
action_mask[last_opp_move] *= block_vert_coef
action = random.choices(list(range(self.num_cols)), action_mask)[0]
return action
def weak_rule_based_opponent(self, player: int) -> int:
"""Takes move for weak rule-based opponent.
:param player: Player who we are checking, 0 or 1
:type player: int
"""
self.update_top()
max_length = -1
best_actions = []
for action in range(self.num_cols):
possible, reward, ended, lengths = self.outcome(
action, player, return_length=True
)
if possible and lengths.sum() > max_length:
best_actions = []
max_length = lengths.sum()
if possible and lengths.sum() == max_length:
best_actions.append(action)
best_action = random.choice(best_actions)
return best_action
def strong_rule_based_opponent(self, player: int) -> int:
"""Takes move for strong rule-based opponent.
:param player: Player who we are checking, 0 or 1
:type player: int
"""
self.update_top()
winning_actions = []
for action in range(self.num_cols):
possible, reward, ended = self.outcome(action, player)
if possible and ended:
winning_actions.append(action)
if len(winning_actions) > 0:
winning_action = random.choice(winning_actions)
return winning_action
opp = 1 if player == 0 else 0
loss_avoiding_actions = []
for action in range(self.num_cols):
possible, reward, ended = self.outcome(action, opp)
if possible and ended:
loss_avoiding_actions.append(action)
if len(loss_avoiding_actions) > 0:
loss_avoiding_action = random.choice(loss_avoiding_actions)
return loss_avoiding_action
return self.weak_rule_based_opponent(player) # take best possible move
def outcome(
self, action: int, player: int, return_length: bool = False
) -> tuple[bool, float | None, bool, np.ndarray | None]:
"""Takes move for weak rule-based opponent.
:param action: Action to take in environment
:type action: int
:param player: Player who we are checking, 0 or 1
:type player: int
:param return_length: Return length of outcomes, defaults to False
:type return_length: bool, optional
:return: Tuple containing whether the action is possible, the reward for the action, whether the game has ended, and the length of the outcome
:rtype: tuple[bool, float | None, bool, np.ndarray | None]
"""
if not (self.top[action] < self.num_rows): # action column is full
return (False, None, None) + ((None,) if return_length else ())
row, col = self.top[action], action
piece = player + 1
# down, up, left, right, down-left, up-right, down-right, up-left,
directions = np.array(
[
[[-1, 0], [1, 0]],
[[0, -1], [0, 1]],
[[-1, -1], [1, 1]],
[[-1, 1], [1, -1]],
]
) # |4x2x2|
positions = np.array([row, col]).reshape(1, 1, 1, -1) + np.expand_dims(
directions, -2
) * np.arange(1, self.length).reshape(
1, 1, -1, 1
) # |4x2x3x2|
valid_positions = np.logical_and(
np.logical_and(
positions[:, :, :, 0] >= 0, positions[:, :, :, 0] < self.num_rows
),
np.logical_and(
positions[:, :, :, 1] >= 0, positions[:, :, :, 1] < self.num_cols
),
) # |4x2x3|
d0 = np.where(valid_positions, positions[:, :, :, 0], 0)
d1 = np.where(valid_positions, positions[:, :, :, 1], 0)
board = np.array(self.env.env.board).reshape(self.num_rows, self.num_cols)
board_values = np.where(valid_positions, board[d0, d1], 0)
a = (board_values == piece).astype(int)
b = np.concatenate(
(a, np.zeros_like(a[:, :, :1])), axis=-1
) # padding with zeros to compute length
lengths = np.argmin(b, -1)
ended = False
# check if winnable in any direction
for both_dir in board_values:
# |2x3|
line = np.concatenate((both_dir[0][::-1], [piece], both_dir[1]))
if "".join(map(str, [piece] * self.length)) in "".join(map(str, line)):
ended = True
break
# ended = np.any(np.greater_equal(np.sum(lengths, 1), self.length - 1))
draw = True
for c, v in enumerate(self.top):
draw &= (v == self.num_rows) if c != col else (v == (self.num_rows - 1))
ended |= draw
reward = (-1) ** (player) if ended and not draw else 0
return (True, reward, ended) + ((lengths,) if return_length else ())
General Setup¶
Before we go any further in this tutorial, it would be helpful to define and set up everything remaining we need for training.
Setup code
from gymnasium import spaces
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print("===== AgileRL Curriculum Learning Demo =====")
lesson_number = 1
# Load lesson for curriculum
with open(f"./curriculums/connect_four/lesson{lesson_number}.yaml") as file:
LESSON = yaml.safe_load(file)
# Define the connect four environment
env = connect_four_v3.env()
env.reset()
# Configure the algo input arguments
observation_spaces = [
env.observation_space(agent)["observation"] for agent in env.agents
]
action_spaces = [env.action_space(agent) for agent in env.agents]
# Warp the environment in the curriculum learning wrapper
env = CurriculumEnv(env, LESSON)
# Pre-process dimensions for PyTorch layers (channels-first: C, H, W)
# We only need to worry about the state dim of a single agent
_ob = observation_spaces[0]
observation_space = spaces.Box(
low=np.moveaxis(_ob.low, -1, -3),
high=np.moveaxis(_ob.high, -1, -3),
dtype=_ob.dtype,
)
action_space = action_spaces[0]
# Configure network architecture
net_config = {
"encoder_config": {
"channel_size": [128], # CNN channel size
"kernel_size": [4], # CNN kernel size
"stride_size": [1], # CNN stride size
},
"head_config": {"hidden_size": [64, 64]}, # Network head hidden size
}
# Algorithm hyperparameters
init_hp = {
"double": True,
"batch_size": 256,
"lr": 1e-4,
"gamma": 0.99,
"learn_step": 2,
"tau": 0.01,
}
# Initialize population
population_size = 6
pop = DQN.population(
size=population_size,
observation_space=observation_space,
action_space=action_space,
net_config=net_config,
device=device,
**init_hp,
)
memory_size = 20000
# Configure the replay buffer
memory = ReplayBuffer(
max_size=memory_size, # Max replay buffer size
device=device,
)
# Instantiate a tournament selection object (used for HPO)
tournament = TournamentSelection(
tournament_size=2, # Tournament selection size
elitism=True, # Elitism in tournament selection
population_size=population_size, # Population size
)
# Instantiate a mutations object (used for HPO)
mutations = Mutations(
no_mutation=0.2, # Probability of no mutation
architecture=0, # Probability of architecture mutation
new_layer_prob=0.2, # Probability of new layer mutation
parameters=0.2, # Probability of parameter mutation
activation=0, # Probability of activation function mutation
rl_hp=0.2, # Probability of RL hyperparameter mutation
mutation_sd=0.1, # Mutation strength
rand_seed=1,
device=device,
)
# Define training loop parameters
episodes_per_epoch = 10
max_episodes = LESSON["max_train_episodes"] # Total episodes
max_steps = 500 # Maximum steps to take in each episode
evo_epochs = 20 # Evolution frequency
evo_loop = 50 # Number of evaluation episodes
elite = pop[0] # Assign a placeholder "elite" agent
epsilon = 1.0 # Starting epsilon value
eps_end = 0.1 # Final epsilon value
eps_decay = 0.9998 # Epsilon decays
opp_update_counter = 0
As part of the curriculum, we may also choose to fill the replay buffer with random experiences, and also train on these offline.
Fill Replay Buffer
# Perform buffer and agent warmups if desired
if LESSON["buffer_warm_up"]:
warm_up_opponent = Opponent(env, difficulty=LESSON["warm_up_opponent"])
memory = env.fill_replay_buffer(
memory, warm_up_opponent
) # Fill replay buffer with transitions
if LESSON["agent_warm_up"] > 0:
print("Warming up agents ...")
agent = pop[0]
# Train on randomly collected samples
for epoch in trange(LESSON["agent_warm_up"]):
experiences = memory.sample(agent.batch_size)
agent.learn(experiences)
pop = [agent.clone() for _ in pop]
elite = agent
print("Agent population warmed up.")
The observation space of Connect Four is (6, 7, 2), where the first two dimensions represent the board and the third dimension represents the player. As PyTorch uses channels-first by default, we need to preprocess the observation. Moreover, we need to flip and swap the planes of the observation to account for the fact that the agent will play as both player 0 and player 1. We can define a function to do this as follows:
Transform and Flip
def transform_and_flip(observation, player):
"""Transforms and flips observation for input to agent's neural network.
:param observation: Observation to preprocess
:type observation: dict[str, np.ndarray]
:param player: Player, 0 or 1
:type player: int
"""
state = observation["observation"]
# Pre-process dimensions for PyTorch (N, C, H, W)
state = np.moveaxis(state, -1, -3)
if player == 1:
# Swap pieces so that the agent always sees the board from the same perspective
state[[0, 1], :, :] = state[[1, 0], :, :]
state_flipped = np.expand_dims(np.flip(state, 2), 0)
state = np.expand_dims(state, 0)
return state, state_flipped
Self-play¶
In this tutorial, we use self-play as the final lesson in our curriculum. By iteratively improving our agent and making it learn to win against itself, we can allow it to discover new strategies and achieve higher performance. The weights of our pretrained agent from an earlier lesson can be loaded to the population as follows:
Load Pretrained Weights
from agilerl.algorithms.core.wrappers import OptimizerWrapper
if LESSON["pretrained_path"] is not None:
for agent in pop:
# Load pretrained checkpoint
agent.load_checkpoint(LESSON["pretrained_path"])
# Reinit optimizer for new task
agent.lr = init_hp["lr"]
agent.optimizer = OptimizerWrapper(
torch.optim.Adam,
networks=agent.actor,
lr=agent.lr,
network_names=agent.optimizer.network_names,
lr_name=agent.optimizer.lr_name,
optimizer_kwargs={"capturable": agent.capturable},
)
To train against an old version of our agent, we create a pool of opponents. At training time, we randomly select an opponent from this pool. At regular intervals, we update the opponent pool by removing the oldest opponent and adding a copy of the latest version of our agent. This provides a balance between training against an increasingly difficult opponent and providing variety in the moves an opponent might make.
Create Opponent Pool
if LESSON["opponent"] == "self":
# Create initial pool of opponents
opponent_pool = deque(maxlen=LESSON["opponent_pool_size"])
for _ in range(LESSON["opponent_pool_size"]):
opp = copy.deepcopy(pop[0])
opp.actor.load_state_dict(pop[0].actor.state_dict())
opp.actor.eval()
opponent_pool.append(opp)
A sample lesson config for self-play training could be defined as follows:
Lesson 4
---
# Connect Four Lesson 4
# Train against random agent: 'random', weak opponent: 'weak', strong opponent: 'strong', or use self-play: 'self'
opponent: self
opponent_pool_size: 6 # Size of opponent pool for self-play
opponent_upgrade: 6000 # Epoch frequency to update opponent pool
eval_opponent: strong # 'random', 'weak' or 'strong'
pretrained_path: models/DQN/lesson3_trained_agent.pt # Path to pretrained model weights
save_path: models/DQN/lesson4_trained_agent.pt # Path to save trained model
max_train_episodes: 600000 # Maximum number of training episodes in environment
## Game specific:
buffer_warm_up: false # Fill replay buffer with random experiences
warm_up_opponent: # Difficulty level of warm up experiences
agent_warm_up: 0 # Number of epochs to warm up agent by training on random experiences
block_vert_coef: 1 # How many times more likely to block vertically if playing random opponent
rewards: # Rewards for different outcomes
win: 1
vertical_win: 1
three_in_row: 0.01
opp_three_in_row: -0.01
lose: -1
play_continues: 0
It could also be possible to train an agent through self-play only, without using any previous lessons in the curriculum. This would require significant training time, but could ultimately result in better performance than other methods, and could avoid some of the mistakes discussed in The Bitter Lesson.
Training Loop¶
The Connect Four training loop must take into account that the agent only takes an action every other interaction with the environment (the opponent takes alternating turns). This must be considered when saving transitions to the replay buffer. Equally, we must wait for the outcome of the next player’s turn before determining what the reward should be for a transition. This is not a true Markov Decision Process for this reason, but we can still train a reinforcement learning agent reasonably successfully in these non-stationary conditions.
At regular intervals, we evaluate the performance, or ‘fitness’, of the agents in our population, and do an evolutionary step. Those which perform best are more likely to become members of the next generation, and the hyperparameters and neural architectures of agents in the population are mutated. This evolution allows us to optimize hyperparameters and maximise the performance of our agents in a single training run.
Training Loop
if max_episodes > 0:
wandb.init(
# set the wandb project where this run will be logged
project="AgileRL",
name="{}-EvoHPO-{}-{}Opposition-CNN-{}".format(
"connect_four_v3",
"DQN",
LESSON["opponent"],
datetime.now().strftime("%m%d%Y%H%M%S"),
),
# track hyperparameters and run metadata
config={
"algo": "Evo HPO Rainbow DQN",
"env": "connect_four_v3",
"init_hp": init_hp,
"lesson": LESSON,
},
)
total_steps = 0
total_episodes = 0
pbar = trange(int(max_episodes / episodes_per_epoch))
# Training loop
for idx_epi in pbar:
turns_per_episode = []
train_actions_hist = [0] * action_spaces[0].n
for agent in pop: # Loop through population
for episode in range(episodes_per_epoch):
env.reset() # Reset environment at start of episode
observation, cumulative_reward, done, truncation, _ = env.last()
(
p1_state,
p1_state_flipped,
p1_action,
p1_next_state,
p1_next_state_flipped,
) = (None, None, None, None, None)
if LESSON["opponent"] == "self":
# Randomly choose opponent from opponent pool if using self-play
opponent = random.choice(opponent_pool)
else:
# Create opponent of desired difficulty
opponent = Opponent(env, difficulty=LESSON["opponent"])
# Randomly decide whether agent will go first or second
opponent_first = random.random() > 0.5
score = 0
turns = 0 # Number of turns counter
for idx_step in range(max_steps):
# Player 0"s turn
p0_action_mask = observation["action_mask"]
p0_state, p0_state_flipped = transform_and_flip(
observation, player=0
)
if opponent_first:
if LESSON["opponent"] == "self":
p0_action = opponent.get_action(
p0_state, 0, p0_action_mask
)[0]
elif LESSON["opponent"] == "random":
p0_action = opponent.get_action(
p0_action_mask, p1_action, LESSON["block_vert_coef"]
)
else:
p0_action = opponent.get_action(player=0)
else:
p0_action = agent.get_action(
p0_state, epsilon, p0_action_mask
)[
0
] # Get next action from agent
train_actions_hist[p0_action] += 1
env.step(p0_action) # Act in environment
observation, cumulative_reward, done, truncation, _ = env.last()
p0_next_state, p0_next_state_flipped = transform_and_flip(
observation, player=0
)
if not opponent_first:
score = cumulative_reward
turns += 1
# Check if game is over (Player 0 win)
if done or truncation:
reward = env.reward(done=True, player=0)
transition = Transition(
obs=np.concatenate(
(
p0_state,
p1_state,
p0_state_flipped,
p1_state_flipped,
)
),
action=np.array(
[p0_action, p1_action, 6 - p0_action, 6 - p1_action]
),
reward=np.array(
[
reward,
LESSON["rewards"]["lose"],
reward,
LESSON["rewards"]["lose"],
]
),
next_obs=np.concatenate(
(
p0_next_state,
p1_next_state,
p0_next_state_flipped,
p1_next_state_flipped,
)
),
done=np.array([done, done, done, done]),
batch_size=[4],
)
memory.add(transition.to_tensordict(), is_vectorised=True)
else: # Play continues
if p1_state is not None:
reward = env.reward(done=False, player=1)
transition = Transition(
obs=np.concatenate((p1_state, p1_state_flipped)),
action=np.array([p1_action, 6 - p1_action]),
reward=np.array([reward, reward]),
next_obs=np.concatenate(
(p1_next_state, p1_next_state_flipped)
),
done=np.array([done, done]),
batch_size=[2],
)
memory.add(
transition.to_tensordict(), is_vectorised=True
)
# Player 1"s turn
p1_action_mask = observation["action_mask"]
p1_state, p1_state_flipped = transform_and_flip(
observation, player=1
)
if not opponent_first:
if LESSON["opponent"] == "self":
p1_action = opponent.get_action(
p1_state, 0, p1_action_mask
)[0]
elif LESSON["opponent"] == "random":
p1_action = opponent.get_action(
p1_action_mask,
p0_action,
LESSON["block_vert_coef"],
)
else:
p1_action = opponent.get_action(player=1)
else:
p1_action = agent.get_action(
p1_state, epsilon, p1_action_mask
)[
0
] # Get next action from agent
train_actions_hist[p1_action] += 1
env.step(p1_action) # Act in environment
observation, cumulative_reward, done, truncation, _ = (
env.last()
)
p1_next_state, p1_next_state_flipped = transform_and_flip(
observation, player=1
)
if opponent_first:
score = cumulative_reward
turns += 1
# Check if game is over (Player 1 win)
if done or truncation:
reward = env.reward(done=True, player=1)
transition = Transition(
obs=np.concatenate(
(
p0_state,
p1_state,
p0_state_flipped,
p1_state_flipped,
)
),
action=np.array(
[
p0_action,
p1_action,
6 - p0_action,
6 - p1_action,
]
),
reward=np.array(
[
reward,
LESSON["rewards"]["lose"],
reward,
LESSON["rewards"]["lose"],
]
),
next_obs=np.concatenate(
(
p0_next_state,
p1_next_state,
p0_next_state_flipped,
p1_next_state_flipped,
)
),
done=np.array([done, done, done, done]),
batch_size=[4],
)
memory.add(
transition.to_tensordict(), is_vectorised=True
)
else: # Play continues
reward = env.reward(done=False, player=0)
transition = Transition(
obs=np.concatenate((p0_state, p0_state_flipped)),
action=np.array([p0_action, 6 - p0_action]),
reward=np.array([reward, reward]),
next_obs=np.concatenate(
(p0_next_state, p0_next_state_flipped)
),
done=np.array([done, done]),
batch_size=[2],
)
memory.add(
transition.to_tensordict(), is_vectorised=True
)
# Learn according to learning frequency
if (memory.counter % agent.learn_step == 0) and (
len(memory) >= agent.batch_size
):
# Sample replay buffer
# Learn according to agent"s RL algorithm
experiences = memory.sample(agent.batch_size)
agent.learn(experiences)
# Stop episode if any agents have terminated
if done or truncation:
break
total_steps += idx_step + 1
total_episodes += 1
turns_per_episode.append(turns)
# Save the total episode reward
agent.scores.append(score)
if LESSON["opponent"] == "self":
if (total_episodes % LESSON["opponent_upgrade"] == 0) and (
(idx_epi + 1) > evo_epochs
):
elite_opp, _, _ = tournament._elitism(pop)
elite_opp.actor.eval()
opponent_pool.append(elite_opp)
opp_update_counter += 1
# Update epsilon for exploration
epsilon = max(eps_end, epsilon * eps_decay)
mean_turns = np.mean(turns_per_episode)
# Now evolve population if necessary
if (idx_epi + 1) % evo_epochs == 0:
# Evaluate population vs random actions
fitnesses = []
win_rates = []
eval_actions_hist = [0] * action_spaces[0].n # Eval actions histogram
eval_turns = 0 # Eval turns counter
for agent in pop:
with torch.no_grad():
rewards = []
for i in range(evo_loop):
env.reset() # Reset environment at start of episode
observation, cumulative_reward, done, truncation, _ = (
env.last()
)
player = -1 # Tracker for which player"s turn it is
# Create opponent of desired difficulty
opponent = Opponent(env, difficulty=LESSON["eval_opponent"])
# Randomly decide whether agent will go first or second
if random.random() > 0.5:
opponent_first = False
else:
opponent_first = True
score = 0
for idx_step in range(max_steps):
action_mask = observation["action_mask"]
if player < 0:
if opponent_first:
if LESSON["eval_opponent"] == "random":
action = opponent.get_action(action_mask)
else:
action = opponent.get_action(player=0)
else:
state = np.moveaxis(
observation["observation"], [-1], [-3]
)
state = np.expand_dims(state, 0)
action = agent.get_action(
state, 0, action_mask
)[
0
] # Get next action from agent
eval_actions_hist[action] += 1
if player > 0:
if not opponent_first:
if LESSON["eval_opponent"] == "random":
action = opponent.get_action(action_mask)
else:
action = opponent.get_action(player=1)
else:
state = np.moveaxis(
observation["observation"], [-1], [-3]
)
state[[0, 1], :, :] = state[[1, 0], :, :]
state = np.expand_dims(state, 0)
action = agent.get_action(
state, 0, action_mask
)[
0
] # Get next action from agent
eval_actions_hist[action] += 1
env.step(action) # Act in environment
observation, cumulative_reward, done, truncation, _ = (
env.last()
)
if (player > 0 and opponent_first) or (
player < 0 and not opponent_first
):
score = cumulative_reward
eval_turns += 1
if done or truncation:
break
player *= -1
rewards.append(score)
mean_fit = np.mean(rewards)
agent.fitness.append(mean_fit)
fitnesses.append(mean_fit)
eval_turns = eval_turns / len(pop) / evo_loop
pbar.set_postfix_str(
f"Train Mean Score: {np.mean(agent.scores[-episodes_per_epoch:])} "
f"Train Mean Turns: {mean_turns} "
f"Eval Mean Fitness: {np.mean(fitnesses)} "
f"Eval Best Fitness: {np.max(fitnesses)} "
f"Eval Mean Turns: {eval_turns} "
f"Total Steps: {total_steps}"
)
pbar.update(0)
# Tournament selection and population mutation
elite, pop, _ = tournament.select(pop)
pop = mutations.mutation(pop)
# Save the trained agent
save_path = LESSON["save_path"]
os.makedirs(os.path.dirname(save_path), exist_ok=True)
elite.save_checkpoint(save_path)
print(f"Elite agent saved to '{save_path}'.")
pbar.close()
Trained Model Weights¶
Trained model weights are provided at AgileRL/tutorials/pettingzoo/models. Take a look, train against these models, and see if you can beat them!
Watch the Trained Agents Play¶
The following code allows you to load your saved DQN agent from the previous training block, test the agent’s performance, and then visualise a number of episodes as a gif.
Render Trained Agents
# Copyright 2026 AgileRL
# SPDX-License-Identifier: Apache-2.0
import os
import imageio
import numpy as np
import torch
from dqn_curriculum import Opponent, transform_and_flip
from pettingzoo.classic import connect_four_v3
from PIL import Image, ImageDraw, ImageFont
from agilerl.algorithms.dqn import DQN
# Define function to return image
def _label_with_episode_number(frame, episode_num, frame_no, p):
im = Image.fromarray(frame)
drawer = ImageDraw.Draw(im)
text_color = (255, 255, 255)
font = ImageFont.truetype("arial.ttf", size=45)
drawer.text(
(100, 5),
f"Episode: {episode_num + 1} Frame: {frame_no}",
fill=text_color,
font=font,
)
if p == 1:
player = "Player 1"
color = (255, 0, 0)
if p == 2:
player = "Player 2"
color = (100, 255, 150)
if p is None:
player = "Self-play"
color = (255, 255, 255)
drawer.text((700, 5), f"Agent: {player}", fill=color, font=font)
return im
# Resizes frames to make file size smaller
def resize_frames(frames, fraction):
resized_frames = []
for img in frames:
new_width = int(img.width * fraction)
new_height = int(img.height * fraction)
img_resized = img.resize((new_width, new_height))
resized_frames.append(np.array(img_resized))
return resized_frames
if __name__ == "__main__":
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
path = "./models/DQN/lesson3_trained_agent.pt" # Path to saved agent checkpoint
env = connect_four_v3.env(render_mode="rgb_array")
env.reset()
# Load the saved agent
dqn = DQN.load(path, device)
for opponent_difficulty in ["random", "weak", "strong", "self"]:
# Create opponent
if opponent_difficulty == "self":
opponent = dqn
else:
opponent = Opponent(env, opponent_difficulty)
# Define test loop parameters
episodes = 2 # Number of episodes to test agent on
max_steps = (
500 # Max number of steps to take in the environment in each episode
)
frames = [] # List to collect frames
print("============================================")
print(f"Agent: {path}")
print(f"Opponent: {opponent_difficulty}")
# Test loop for inference
for ep in range(episodes):
if ep / episodes < 0.5:
opponent_first = False
p = 1
else:
opponent_first = True
p = 2
if opponent_difficulty == "self":
p = None
env.reset() # Reset environment at start of episode
frame = env.render()
frames.append(
_label_with_episode_number(frame, episode_num=ep, frame_no=0, p=p),
)
observation, reward, done, truncation, _ = env.last()
player = -1 # Tracker for which player's turn it is
score = 0
for step in range(max_steps):
action_mask = observation["action_mask"]
if player < 0:
state, _ = transform_and_flip(observation, player=0)
if opponent_first:
if opponent_difficulty == "self":
action = opponent.get_action(
state,
epsilon=0,
action_mask=action_mask,
)[0]
elif opponent_difficulty == "random":
action = opponent.get_action(action_mask)
else:
action = opponent.get_action(player=0)
else:
action = dqn.get_action(
state,
epsilon=0,
action_mask=action_mask,
)[0] # Get next action from agent
if player > 0:
state, _ = transform_and_flip(observation, player=1)
if not opponent_first:
if opponent_difficulty == "self":
action = opponent.get_action(
state,
epsilon=0,
action_mask=action_mask,
)[0]
elif opponent_difficulty == "random":
action = opponent.get_action(action_mask)
else:
action = opponent.get_action(player=1)
else:
action = dqn.get_action(
state,
epsilon=0,
action_mask=action_mask,
)[0] # Get next action from agent
env.step(action) # Act in environment
observation, reward, termination, truncation, _ = env.last()
# Save the frame for this step and append to frames list
frame = env.render()
frames.append(
_label_with_episode_number(
frame,
episode_num=ep,
frame_no=step,
p=p,
),
)
if (player > 0 and opponent_first) or (
player < 0 and not opponent_first
):
score += reward
else:
score -= reward
# Stop episode if any agents have terminated
if truncation or termination:
break
player *= -1
print("-" * 15, f"Episode: {ep + 1}", "-" * 15)
print(f"Episode length: {step}")
print(f"Score: {score}")
print("============================================")
frames = resize_frames(frames, 0.5)
# Save the gif to specified path
gif_path = "./videos/"
os.makedirs(gif_path, exist_ok=True)
imageio.mimwrite(
os.path.join("./videos/", f"connect_four_{opponent_difficulty}_opp.gif"),
frames,
duration=400,
loop=True,
)
env.close()
Full Training Code¶
Full code
# Copyright 2026 AgileRL
# SPDX-License-Identifier: Apache-2.0
"""This tutorial shows how to train a DQN agent on the connect four environment, using curriculum learning and self play.
Author: Nick (https://github.com/nicku-a)
"""
import os
import random
from collections import deque
from datetime import datetime, timezone
import gymnasium as gym
import numpy as np
import torch
import wandb
import yaml
from pettingzoo import ParallelEnv
from pettingzoo.classic import connect_four_v3
from tensordict import TensorDict
from tqdm import trange
from agilerl.algorithms import DQN
from agilerl.algorithms.core.registry import HyperparameterConfig, RLParameter
from agilerl.components.replay_buffer import ReplayBuffer
from agilerl.hpo.mutation import Mutations
from agilerl.hpo.tournament import TournamentSelection
class CurriculumEnv:
"""Wrapper around environment to modify reward for curriculum learning.
:param env: Environment to learn in
:type env: PettingZoo-style environment
:param lesson: Lesson settings for curriculum learning
:type lesson: dict
"""
def __init__(self, env: ParallelEnv, lesson: dict):
self.env = env
self.lesson = lesson
def check_winnable(self, lst: list[int], piece: int) -> bool:
"""Checks if four pieces in a row represent a winnable opportunity, e.g. [1, 1, 1, 0] or [2, 0, 2, 2].
:param lst: List of pieces in row
:type lst: List
:param piece: Player piece we are checking (1 or 2)
:type piece: int
"""
return lst.count(piece) == 3 and lst.count(0) == 1
def check_vertical_win(self, player: int) -> bool:
"""Checks if a win is vertical.
:param player: Player who we are checking, 0 or 1
:type player: int
"""
board = np.array(self.env.env.board).reshape(6, 7)
piece = player + 1
column_count = 7
row_count = 6
# Check vertical locations for win
for c in range(column_count):
for r in range(row_count - 3):
if (
board[r][c] == piece
and board[r + 1][c] == piece
and board[r + 2][c] == piece
and board[r + 3][c] == piece
):
return True
return False
def check_three_in_row(self, player: int) -> int:
"""Checks if there are three pieces in a row and a blank space next, or two pieces - blank - piece.
:param player: Player who we are checking, 0 or 1
:type player: int
"""
board = np.array(self.env.env.board).reshape(6, 7)
piece = player + 1
# Check horizontal locations
column_count = 7
row_count = 6
three_in_row_count = 0
# Check vertical locations
for c in range(column_count):
for r in range(row_count - 3):
if self.check_winnable(board[r : r + 4, c].tolist(), piece):
three_in_row_count += 1
# Check horizontal locations
for r in range(row_count):
for c in range(column_count - 3):
if self.check_winnable(board[r, c : c + 4].tolist(), piece):
three_in_row_count += 1
# Check positively sloped diagonals
for c in range(column_count - 3):
for r in range(row_count - 3):
if self.check_winnable(
[
board[r, c],
board[r + 1, c + 1],
board[r + 2, c + 2],
board[r + 3, c + 3],
],
piece,
):
three_in_row_count += 1
# Check negatively sloped diagonals
for c in range(column_count - 3):
for r in range(3, row_count):
if self.check_winnable(
[
board[r, c],
board[r - 1, c + 1],
board[r - 2, c + 2],
board[r - 3, c + 3],
],
piece,
):
three_in_row_count += 1
return three_in_row_count
def reward(self, done: bool, player: int) -> float:
"""Processes and returns reward from environment according to lesson criteria.
:param done: Environment has terminated
:type done: bool
:param player: Player who we are checking, 0 or 1
:type player: int
"""
if done:
reward = (
self.lesson["rewards"]["vertical_win"]
if self.check_vertical_win(player)
else self.lesson["rewards"]["win"]
)
else:
agent_three_count = self.check_three_in_row(1 - player)
opp_three_count = self.check_three_in_row(player)
if (agent_three_count + opp_three_count) == 0:
reward = self.lesson["rewards"]["play_continues"]
else:
reward = (
self.lesson["rewards"]["three_in_row"] * agent_three_count
+ self.lesson["rewards"]["opp_three_in_row"] * opp_three_count
)
return reward
def last(self) -> tuple[dict, float, bool, bool, dict]:
"""Wrapper around PettingZoo env last method."""
return self.env.last()
def step(self, action: int) -> None:
"""Wrapper around PettingZoo env step method."""
self.env.step(action)
def reset(self) -> None:
"""Wrapper around PettingZoo env reset method."""
self.env.reset()
class Opponent:
"""Connect 4 opponent to train and/or evaluate against.
:param env: Environment to learn in
:type env: PettingZoo-style environment
:param difficulty: Difficulty level of opponent, 'random', 'weak' or 'strong'
:type difficulty: str
"""
def __init__(self, env: ParallelEnv, difficulty: str):
self.env = env.env
self.difficulty = difficulty
if self.difficulty == "random":
self.get_action = self.random_opponent
elif self.difficulty == "weak":
self.get_action = self.weak_rule_based_opponent
else:
self.get_action = self.strong_rule_based_opponent
self.num_cols = 7
self.num_rows = 6
self.length = 4
self.top = [0] * self.num_cols
def update_top(self) -> None:
"""Updates self.top, a list which tracks the row on top of the highest piece in each column."""
board = np.array(self.env.env.board).reshape(self.num_rows, self.num_cols)
non_zeros = np.where(board != 0)
rows, cols = non_zeros
top = np.zeros(board.shape[1], dtype=int)
for col in range(board.shape[1]):
column_pieces = rows[cols == col]
if len(column_pieces) > 0:
top[col] = np.min(column_pieces) - 1
else:
top[col] = 5
full_columns = np.all(board != 0, axis=0)
top[full_columns] = 6
self.top = top
def random_opponent(
self,
action_mask: list[int],
last_opp_move: int | None = None,
block_vert_coef: float = 1,
) -> int:
"""Takes move for random opponent. If the lesson aims to randomly block vertical
wins with a higher probability, this is done here too.
:param action_mask: Mask of legal actions: 1=legal, 0=illegal
:type action_mask: List
:param last_opp_move: Most recent action taken by agent against this opponent
:type last_opp_move: int
:param block_vert_coef: How many times more likely to block vertically
:type block_vert_coef: float
"""
if last_opp_move is not None:
action_mask[last_opp_move] *= block_vert_coef
return random.choices(list(range(self.num_cols)), action_mask)[0]
def weak_rule_based_opponent(self, player: int) -> int:
"""Takes move for weak rule-based opponent.
:param player: Player who we are checking, 0 or 1
:type player: int
"""
self.update_top()
max_length = -1
best_actions = []
for action in range(self.num_cols):
possible, reward, ended, lengths = self.outcome(
action,
player,
return_length=True,
)
if possible and lengths.sum() > max_length:
best_actions = []
max_length = lengths.sum()
if possible and lengths.sum() == max_length:
best_actions.append(action)
return random.choice(best_actions)
def strong_rule_based_opponent(self, player: int) -> int:
"""Takes move for strong rule-based opponent.
:param player: Player who we are checking, 0 or 1
:type player: int
"""
self.update_top()
winning_actions = []
for action in range(self.num_cols):
possible, reward, ended = self.outcome(action, player)
if possible and ended:
winning_actions.append(action)
if len(winning_actions) > 0:
return random.choice(winning_actions)
opp = 1 if player == 0 else 0
loss_avoiding_actions = []
for action in range(self.num_cols):
possible, reward, ended = self.outcome(action, opp)
if possible and ended:
loss_avoiding_actions.append(action)
if len(loss_avoiding_actions) > 0:
return random.choice(loss_avoiding_actions)
return self.weak_rule_based_opponent(player) # take best possible move
def outcome(
self,
action: int,
player: int,
return_length: bool = False,
) -> tuple[bool, float | None, bool, np.ndarray | None]:
"""Takes move for weak rule-based opponent.
:param action: Action to take in environment
:type action: int
:param player: Player who we are checking, 0 or 1
:type player: int
:param return_length: Return length of outcomes, defaults to False
:type return_length: bool, optional
"""
if not (self.top[action] < self.num_rows): # action column is full
return (False, None, None) + ((None,) if return_length else ())
row, col = self.top[action], action
piece = player + 1
# down, up, left, right, down-left, up-right, down-right, up-left,
directions = np.array(
[
[[-1, 0], [1, 0]],
[[0, -1], [0, 1]],
[[-1, -1], [1, 1]],
[[-1, 1], [1, -1]],
],
) # |4x2x2|
positions = np.array([row, col]).reshape(1, 1, 1, -1) + np.expand_dims(
directions,
-2,
) * np.arange(1, self.length).reshape(
1,
1,
-1,
1,
) # |4x2x3x2|
valid_positions = np.logical_and(
np.logical_and(
positions[:, :, :, 0] >= 0,
positions[:, :, :, 0] < self.num_rows,
),
np.logical_and(
positions[:, :, :, 1] >= 0,
positions[:, :, :, 1] < self.num_cols,
),
) # |4x2x3|
d0 = np.where(valid_positions, positions[:, :, :, 0], 0)
d1 = np.where(valid_positions, positions[:, :, :, 1], 0)
board = np.array(self.env.env.board).reshape(self.num_rows, self.num_cols)
board_values = np.where(valid_positions, board[d0, d1], 0)
a = (board_values == piece).astype(int)
b = np.concatenate(
(a, np.zeros_like(a[:, :, :1])),
axis=-1,
) # padding with zeros to compute length
lengths = np.argmin(b, -1)
ended = False
# check if winnable in any direction
for both_dir in board_values:
# |2x3|
line = np.concatenate((both_dir[0][::-1], [piece], both_dir[1]))
if "".join(map(str, [piece] * self.length)) in "".join(map(str, line)):
ended = True
break
# ended = np.any(np.greater_equal(np.sum(lengths, 1), self.length - 1))
draw = True
for c, v in enumerate(self.top):
draw &= (v == self.num_rows) if c != col else (v == (self.num_rows - 1))
ended |= draw
reward = (-1) ** (player) if ended and not draw else 0
return (True, reward, ended) + ((lengths,) if return_length else ())
def agent_state(observation: dict) -> np.ndarray:
"""Player-perspective CHW float state from a Connect-Four observation.
PettingZoo returns the observation from the *current* player's point of view,
so a plain channel-move is enough for whoever is to move (no plane swap).
:param observation: Raw PettingZoo observation dict.
:type observation: dict
:return: (channels, height, width) float32 array.
:rtype: numpy.ndarray
"""
return np.moveaxis(observation["observation"], -1, -3).astype(np.float32)
class ConnectFourVecEnv:
"""Vectorized self-play Connect Four as a single-agent MDP.
``num_envs`` games are stepped in lockstep with the opponent embedded in
:meth:`step`, so one batched ``get_action`` drives them all. Buffers are
preallocated and written in place, and terminated games auto-reset. A fixed
``num_envs`` keeps the agent's batch shapes static for CUDA graph capture.
:param num_envs: Number of parallel games.
:type num_envs: int
:param lesson: Curriculum lesson settings (opponent, rewards, ...).
:type lesson: dict
:param opponent_policy: Frozen agent used as the opponent for self-play
(``lesson['opponent'] == 'self'``); ``None`` for rule-based opponents.
:type opponent_policy: DQN | None
"""
def __init__(self, num_envs: int, lesson: dict, opponent_policy=None):
self.num_envs = num_envs
self.lesson = lesson
self.opponent_policy = opponent_policy
raw = connect_four_v3.env().observation_space("player_0")["observation"]
self.single_observation_space = gym.spaces.Box(
low=raw.low.transpose(2, 0, 1),
high=raw.high.transpose(2, 0, 1),
dtype=np.float32,
)
self.single_action_space = connect_four_v3.env().action_space("player_0")
self.observations = np.zeros(
(num_envs, *self.single_observation_space.shape), dtype=np.float32
)
self.masks = np.ones((num_envs, 7), dtype=np.int8)
self.rewards = np.zeros(num_envs, dtype=np.float32)
self.terminals = np.zeros(num_envs, dtype=bool)
self.games = [
CurriculumEnv(connect_four_v3.env(), lesson) for _ in range(num_envs)
]
self.rule_opponents = (
[Opponent(g, difficulty=lesson["opponent"]) for g in self.games]
if lesson["opponent"] != "self"
else None
)
def transition(
self, prev_observations: np.ndarray, actions: np.ndarray
) -> TensorDict:
"""Build one batched replay transition from the pre-step obs and current buffers."""
n = self.num_envs
return TensorDict(
{
"obs": torch.from_numpy(prev_observations.copy()),
"action": torch.from_numpy(np.asarray(actions, dtype=np.int64)).reshape(
n, 1
),
"reward": torch.from_numpy(self.rewards.copy()).reshape(n, 1),
"next_obs": torch.from_numpy(self.observations.copy()),
"done": torch.from_numpy(self.terminals.astype(np.float32)).reshape(
n, 1
),
},
batch_size=[n],
)
def _write_obs(self, i: int) -> None:
obs, _, _, _, _ = self.games[i].last()
self.observations[i] = agent_state(obs)
self.masks[i] = obs["action_mask"]
def _reset_game(self, i: int) -> None:
self.games[i].reset()
self._write_obs(i)
def reset(self, seed: int | None = None) -> None:
"""Reset every game and fill the observation/mask buffers in place."""
for i in range(self.num_envs):
self._reset_game(i)
self.terminals[:] = False
def _opponent_actions(self, pending: list[int]) -> list[int]:
"""Opponent (player_1) moves for the games in ``pending``."""
if not pending:
return []
if self.opponent_policy is not None:
# self-play: one batched greedy forward through the frozen opponent net
obs = np.stack([agent_state(self.games[i].last()[0]) for i in pending])
masks = np.stack([self.games[i].last()[0]["action_mask"] for i in pending])
return list(
self.opponent_policy.get_action(obs, epsilon=0.0, action_mask=masks)
)
# 'random' takes the action mask; 'weak'/'strong' introspect the board.
if self.lesson["opponent"] == "random":
block = self.lesson.get("block_vert_coef", 1)
return [
self.rule_opponents[i].get_action(
self.games[i].last()[0]["action_mask"], None, block
)
for i in pending
]
return [self.rule_opponents[i].get_action(player=1) for i in pending]
def step(self, actions: np.ndarray) -> None:
"""Apply agent actions, play the opponent's reply, write buffers, auto-reset."""
pending = []
for i, g in enumerate(self.games):
g.step(int(actions[i]))
_, _, done, trunc, _ = g.last()
if done or trunc:
self.rewards[i] = g.reward(done=True, player=0)
self.terminals[i] = True
self._reset_game(i)
else:
pending.append(i)
opp_actions = self._opponent_actions(pending)
for j, i in enumerate(pending):
g = self.games[i]
g.step(int(opp_actions[j]))
_, _, done, trunc, _ = g.last()
if done or trunc:
self.rewards[i] = self.lesson["rewards"]["lose"]
self.terminals[i] = True
self._reset_game(i)
else:
self.rewards[i] = g.reward(done=False, player=0)
self.terminals[i] = False
self._write_obs(i)
@torch.no_grad()
def evaluate(agent, lesson, num_envs: int, n_games: int = 192) -> float:
"""Greedy win-rate vs the eval opponent, played on ``num_envs`` parallel games.
Uses the same ``num_envs`` as training so the CUDA-graph-captured
``get_action`` always sees the same static batch shape.
"""
eval_lesson = dict(lesson, opponent=lesson["eval_opponent"])
venv = ConnectFourVecEnv(num_envs, eval_lesson)
venv.reset()
wins = done = 0
while done < n_games:
actions = agent.get_action(
venv.observations, epsilon=0.0, action_mask=venv.masks
)
venv.step(np.asarray(actions))
for i in range(num_envs):
if venv.terminals[i]:
wins += venv.rewards[i] >= lesson["rewards"]["win"] * 0.9
done += 1
return wins / max(done, 1)
if __name__ == "__main__":
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print("===== AgileRL Curriculum Learning Demo (vectorized + CUDA graphs) =====")
# Fixed so the agent's batch shapes stay static for CUDA graph capture.
NUM_ENVS = 32
# Replays a step's kernel launches as one call, removing the CPU dispatch
# overhead that dominates small networks. Needs static shapes.
USE_CUDAGRAPHS = device.type == "cuda"
for lesson_number in range(1, 5):
with open(f"./curriculums/connect_four/lesson{lesson_number}.yaml") as file:
LESSON = yaml.safe_load(file)
net_config = {
"encoder_config": {
"channel_size": [128],
"kernel_size": [4],
"stride_size": [1],
},
"head_config": {"hidden_size": [64, 64]},
}
init_hp = {
"double": True,
"batch_size": 256, # fixed (not mutated) so the CUDA graph stays valid
"lr": 1e-4,
"gamma": 0.99,
"learn_step": 1,
"tau": 0.01,
}
population_size = 6
memory_size = 10000
probe = ConnectFourVecEnv(NUM_ENVS, LESSON)
observation_space = probe.single_observation_space
action_space = probe.single_action_space
# batch_size is not mutated: it would invalidate the captured CUDA graph.
hp_config = HyperparameterConfig(
lr=RLParameter(min=1e-4, max=1e-2),
learn_step=RLParameter(
min=1, max=120, dtype=int, grow_factor=1.5, shrink_factor=0.75
),
)
pop = DQN.population(
size=population_size,
observation_space=observation_space,
action_space=action_space,
net_config=net_config,
hp_config=hp_config,
device=device,
cudagraphs=USE_CUDAGRAPHS,
**init_hp,
)
memory = ReplayBuffer(max_size=memory_size, device=device)
tournament = TournamentSelection(
tournament_size=2, elitism=True, population_size=population_size
)
mutations = Mutations(
no_mutation=0.2,
architecture=0,
new_layer_prob=0.2,
parameters=0.2,
activation=0,
rl_hp=0.2,
mutation_sd=0.1,
rand_seed=1,
device=device,
)
# Training-loop parameters
max_episodes = LESSON["max_train_episodes"]
evo_epochs = 5 # evolve every N vectorized rollout blocks
block_steps = 100 # batched steps per agent per block
epsilon, eps_end, eps_decay = 1.0, 0.1, 0.9995
if LESSON["pretrained_path"] is not None:
for agent in pop:
agent.load_checkpoint(LESSON["pretrained_path"])
opponent_pool = None
if LESSON["opponent"] == "self":
opponent_pool = deque(maxlen=LESSON["opponent_pool_size"])
for _ in range(LESSON["opponent_pool_size"]):
opponent_pool.append(pop[0].clone())
# Buffer + agent warm-up (vectorized random rollout)
if LESSON["buffer_warm_up"]:
warm_lesson = dict(LESSON, opponent=LESSON["warm_up_opponent"])
warm_env = ConnectFourVecEnv(NUM_ENVS, warm_lesson)
warm_env.reset()
print("Filling replay buffer ...")
while len(memory) < memory.max_size:
acts = np.array(
[
random.choices(range(7), warm_env.masks[i])[0]
for i in range(NUM_ENVS)
]
)
prev = warm_env.observations.copy()
warm_env.step(acts)
memory.add(warm_env.transition(prev, acts))
if LESSON["agent_warm_up"] > 0:
print("Warming up agents ...")
for agent in pop:
for _ in range(LESSON["agent_warm_up"]):
agent.learn(memory.sample(agent.batch_size))
if max_episodes > 0:
wandb.init(
project="AgileRL",
name="{}-EvoHPO-{}-{}Opposition-CNN-{}".format(
"connect_four_v3",
"DQN",
LESSON["opponent"],
datetime.now(tz=timezone.utc).strftime("%m%d%Y%H%M%S"),
),
config={
"algo": "Evo HPO DQN",
"env": "connect_four_v3",
"init_hp": init_hp,
"lesson": LESSON,
},
)
elite = pop[0]
total_steps = 0
# Each step collects NUM_ENVS transitions.
n_blocks = 0 if max_episodes == 0 else 25
pbar = trange(n_blocks)
for block in pbar:
for agent in pop:
opp = (
random.choice(opponent_pool) if opponent_pool is not None else None
)
venv = ConnectFourVecEnv(NUM_ENVS, LESSON, opponent_policy=opp)
venv.reset()
for agent_step in range(block_steps):
prev = venv.observations.copy()
actions = agent.get_action(venv.observations, epsilon, venv.masks)
venv.step(np.asarray(actions))
memory.add(venv.transition(prev, np.asarray(actions)))
if (
len(memory) >= agent.batch_size
and agent_step % agent.learn_step == 0
):
agent.learn(memory.sample(agent.batch_size))
total_steps += NUM_ENVS
epsilon = max(eps_end, epsilon * eps_decay)
# Self-play: refresh the opponent pool with the current elite
if opponent_pool is not None and (block + 1) % evo_epochs == 0:
elite_opp, _, _ = tournament._elitism(pop)
opponent_pool.append(elite_opp.clone())
if (block + 1) % evo_epochs == 0:
fitnesses = [evaluate(agent, LESSON, NUM_ENVS) for agent in pop]
for agent, fit in zip(pop, fitnesses):
agent.metrics.add_fitness(fit)
elite, pop, _ = tournament.select(pop)
pop = mutations.mutation(pop)
pbar.set_postfix_str(
f"Lesson {lesson_number} Eval win-rate (best): {max(fitnesses):.2f} "
f"Total steps: {total_steps}"
)
if max_episodes > 0:
wandb.finish()
save_path = LESSON["save_path"]
os.makedirs(os.path.dirname(save_path), exist_ok=True)
elite.save_checkpoint(save_path)
print(f"Elite agent saved to '{save_path}'.")