Rate this Page

★ ★ ★ ★ ★

reference/generated/torchrl.collectors.Evaluator

Run in Google Colab

Evaluator#

class torchrl.collectors.Evaluator(env: EnvBase | Callable[[], EnvBase], policy: TensorDictModuleBase | Callable | None = None, *, policy_factory: Callable[[...], Callable] | None = None, num_trajectories: int = 10, max_steps: int | None = None, frames_per_batch: int | None = None, collector_cls: type | str | None = None, collector_kwargs: dict | None = None, weight_sync_schemes: dict[str, Any] | None = None, log_prefix: str = 'eval', reward_keys: str | tuple[str, ...] = ('next', 'reward'), done_keys: str | tuple[str, ...] = ('next', 'done'), device: device | str | None = None, exploration_type: InteractionType = InteractionType.DETERMINISTIC, metrics_fn: Callable[[TensorDictBase], dict[str, float]] | None = None, dump_video: bool = True, on_result: Callable[[TensorDictBase], None] | None = None, episode_sampling: Literal['first_done', 'per_env'] = 'first_done', busy_policy: str = 'skip', backend: str = 'thread', init_fn: Callable[[], None] | None = None, num_gpus: int = 1, ray_kwargs: dict | None = None)[source]#

Unified sync / async evaluator with pluggable backend.

The evaluator wraps an environment and a policy and provides two modes of operation:

Synchronous – call evaluate() to run a blocking evaluation and get metrics back immediately.
Asynchronous – call trigger_eval() to kick off an evaluation in the background, then poll() (non-blocking) or wait() (blocking) to retrieve the result. Use the pending property to check whether an evaluation is currently in progress. Results can also be consumed via an on_result callback.

Internally, a Collector is used with trajs_per_batch=num_trajectories to collect complete episodes. The collector pre-allocates buffers and writes in-place — O(1) GPU allocations vs O(n) per step — yielding significant speedups for batched eval environments.

Three backends are available:

"thread" (default) – runs in a daemon thread. Low overhead, well suited for GPU-bound evaluation where the GIL is released by CUDA ops. When env is a callable and policy_factory is provided, both are created lazily inside the worker thread, which is useful for dedicated eval devices.
"process" – runs in a child process (spawn context). The env and policy are always created inside the child process, giving full CUDA context isolation and avoiding the GIL entirely. Requires env to be a callable and policy_factory to be provided.
"ray" – runs in a Ray actor, suitable for distributed setups. Requires env to be a callable and policy_factory to be provided.

Backpressure / overlap policy: calling trigger_eval() while a previous evaluation is still running skips the new request by default (busy_policy="skip"), raises immediately (busy_policy="error"), or queues the new request (busy_policy="queue"). trigger_eval() returns True when a request was accepted and False when it was skipped.

Callback thread-safety: when on_result is provided, it is invoked from the evaluator’s async coordination thread after the rollout completes. If the callback writes to a logger, the callback is responsible for any locking it needs.

Dedicated eval device (multi-GPU example):

evaluator = Evaluator(
    lambda: make_env(device="cuda:7"),
    policy_factory=lambda env: make_policy(env).to("cuda:7"),
    max_steps=1000,
    backend="process",  # or "thread"
)

Batched eval environments: for best results, add a RewardSum transform to the eval env so that per-episode returns are tracked. Without it, the evaluator falls back to summing raw rewards over each trajectory.

Eval video logging: when the eval env contains a VideoRecorder, each evaluation ends with a dump(step=...) on the env transforms of the collector that ran the rollout (disable with dump_video=False). The recorder must be attached to the outermost env: recorders nested inside the worker envs of a SerialEnv / ParallelEnv are not reached by the dump and would accumulate frames indefinitely. With backend="process" the recorder lives inside the collector worker, so build it in the env factory around a picklable logger client such as ProcessLogger.client():

from torchrl.record import VideoRecorder
from torchrl.record.loggers import CSVLogger, ProcessLogger

logger = ProcessLogger(CSVLogger, exp_name="run", log_dir="logs")
video_client = logger.client()

def make_eval_env():
    env = make_env()
    env.append_transform(VideoRecorder(video_client, tag="eval/video"))
    return env

evaluator = Evaluator(
    make_eval_env,
    policy_factory=make_eval_policy,
    max_steps=1000,
    backend="process",
)

Parameters:

