"""Evaluating agents.
During training, we periodically evaluate the current actor on a separate
set of environments (using its *deterministic* behavior, `actor.act`), and
we keep a copy of the **best actor seen so far**. The `Evaluator` handles
this: call `run_if_needed(steps, actor)` inside the training loop; it does
nothing until `every` new steps have been performed since the last
evaluation.
All the evaluations are kept in `history` (for plots and statistical
tests), and `visualize_best()` plays an episode of the best actor (in a
notebook widget, or as a video file when run as a script).
"""
from rl_mind.core import ActionT, Actor
from rl_mind.env import VecEnv
from rl_mind.notebook import is_notebook, video_display
from rl_mind.rendering import Trimmer, animate_frames, rollout_frames
import copy
from dataclasses import dataclass
from pathlib import Path
from typing import Generic
import gymnasium as gym
import numpy as np
import torch
from torch import Tensor
from torch.utils.tensorboard import SummaryWriter
__all__ = ["EvalResult", "Evaluator", "actor_policy", "record_video"]
[docs]
@dataclass(frozen=True)
class EvalResult:
"""The result of one evaluation
:param step: The number of training steps when the evaluation was run
:param rewards: The cumulated reward of each evaluation episode, shape
`[num_envs]`
:param is_best: True if this evaluation is the best one so far
"""
step: int
rewards: Tensor
is_best: bool
@property
def mean(self) -> float:
"""The mean cumulated reward"""
return float(self.rewards.mean())
[docs]
class Evaluator(Generic[ActionT]):
"""Periodically evaluates an actor, keeping track of the best one
The evaluation environments are deliberately *separate* from the training
ones (their own seed, their own episode count, no same-step reset), so
that what is measured does not depend on the data the agent is training
on. Evaluation uses `Actor.act`, the deterministic behavior.
"""
def __init__(
self,
env: VecEnv,
every: int,
run_dir: Path | None = None,
writer: SummaryWriter | None = None,
):
"""Create a new evaluator
:param env: The evaluation environments (each evaluation runs one
episode per environment)
:param every: The number of training steps between two evaluations
:param run_dir: Where to record videos (optional)
:param writer: A tensorboard writer where the evaluation rewards are
logged (optional)
"""
#: The evaluation environments
self.env = env
#: The number of training steps between two evaluations
self.every = every
#: Where videos are recorded (see `visualize_best`)
self.run_dir = run_dir
#: The tensorboard writer the mean reward is logged to, if any
self.writer = writer
#: All the evaluations so far
self.history: list[EvalResult] = []
#: A copy of the best actor so far
self.best_actor: Actor[ActionT] | None = None
#: The mean reward of the best actor
self.best_reward = -float("inf")
self._last_eval_step = 0
[docs]
def run_if_needed(self, steps: int, actor: Actor[ActionT]) -> EvalResult | None:
"""Evaluate `actor` if it is time to do so
The call to make from the training loop: it does nothing until `every`
steps have been performed since the last evaluation. When it does
evaluate, the result is appended to `history`, logged to `writer`, and
a deep copy of the actor is kept in `best_actor` if it improves on
`best_reward`.
:param steps: The number of environment steps performed so far
(typically a collector's `steps`)
:param actor: The actor to evaluate
:returns: The `EvalResult`, or `None` when no evaluation was run
"""
if steps - self._last_eval_step < self.every:
return None
self._last_eval_step = steps
rewards = self.run(actor)
is_best = float(rewards.mean()) > self.best_reward
if is_best:
self.best_reward = float(rewards.mean())
self.best_actor = copy.deepcopy(actor)
result = EvalResult(step=steps, rewards=rewards, is_best=is_best)
self.history.append(result)
if self.writer is not None:
self.writer.add_scalar("eval/reward", result.mean, steps)
return result
[docs]
def run(self, actor: Actor[ActionT]) -> Tensor:
"""Run one evaluation episode in each environment
Steps every environment until *all* of them are done, ignoring the
rewards of an environment once its episode has ended. Does not update
`history` nor `best_actor` (see `run_if_needed`).
:param actor: The actor to evaluate (its `act`, deterministic,
behavior is used)
:returns: The cumulated reward of each episode, shape `[num_envs]`
"""
obs = self.env.reset()
cumulated = torch.zeros(self.env.num_envs)
finished = torch.zeros(self.env.num_envs, dtype=torch.bool)
with torch.no_grad():
while not bool(finished.all()):
step = self.env.step(actor.act(obs))
cumulated += step.reward * ~finished
finished |= step.done
obs = step.obs
return cumulated
[docs]
def visualize_best(
self,
video: bool = False,
seed: int = 0,
trimmer: Trimmer | None = None,
fps: int = 30,
width: str = "400px",
):
"""Play one episode of the best actor
By default, in a notebook, the episode is shown in the same player as
the one of `rl_mind.rendering` (`rollout_frames` + `animate_frames`:
play/pause and a slider to scrub through it). The player only works
while the kernel runs: with `video=True` (and always when run as a
script), a video is recorded under `<run_dir>/videos` (the working
directory if no `run_dir` was given) and embedded in the notebook,
where it is kept when the notebook is saved (a script prints its
path). Either way, the environment has the same wrappers and
arguments as the evaluation environments, and only `actor.act`
(deterministic) is used.
:param video: Record a video instead of showing the player
:param seed: The seed of the episode (player only)
:param trimmer: Bounds the frames kept in memory (player only;
`None` keeps every frame, ~10 KB each once encoded)
:param fps: The playback speed (player only)
:param width: The CSS width of the player (player only)
:returns: An `IPython.display.Video` when a video is shown in a
notebook, `None` otherwise
:raises AssertionError: If no evaluation has been run yet
"""
assert self.best_actor is not None, "No actor has been evaluated yet"
if not video and is_notebook():
history = rollout_frames(
self.env.env_name,
actor_policy(self.best_actor, self.env.gym_env.single_action_space),
seed=seed,
trimmer=trimmer,
wrappers=self.env.wrappers,
**self.env.env_kwargs,
)
animate_frames(history, fps=fps, width=width)
return None
directory = (self.run_dir or Path(".")) / "videos"
path = record_video(
self.best_actor,
self.env.env_name,
directory,
wrappers=self.env.wrappers,
**self.env.env_kwargs,
)
return video_display(path)
[docs]
def actor_policy(actor: Actor, action_space: gym.Space):
"""Turn an actor into a single-environment policy `obs -> action`
The observation (a numpy array, no batch dimension) is batched, given to
`actor.act` (the deterministic behavior), and the action converted back
to what `gymnasium` expects for `action_space`. This is what
`rl_mind.rendering.rollout_frames` needs to play a trained actor.
:param actor: The actor
:param action_space: The action space of the (single) environment
:returns: A callable `obs -> action`
"""
def policy(obs):
with torch.no_grad():
obs = torch.as_tensor(np.asarray(obs), dtype=torch.float32)
action = actor.act(obs.unsqueeze(0))[0]
if isinstance(action_space, gym.spaces.Discrete):
return int(action)
return action.numpy().astype(action_space.dtype)
return policy
[docs]
def record_video(
actor: Actor, env_name: str, directory: Path, wrappers=(), **kwargs
) -> Path | None:
"""Record a video of one episode played by `actor` (deterministic mode)
A fresh single environment is created with `render_mode="rgb_array"` and
played until the episode ends; only `actor.act` is used, so the video
shows the behavior the actor has learned, without exploration noise.
:param actor: The actor to film
:param env_name: The gymnasium environment name
:param directory: Where to write the video file
:param wrappers: gymnasium wrappers applied to the environment (must match
the ones used during training)
:param kwargs: Extra arguments passed to `gymnasium.make`
:returns: The path to the recorded video, or `None` if no file was written
"""
env = gym.make(env_name, render_mode="rgb_array", **kwargs)
for wrapper in wrappers:
env = wrapper(env)
env = gym.wrappers.RecordVideo(
env,
video_folder=str(directory),
episode_trigger=lambda i: True,
disable_logger=True,
)
policy = actor_policy(actor, env.action_space)
obs, _ = env.reset()
done = False
while not done:
obs, _, terminated, truncated, _ = env.step(policy(obs))
done = terminated or truncated
env.close()
videos = sorted(Path(directory).glob("*.mp4"), key=lambda p: p.stat().st_mtime)
return videos[-1] if videos else None