4️⃣ Training
Live training visualisation (optional)
To watch training, rl_utils.log_grid_video(obs, env.draw, dones=dones, live=...) renders the first 16
environments of a rollout as a 4x4 grid (each environment draws its own cell via its draw method - this is
the same renderer used by DQN and PPO), encodes the frames as a single
autoplaying/looping MP4 (one ffmpeg encode, ~0.2s), logs it to wandb if a run is active, and displays it inline
if live. The training loop below calls it every video_log_freq gradient steps.
Trainer
This is the function that will handle the full training loop. We've provided you with the template of a training loop which should be very similar to the DQN section's.
Exercise - implement VPGTrainer
Difficulty:
🔴🔴🔴🔴🔴
Importance:
🔵🔵🔵🔵🔵
You should spend up to 45 minutes on this exercise.
You should fill in the following methods. Ignore logging, can just copy from the solution later.
compute_loss- this method should compute the loss for the VPG objective function.
The training loop is rather standard once everything else is done: we do a rollout, we cut the result into batches, compute the loss, and update the weights from each batch, so we've provided it for you.
class VPGTrainer:
def __init__(self, args: VPGArgs):
set_global_seeds(args.seed)
self.args = args
device = args.device
self.run_name = f"{args.env_id}__{args.wandb_project_name}__seed{args.seed}__{time.strftime('%Y%m%d-%H%M%S')}"
# Create our environments (any `env_id` registered in `ENVS`, see the setup cell)
self.envs = make_envs(args.env_id, args.num_envs, args.seed, device)
# Define some basic variables from our environment (note, we assume a single discrete action space)
self.num_envs = args.num_envs
self.action_shape = self.envs.action_space.shape
self.num_actions = self.envs.action_space.n
self.obs_shape = self.envs.observation_space.shape
# Create our networks & optimizer
self.policy_network = PolicyNetwork(self.obs_shape, self.num_actions).to(device)
self.optimizer = t.optim.Adam(self.policy_network.parameters(), lr=args.lr, eps=1e-5, maximize=True)
self.optimizer.zero_grad()
# Create our agent
self.agent = VPGAgent(envs=self.envs, policy_network=self.policy_network, args=self.args)
def compute_loss(self, tau: RolloutTensors) -> tuple[t.Tensor, dict[str, Any]]:
raise NotImplementedError()
info = {
"r_joy": joy.item(),
"iw": iw.mean().item() if self.args.use_iw else None,
}
return joy, info
def update_learning_rate(self, time_steps, args):
if args.use_lr_decay and args.lr_frac > 0:
progress = min(1.0, max(time_steps / args.total_timesteps, 0) / args.lr_frac)
return (progress * args.lr_end) + ((1 - progress) * args.lr)
return args.lr
def train(self) -> None:
"""
Trains the agent by generating rollouts and updating the policy.
The progress bar tracks total environment steps.
"""
if self.args.use_wandb:
wandb.init(
project=self.args.wandb_project_name,
entity=self.args.wandb_entity,
name=self.run_name,
)
wandb.watch(self.policy_network, log="all", log_freq=50)
# --- Setup ---
rollout = Rollout(max_steps=self.args.num_steps_per_rollout)
# Calculate the total number of rollouts to perform
env_steps_per_rollout = self.args.num_steps_per_rollout * self.args.num_envs
num_updates = self.args.total_timesteps // env_steps_per_rollout
train_steps = 0 # Counter for gradient updates
next_video_at = 0 # gradient-step count at which the next grid video is due
video_slot = LiveVideo(overwrite=self.args.overwrite_video) # one display slot per run (see rl_utils.LiveVideo)
# --- Training Loop ---
# The progress bar is managed manually with a `with` statement.
# `total` is set to the total environment steps we want to run.
# The loop iterates `num_updates` times, not `total_timesteps` times.
with tqdm(
total=self.args.total_timesteps,
unit=" env steps",
unit_scale=True,
desc="Training",
miniters=1,
mininterval=0.02,
) as pbar:
env_steps_consumed = 0
for update_num in range(num_updates):
# 1. Generate a new rollout from the environment
rollout = self.agent.gen_rollout(rollout)
# 2. Split the rollout into batches along the num_envs dimension
rollout_batches = rollout.get_batches(self.args.batch_size)
# 3. Logging and Progress Bar Update
# Undiscounted return of each env's *first* episode in this rollout
tau = rollout.get()
first_episode = (tau.dones.int().cumsum(dim=1) - tau.dones.int()) == 0 # no `done` strictly before this step
episode_return = (tau.rewards * first_episode).sum(dim=1) # (num_envs,)
mean_return = episode_return.mean().item()
# 95% confidence interval of average of return on first episode
ci95_return = 1.96 * (episode_return.std() / self.args.num_envs**0.5).item()
max_return = episode_return.max().item()
# Log a 4x4 grid video of the rollout we already have (no extra env steps) every `video_log_freq`
# gradient steps - to wandb if use_wandb, inline if live_viz - via the shared `rl_utils.log_grid_video`.
# With rollout_use_count=1 and one batch per rollout that is one video every `video_log_freq` rollouts.
want_video = self.args.video_log_freq and (self.args.use_wandb or self.args.live_viz)
if want_video and train_steps >= next_video_at:
caption = f"rollout {rollout.timestep} steps | mean first-episode return {mean_return:.1f}"
log_grid_video(tau.obs, self.envs.draw, dones=tau.dones, step=env_steps_consumed, live=self.args.live_viz, slot=video_slot, caption=caption)
next_video_at += self.args.video_log_freq
# 4. Advance env-step counter before gradient updates (one rollout collected)
env_steps_consumed += self.args.num_steps_per_rollout * self.args.num_envs
# 5. For each batch, perform multiple gradient updates
for i in range(self.args.rollout_use_count):
for batch in rollout_batches:
loss, reinforce_info = self.compute_loss(batch)
info = reinforce_info
loss.backward()
self.optimizer.step()
self.optimizer.zero_grad()
train_steps += 1
new_lr = self.update_learning_rate(env_steps_consumed, self.args)
for pg in self.optimizer.param_groups:
pg["lr"] = new_lr
# Create info string to display in the progress bar
current_lr = self.optimizer.param_groups[0]["lr"]
info_dict = {
"joy": f"{info['r_joy']:.4f}",
"G": f"{mean_return:.2f} ± {ci95_return:.2f} (max: {max_return:.0f})",
"iw": f"{info['iw']:.4f}" if self.args.use_iw else None,
"lr": f"{current_lr:.2e}",
}
pbar.set_postfix(info_dict)
# Progress bar advances once per rollout (env steps actually collected)
pbar.update(self.args.num_steps_per_rollout * self.args.num_envs)
# --- Cleanup ---
self.envs.close()
if self.args.use_wandb:
wandb.finish()
tests.test_compute_loss(VPGTrainer, VPGArgs, Rollout)
Solution
class VPGTrainer:
def __init__(self, args: VPGArgs):
set_global_seeds(args.seed)
self.args = args
device = args.device
self.run_name = f"{args.env_id}__{args.wandb_project_name}__seed{args.seed}__{time.strftime('%Y%m%d-%H%M%S')}"
# Create our environments (any `env_id` registered in `ENVS`, see the setup cell)
self.envs = make_envs(args.env_id, args.num_envs, args.seed, device)
# Define some basic variables from our environment (note, we assume a single discrete action space)
self.num_envs = args.num_envs
self.action_shape = self.envs.action_space.shape
self.num_actions = self.envs.action_space.n
self.obs_shape = self.envs.observation_space.shape
# Create our networks & optimizer
self.policy_network = PolicyNetwork(self.obs_shape, self.num_actions).to(device)
self.optimizer = t.optim.Adam(self.policy_network.parameters(), lr=args.lr, eps=1e-5, maximize=True)
self.optimizer.zero_grad()
# Create our agent
self.agent = VPGAgent(envs=self.envs, policy_network=self.policy_network, args=self.args)
def compute_loss(self, tau: RolloutTensors) -> tuple[t.Tensor, dict[str, Any]]:
returns = compute_returns(tau.rewards, tau.dones, self.args.gamma) # (num_envs, timestep)
if self.args.normalize_returns:
returns = normalize_returns(returns)
logprobs_taken = compute_logprobs(tau, self.policy_network)
iw = compute_importance_weights(logprobs_taken, tau, self.args.clip_coef) if self.args.use_iw else t.ones_like(logprobs_taken)
joy = compute_reinforce_loss(returns, logprobs_taken, iw)
info = {
"r_joy": joy.item(),
"iw": iw.mean().item() if self.args.use_iw else None,
}
return joy, info
def update_learning_rate(self, time_steps, args):
if args.use_lr_decay and args.lr_frac > 0:
progress = min(1.0, max(time_steps / args.total_timesteps, 0) / args.lr_frac)
return (progress * args.lr_end) + ((1 - progress) * args.lr)
return args.lr
def train(self) -> None:
"""
Trains the agent by generating rollouts and updating the policy.
The progress bar tracks total environment steps.
"""
if self.args.use_wandb:
wandb.init(
project=self.args.wandb_project_name,
entity=self.args.wandb_entity,
name=self.run_name,
)
wandb.watch(self.policy_network, log="all", log_freq=50)
# --- Setup ---
rollout = Rollout(max_steps=self.args.num_steps_per_rollout)
# Calculate the total number of rollouts to perform
env_steps_per_rollout = self.args.num_steps_per_rollout * self.args.num_envs
num_updates = self.args.total_timesteps // env_steps_per_rollout
train_steps = 0 # Counter for gradient updates
next_video_at = 0 # gradient-step count at which the next grid video is due
video_slot = LiveVideo(overwrite=self.args.overwrite_video) # one display slot per run (see rl_utils.LiveVideo)
# --- Training Loop ---
# The progress bar is managed manually with a `with` statement.
# `total` is set to the total environment steps we want to run.
# The loop iterates `num_updates` times, not `total_timesteps` times.
with tqdm(
total=self.args.total_timesteps,
unit=" env steps",
unit_scale=True,
desc="Training",
miniters=1,
mininterval=0.02,
) as pbar:
env_steps_consumed = 0
for update_num in range(num_updates):
# 1. Generate a new rollout from the environment
rollout = self.agent.gen_rollout(rollout)
# 2. Split the rollout into batches along the num_envs dimension
rollout_batches = rollout.get_batches(self.args.batch_size)
# 3. Logging and Progress Bar Update
# Undiscounted return of each env's *first* episode in this rollout
tau = rollout.get()
first_episode = (tau.dones.int().cumsum(dim=1) - tau.dones.int()) == 0 # no `done` strictly before this step
episode_return = (tau.rewards * first_episode).sum(dim=1) # (num_envs,)
mean_return = episode_return.mean().item()
# 95% confidence interval of average of return on first episode
ci95_return = 1.96 * (episode_return.std() / self.args.num_envs**0.5).item()
max_return = episode_return.max().item()
# Log a 4x4 grid video of the rollout we already have (no extra env steps) every `video_log_freq`
# gradient steps - to wandb if use_wandb, inline if live_viz - via the shared `rl_utils.log_grid_video`.
# With rollout_use_count=1 and one batch per rollout that is one video every `video_log_freq` rollouts.
want_video = self.args.video_log_freq and (self.args.use_wandb or self.args.live_viz)
if want_video and train_steps >= next_video_at:
caption = f"rollout {rollout.timestep} steps | mean first-episode return {mean_return:.1f}"
log_grid_video(tau.obs, self.envs.draw, dones=tau.dones, step=env_steps_consumed, live=self.args.live_viz, slot=video_slot, caption=caption)
next_video_at += self.args.video_log_freq
# 4. Advance env-step counter before gradient updates (one rollout collected)
env_steps_consumed += self.args.num_steps_per_rollout * self.args.num_envs
# 5. For each batch, perform multiple gradient updates
for i in range(self.args.rollout_use_count):
for batch in rollout_batches:
loss, reinforce_info = self.compute_loss(batch)
info = reinforce_info
loss.backward()
self.optimizer.step()
self.optimizer.zero_grad()
train_steps += 1
new_lr = self.update_learning_rate(env_steps_consumed, self.args)
for pg in self.optimizer.param_groups:
pg["lr"] = new_lr
# Create info string to display in the progress bar
current_lr = self.optimizer.param_groups[0]["lr"]
info_dict = {
"joy": f"{info['r_joy']:.4f}",
"G": f"{mean_return:.2f} ± {ci95_return:.2f} (max: {max_return:.0f})",
"iw": f"{info['iw']:.4f}" if self.args.use_iw else None,
"lr": f"{current_lr:.2e}",
}
pbar.set_postfix(info_dict)
# Progress bar advances once per rollout (env steps actually collected)
pbar.update(self.args.num_steps_per_rollout * self.args.num_envs)
# --- Cleanup ---
self.envs.close()
if self.args.use_wandb:
wandb.finish()
tests.test_compute_loss(VPGTrainer, VPGArgs, Rollout)
Probes
As in the DQN section, we will be using probes to test our model. They've been implemented for you.
def test_probe(probe_idx: int):
"""
Tests a probe environment by training a network on it & verifying that the value functions are
in the expected range.
"""
# Train our network
args = VPGArgs(
env_id=f"Probe{probe_idx}-v0",
wandb_project_name=f"test-probe-{probe_idx}",
total_timesteps=[None, None, None, 2500, 5000][probe_idx - 1],
num_steps_per_rollout=16, # one-step episodes: short rollouts = more (cheap) updates per env step
lr=5e-3,
num_envs=4,
video_log_freq=None,
use_wandb=False,
device="cpu",
clip_coef=None,
normalize_returns=False,
rollout_use_count=1,
)
trainer = VPGTrainer(args)
trainer.train()
agent = trainer.agent
# Get the correct set of observations, and corresponding values we expect
obs_for_probes = [[[0.0]], [[-1.0], [+1.0]], [[0.0], [1.0]], [[0.0]], [[0.0], [1.0]]]
expected_probs_for_probes = [None, None, None, [[0.0, 1.0]], [[1.0, 0.0], [0.0, 1.0]]]
tolerances = [1e-3, 1e-3, 1e-3, 2e-3, 2e-3]
obs = t.tensor(obs_for_probes[probe_idx - 1]).to(args.device)
# Calculate the actual value & probs, and verify them
with t.inference_mode():
probs = agent.policy_network(obs).softmax(-1)
expected_probs = expected_probs_for_probes[probe_idx - 1]
if expected_probs is not None:
print(f"Probs: {probs}")
print(f"Expected probs: {t.tensor(expected_probs).to(args.device)}")
t.testing.assert_close(probs, t.tensor(expected_probs).to(args.device), atol=tolerances[probe_idx - 1], rtol=0)
print(f"Probe {probe_idx} tests passed!\n")
for probe_idx in [4, 5]:
test_probe(probe_idx)
Training Run
Vanilla Policy Gradient can often be a bit finicky and unstable to train (which is why in practice we use PPO instead).
None-the-less, I've tried to find a good set of hyperparameters such that it trains to an optimal policy in around 10 seconds on a single GPU!
By default (live_viz=True) a 4x4 grid video of the agent is rendered in the notebook every video_log_freq gradient steps as it trains; pass live_viz=False to turn that off (videos are only sent to wandb if use_wandb=True). The background of a cell flashes pink when that environment's agent dies and is reset, so you can easily see when a new episode starts.
args_fast = VPGArgs(
use_wandb=False,
num_envs=1024,
num_batches_per_rollout=1,
total_timesteps=15_000_000,
num_steps_per_rollout=500,
rollout_use_count=1, # this seems to matter a lot
normalize_returns=True,
lr=2e-2, # very high, but with 1024 envs the gradient is low-variance enough to take it
use_lr_decay=False,
use_iw=False, # don't need it if we only use each rollout once
gamma=0.99,
seed=1337,
device=device,
video_log_freq=10,
)
trainer = VPGTrainer(args_fast)
trainer.train()