env – An EnvBase instance or a callable that returns one. For the "process" and "ray" backends the callable form is required. For the "thread" backend, when combined with policy_factory, passing a callable defers construction to the worker thread.
policy – The evaluation policy. Mutually exclusive with policy_factory.

Keyword Arguments:

policy_factory – A callable (env) -> policy used to build the policy. Required for the "process" and "ray" backends. For "thread", if both env (callable) and policy_factory are provided, construction is deferred to the worker thread.
num_trajectories (int) – Number of complete episodes to collect per evaluation round. A Collector is used internally with trajs_per_batch=num_trajectories. Default: 10.
max_steps (int or None) – Maximum environment steps per episode, passed as max_frames_per_traj to the internal collector. When None, episodes run until done with no step limit. Default: None.
frames_per_batch (int or None) – Internal collection batch size (env steps per collector iteration). If None, defaults to max_steps. This is purely internal — output granularity is controlled by num_trajectories.
collector_cls – Which collector class to use. Accepts a class or a string name resolved from torchrl.collectors (e.g. "Collector"). Default: None (uses Collector).
collector_kwargs (dict or None) – Extra keyword arguments forwarded to the collector constructor.
log_prefix (str) – Prefix prepended to all logged metric names. Default: "eval".
reward_keys – Nested key(s) for reading the reward from the tensordict. Default: ("next", "reward").
done_keys – Nested key(s) for reading the done flag. Default: ("next", "done").
device – Device for the evaluation policy. If None, inferred from the policy parameters.
exploration_type – Exploration mode during evaluation. Default: ExplorationType.DETERMINISTIC.
metrics_fn – Optional (TensorDictBase) -> dict[str, float] called on every trajectory batch to extract custom metrics.
dump_video (bool) – Call dump() on VideoRecorder transforms after each evaluation. Process-backed collectors invoke the transform in their worker and can use a service-backed logger. Default: True.
on_result – Optional (TensorDictBase) -> None invoked after each completed evaluation. The callback receives a flat tensordict with the same prefixed metric names returned by evaluate(), poll(), and wait().
episode_sampling (str) – How episodes are gathered from a batched environment. "first_done" (default) keeps the first num_trajectories episodes to finish, which biases the sample towards short episodes, even with identical sub-environments. "per_env" rolls the environment out with break_when_all_done=True so that each rollout yields exactly one episode per sub-environment; it requires the "thread" backend without weight_sync_schemes and an integer max_steps. evaluate_episodes() always samples "per_env". Default: "first_done".
busy_policy (str) –
Behaviour when trigger_eval() is called while another async evaluation is still pending. "skip" returns False without scheduling a new request (default). "error" raises immediately. "queue" enqueues the new request and runs it when the current evaluation finishes.

Warning

With busy_policy="queue", each queued request stores a copy of the weights dict. For large models this can consume significant memory. Prefer checking pending and skipping triggers instead.

weight_sync_schemes (dict or None) –

A dict mapping model IDs to WeightSyncScheme instances. When provided, a MultiSyncCollector with a single worker is used for process-level CUDA isolation and scheme-based weight transfer. Model IDs follow the collector convention: "policy" for the main policy, "env.transform[0]" for env transforms, etc. Example:

from torchrl.weight_update import MultiProcessedWeightSyncScheme
evaluator = Evaluator(
    env=make_eval_env,
    policy_factory=make_eval_policy,
    weight_sync_schemes={
        "policy": MultiProcessedWeightSyncScheme(),
        "env.transform[0]": MultiProcessedWeightSyncScheme(),
    },
    max_steps=1000,
)

backend (str) – "thread" (default), "process", or "ray". The "process" backend is implemented as a thread backend with a MultiSyncCollector (1 worker) running in a child process. This provides full CUDA context isolation without custom queue management.
init_fn – Callable invoked at the start of the worker / actor process, before any evaluation work (and, ideally, before any torch import inside that process). Used by both the "process" and "ray" backends. Typical use case: start Isaac Lab’s AppLauncher in headless mode. Ignored by the "thread" backend because no new process is spawned.
num_gpus (int) – (Ray only) GPUs requested for the actor. Default: 1.
ray_kwargs (dict) – (Ray only) Extra keyword arguments forwarded to ray.remote().

Run a blocking evaluation rollout.

Parameters:

weights – Policy weights to load before the rollout. Accepts a TensorDictBase (e.g. from TensorDict.from_module(policy).data) or an nn.Module (weights are extracted automatically). If None the current policy weights are used.
step – Logging step. If None an internal counter is used.

