3️⃣ Training Loop

class PPOTrainer:
    def __init__(self, args: PPOArgs):
        set_global_seeds(args.seed)
        self.args = args
        self.run_name = f"{args.env_id}__{args.wandb_project_name}__seed{args.seed}__{time.strftime('%Y%m%d-%H%M%S')}"
        # Accelerated vectorised env, chosen by mode. All three expose the same gym-style
        # reset()/step() returning GPU tensors, so PPO never leaves the GPU:
        #   classic-control -> GPU CartPole;  atari -> EnvPool (C++ emulators);  mujoco -> Brax (GPU physics).
        # (AtariEnvs / BraxEnvs are defined in the Atari / MuJoCo bonus sections.)
        self.envs = ENV_DICT[args.mode](args.env_id, args.num_envs, seed=args.seed)

        # Define some basic variables from our environment
        self.num_envs = self.envs.num_envs
        self.action_shape = self.envs.single_action_space.shape
        self.obs_shape = self.envs.single_observation_space.shape

        # Create our replay memory
        self.memory = ReplayMemory(
            self.num_envs,
            self.obs_shape,
            self.action_shape,
            args.batch_size,
            args.minibatch_size,
            args.batches_per_learning_phase,
            args.seed,
        )

        # Create our networks & optimizer
        self.actor, self.critic = get_actor_and_critic(self.envs, mode=args.mode)
        self.optimizer, self.scheduler = make_optimizer(self.actor, self.critic, args.total_training_steps, args.lr)

        # Create our agent
        self.agent = PPOAgent(self.envs, self.actor, self.critic, self.memory)

    def rollout_phase(self) -> dict | None:
        """
        This function populates the memory with a new set of experiences, using self.agent.play_step
        to step through the environment. It also returns a dict of data which you can include in
        your progress bar postfix.
        """
        raise NotImplementedError()

    def learning_phase(self) -> None:
        """
        This function does the following:
            - Generates minibatches from memory
            - Calculates the objective function, and takes an optimization step based on it
            - Clips the gradients (see detail #11)
            - Steps the learning rate scheduler
        """
        raise NotImplementedError()

    def compute_ppo_objective(self, minibatch: ReplayMinibatch) -> Float[Tensor, ""]:
        """
        Handles learning phase for a single minibatch. Returns objective function to be maximized.
        """
        raise NotImplementedError()

    def log_video(self, phase: int) -> None:
        """Render the first 16 envs of the rollout currently sitting in the replay memory as a 4x4
        grid video (drawn with the env's own `draw` method), save it as an HTML <video> under
        `args.video_save_path / run_name`, and log it to wandb if enabled. This is what
        `video_log_freq` does. It reuses the rollout we just collected, so it costs no extra env
        steps; modes whose draw() can't render from observations alone (mujoco) are skipped — use
        `record_brax_video` after training instead."""
        if self.args.mode == "mujoco" or not callable(getattr(self.envs, "draw", None)):
            return
        try:
            obs = t.stack([o[:16] for o in self.memory.obs], dim=1).cpu()           # (16, T, *obs_shape)
            dones = t.stack([d[:16] for d in self.memory.terminated], dim=1).cpu()  # (16, T)
            if self.args.mode == "pendulum":  # pendulum's draw() wants the applied torque appended
                actions = t.stack([a[:16] for a in self.memory.actions], dim=1).cpu()
                obs = t.cat([obs, actions.reshape(*obs.shape[:2], -1)], dim=-1)
            cell_w, cell_h = (84, 84) if self.args.mode == "atari" else (160, 120)
            video = render_rollout_grid_html(obs, self.envs.draw, dones=dones, cell_w=cell_w, cell_h=cell_h)
            video_dir = self.args.video_save_path / self.run_name
            video_dir.mkdir(parents=True, exist_ok=True)
            (video_dir / f"phase{phase:04d}.html").write_text(video.data)
            if self.args.use_wandb:
                wandb.log({"rollout_video": wandb.Html(video.data)}, step=self.agent.step)
        except Exception as e:  # never let visualization break training
            print(f"[video log skipped: {e}]")

    def train(self) -> None:
        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.actor, self.critic], log="all", log_freq=50)

        pbar = tqdm(range(self.args.total_phases), desc=f"training {self.args.mode}")
        last_logged_time = time.time()  # so we don't update the progress bar too much
        data = {}

        for phase in pbar:
            new_data = self.rollout_phase()
            if new_data is not None:
                data = new_data
            # Periodically render the rollout we just collected as a grid video (must happen here,
            # before the learning phase consumes & resets the memory).
            if self.args.video_log_freq and (phase % self.args.video_log_freq == 0):
                self.log_video(phase)
            self.learning_phase()
            # Show episode stats (when available) alongside the latest loss components / KL.
            if time.time() - last_logged_time > 0.5:
                last_logged_time = time.time()
                pbar.set_postfix(phase=phase, **data, **getattr(self, "last_metrics", {}))

        self.envs.close()
        if self.args.use_wandb:
            # Remove the watch() forward hooks BEFORE finishing the run: they log to `wandb.run`,
            # which becomes None after finish(), so any later forward pass through the networks
            # (e.g. rendering a video of the trained agent) would crash with
            # "'NoneType' object has no attribute '_log'".
            wandb.unwatch((self.actor, self.critic))
            wandb.finish()

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 = PPOArgs(
        env_id=f"Probe{probe_idx}-v0",
        mode="probe",  # use the GPU (tensor-native) probe envs, so PPOTrainer's GPU agent can run them
        wandb_project_name=f"test-probe-{probe_idx}",
        total_timesteps=30_000, #adjust up if needed to make tests pass
        num_envs=256,
        num_steps_per_rollout=8,
        num_minibatches=4,
        lr=0.01, # also maybe adjust this 
        video_log_freq=None,
        use_wandb=False,
    )
    trainer = PPOTrainer(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_value_for_probes = [
        [[1.0]],
        [[-1.0], [+1.0]],
        [[args.gamma], [1.0]],
        [[1.0]],
        [[1.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(device)

    # Calculate the actual value & probs, and verify them
    with t.inference_mode():
        value = agent.critic(obs)
        probs = agent.actor(obs).softmax(-1)
    expected_value = t.tensor(expected_value_for_probes[probe_idx - 1]).to(device)
    t.testing.assert_close(value, expected_value, atol=tolerances[probe_idx - 1], rtol=0)
    expected_probs = expected_probs_for_probes[probe_idx - 1]
    if expected_probs is not None:
        t.testing.assert_close(probs, t.tensor(expected_probs).to(device), atol=tolerances[probe_idx - 1], rtol=0)
    print("Probe tests passed!\n")


for probe_idx in range(1, 6):
    test_probe(probe_idx)

args = PPOArgs(use_wandb=True, video_log_freq=50)
trainer = PPOTrainer(args)
trainer.train()
display(record_grid_video(trainer, kind="classic-control"))

def cartpole_reward_function(self, action: Tensor) -> Tensor:
    """Shaped reward for CartPole. Called by `step` just after the physics update, so `self.state`
    is the new (num_envs, 4) state tensor with columns [x, v, theta, omega]. Should return a
    (num_envs,) reward tensor (ideally bounded in [0, 1], for easier comparison with the unshaped
    env)."""
    x, v, theta, omega = self.state.unbind(-1)  # each (num_envs,): position, velocity, angle, angular velocity
    raise NotImplementedError()


class EasyCart(CartPole):
    def reward_function(self, action):
        return cartpole_reward_function(self, action)

# Swap the shaped env into ENV_DICT so `mode="classic-control"` builds it (we restore the
# unshaped CartPole afterwards). `env_id` is just a label for the run name.
ENV_DICT["classic-control"] = EasyCart
args = PPOArgs(env_id="EasyCart", use_wandb=True, video_log_freq=50)
trainer = PPOTrainer(args)
trainer.train()
ENV_DICT["classic-control"] = CartPole
display(record_grid_video(trainer, kind="classic-control"))

def spin_cart_reward_function(self, action: Tensor) -> Tensor:
    """Reward for the spinning-cart task. Like `cartpole_reward_function`, this is called by `step`
    after the physics update (`self.state` is the new (num_envs, 4) tensor [x, v, theta, omega])
    and should return a (num_envs,) reward tensor."""
    x, v, theta, omega = self.state.unbind(-1)  # each (num_envs,): position, velocity, angle, angular velocity
    raise NotImplementedError()


class SpinCart(CartPole):
    def reward_function(self, action):
        return spin_cart_reward_function(self, action)

    def terminated(self):
        # Allow full 360-degree rotation: unlike the parent class, only terminate when the cart
        # leaves the track (the pole angle no longer ends the episode).
        x = self.state[:, 0]
        return x.abs() > self.x_threshold

ENV_DICT["classic-control"] = SpinCart
args = PPOArgs(env_id="SpinCart", use_wandb=True, video_log_freq=50)
trainer = PPOTrainer(args)
trainer.train()
ENV_DICT["classic-control"] = CartPole  # restore the unshaped env
display(record_grid_video(trainer, kind="classic-control"))