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.
What's with the t.set_num_threads line in the setup code?
Note that if you are training on the CPU (which for a network this small can often be faster than the GPU!), having too many threads can slow training as the overhead between splitting the work over many threads can hurt. As such, we've added the line
t.set_num_threads(min(4, t.get_num_threads()))
to the setup code, which caps the number of CPU threads PyTorch uses at 4 (and does nothing if you have fewer). Feel free to adjust it or remove it entirely. On a beefy machine with dozens of threads, I found a x10 slowdown if the number of threads were not capped.
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)}")
# A bf16 network (e.g. under autocast) can't represent these targets to better than its eps
atol = max(tolerances[probe_idx - 1], t.finfo(probs.dtype).eps)
t.testing.assert_close(probs.float(), t.tensor(expected_probs).to(args.device), atol=atol, 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 CPU or a GPU. (This relies on the t.set_num_threads line in the setup cell if you're training on a CPU with too many cores: PyTorch's default of one thread per core is counterproductive for a network this small.)
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=32,
num_batches_per_rollout=1,
total_timesteps=640_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 the per-timestep baseline across envs keeps the gradient 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()