Keyword Arguments:

weights_dict – A dict mapping model_id strings to weight sources (nn.Module or TensorDictBase). Use this to sync multiple models (e.g. policy + env transforms). When provided, weights is treated as weights_dict["policy"] if "policy" is not already in the dict.

Returns:

dict with at least "<prefix>/reward" and "<prefix>/episode_length" keys.

Run a blocking evaluation and return one entry per episode.

Episodes are always sampled "per_env" (see episode_sampling): each rollout of the environment contributes exactly one episode per sub-environment until num_trajectories episodes are gathered. Unlike evaluate(), this does not advance the internal step counter and does not call on_result or metrics_fn. When the env has a device, the policy may live on another one: data is moved between them at each step. The policy’s train/eval mode is restored afterwards, also when a rollout raises. Requires the "thread" backend without weight_sync_schemes and an integer max_steps.

Parameters:

weights (TensorDictBase or nn.Module, optional) – Policy weights to load before the rollouts. See evaluate(). Defaults to None (current weights).
step (int, optional) – Step at which the video of a VideoRecorder on the env is logged. Defaults to None (the recorder’s own counter).

Keyword Arguments:

weights_dict (dict, optional) – Multi-model weights dict. See evaluate(). Defaults to None.
seed (int, optional) – Before each rollout, the torch RNG (CPU and the current accelerator, e.g. CUDA or MPS) is seeded and set_seed() is called with a seed derived from seed and the rollout index only, so two evaluations with the same seed start each episode from the same state whatever the policy, and their results are paired episode by episode. The caller’s torch RNG state is restored afterwards, while the environment keeps the seeds it was given. Defaults to None (RNGs are left untouched).

Returns:

a TensorDict of batch size [num_trajectories] on the env device, with entries "episode_reward" (the final ("next", "episode_reward") of a RewardSum when present, otherwise the sum of rewards), "episode_length" (int64), "truncated" (bool, True when the episode reached max_steps without a done or the env flagged it as truncated) and "env_index" (int64, the flat index of the sub-environment that ran the episode).

Examples

>>> from torchrl.collectors import Evaluator, paired_comparison
>>> from torchrl.envs import GymEnv, SerialEnv
>>> from torchrl.modules import RandomPolicy
>>> env = SerialEnv(4, lambda: GymEnv("Pendulum-v1"))
>>> evaluator = Evaluator(
...     env, RandomPolicy(env.action_spec), num_trajectories=6, max_steps=50
... )
>>> episodes = evaluator.evaluate_episodes(seed=0)
>>> episodes["env_index"], episodes["truncated"].all()
(tensor([0, 1, 2, 3, 0, 1]), tensor(True))
>>> again = evaluator.evaluate_episodes(seed=0)
>>> paired_comparison(again["episode_reward"], episodes["episode_reward"])[
...     "estimate"
... ]
tensor(0.)
>>> evaluator.shutdown()

static extract_weights(policy: Module) → TensorDictBase[source]#

Extract detached, cloned, CPU weights from a policy.

This is a convenience helper; the returned TensorDict is safe to pass across threads.

load_state_dict(state_dict: Mapping[str, Any]) → None[source]#: Restore evaluator progress when no asynchronous work is pending.

property pending: bool#

Return True if an async evaluation is currently in progress.

This can be used to avoid triggering overlapping evaluations:

if not evaluator.pending:
    evaluator.trigger_eval(weights, step=step)

poll(timeout: float = 0) → dict[str, Any] | None[source]#

Return the latest evaluation result if ready, else None.

Parameters:: timeout – Seconds to wait. 0 means non-blocking.

shutdown(timeout: float = 5.0) → None[source]#: Cancel any running evaluation, clean up resources.

state_dict() → dict[str, Any][source]#: Return evaluator progress when no asynchronous work is pending.

Start an async evaluation.

Parameters:

weights – Policy weights to load. See evaluate().
step – Logging step. See evaluate().
weights_dict – Multi-model weights dict. See evaluate().

Returns:

True if an evaluation request was scheduled, False if it was skipped because another request was pending and busy_policy="skip".

wait(timeout: float | None = None) → dict[str, Any] | None[source]#

Block until the current evaluation finishes.

Parameters:: timeout – Max seconds to wait. None waits forever.

Evaluator#

Docs

Tutorials

Resources