synthetic_rollout

synthetic_rollout

Callback for evaluating synthetic benchmark policies during training.

SyntheticRolloutCallback

SyntheticRolloutCallback(task_name, num_modes, num_styles, trajectory_length, noise_std, num_rollouts=50, image_size=64, log_every_n_epochs=1)

Bases: Callback

Run rollouts and log mode coverage metrics at the end of each training epoch.

Puts the policy in eval mode, generates trajectories via closed-loop rollout, computes mode coverage and goal success against regenerated expert demonstrations, and logs metrics + trajectory plots to wandb.

Parameters:

Name	Type	Description	Default
task_name	str	SyntheticTaskName.value string.	required
num_modes	int	Number of behavioral modes to generate for expert reference. Must match the training dataset.	required
num_styles	int	Number of sinusoidal styles per corridor gap. Ignored by tasks that do not use styles.	required
trajectory_length	int	Length of generated expert and rollout trajectories.	required
noise_std	float	Standard deviation of expert trajectory noise.	required
num_rollouts	int	Number of rollout trajectories per evaluation.	50
image_size	int	Side length for rendered observation images.	64
log_every_n_epochs	int	Evaluate every N epochs.	1

Initialize rollout generation and logging parameters.

Source code in src/versatil/training/callbacks/synthetic_rollout.py

def __init__(
    self,
    task_name: str,
    num_modes: int,
    num_styles: int,
    trajectory_length: int,
    noise_std: float,
    num_rollouts: int = 50,
    image_size: int = 64,
    log_every_n_epochs: int = 1,
):
    """Initialize rollout generation and logging parameters."""
    super().__init__()
    self.task_name = task_name
    self.num_modes = num_modes
    self.num_styles = num_styles
    self.trajectory_length = trajectory_length
    self.noise_std = noise_std
    self.num_rollouts = num_rollouts
    self.image_size = image_size
    self.log_every_n_epochs = log_every_n_epochs
    self._training_data_logged = False

on_train_epoch_end

on_train_epoch_end(trainer, pl_module)

Run rollouts, compute metrics, log to wandb and console.

Source code in src/versatil/training/callbacks/synthetic_rollout.py

def on_train_epoch_end(
    self, trainer: pl.Trainer, pl_module: pl.LightningModule
) -> None:
    """Run rollouts, compute metrics, log to wandb and console."""
    if trainer.logger is not None and not self._training_data_logged:
        self._log_training_data(trainer=trainer)
        self._training_data_logged = True

    is_log_epoch = trainer.current_epoch % self.log_every_n_epochs == 0
    is_last_epoch = trainer.current_epoch == trainer.max_epochs - 1
    if not is_log_epoch and not is_last_epoch:
        return

    policy: Policy = pl_module.policy
    was_training = policy.training
    policy.eval()

    context_modes = self._resolve_context_modes(policy=policy)
    precision_type = PrecisionType(str(trainer.precision))
    with (
        torch.no_grad(),
        precision_type.autocast(device_type=pl_module.device.type),
    ):
        per_mode_trajectories = [
            run_rollouts(
                policy=policy,
                task_name=self.task_name,
                num_rollouts=self.num_rollouts,
                image_size=self.image_size,
                context_mode=mode,
                temporal_aggregation=False,  # open-loop
            )
            for mode in context_modes
        ]
    trajectories = (
        per_mode_trajectories[0]
        if len(per_mode_trajectories) == 1
        else np.concatenate(per_mode_trajectories, axis=0)
    )

    results = evaluate_rollouts(
        rollout_trajectories=trajectories,
        task_name=self.task_name,
        image_size=self.image_size,
        num_modes=self.num_modes,
        num_styles=self.num_styles,
        trajectory_length=self.trajectory_length,
        noise_std=self.noise_std,
    )

    epoch = trainer.current_epoch
    mode_coverage = results["mode_coverage"]
    entropy_ratio = results["mode_entropy_ratio"]
    valid_mode_coverage = results["valid_mode_coverage"]
    valid_entropy_ratio = results["valid_mode_entropy_ratio"]
    per_mode = results["per_mode_count"]
    success_rate = results["success_rate"]
    collision_rate = results["collision_rate"]
    endpoint_reach_rate = results["endpoint_reach_rate"]
    path_length_rate = results["path_length_rate"]

    log_parts = [
        f"epoch {epoch}",
        f"success={success_rate:.2f}",
        f"collision={collision_rate:.2f}",
        f"endpoint_reach={endpoint_reach_rate:.2f}",
        f"path_length={path_length_rate:.2f}",
        f"valid_mode_coverage={valid_mode_coverage:.2f}",
        f"valid_entropy={valid_entropy_ratio:.2f}",
        f"raw_mode_coverage={mode_coverage:.2f}",
        f"raw_entropy={entropy_ratio:.2f}",
        f"per_mode={per_mode}",
    ]
    logging.info(f"Synthetic rollout: {', '.join(log_parts)}")

    if trainer.logger is not None:
        metrics: dict[str, float | wandb.Image] = {
            "synthetic/success_rate": success_rate,
            "synthetic/collision_rate": collision_rate,
            "synthetic/endpoint_reach_rate": endpoint_reach_rate,
            "synthetic/path_length_rate": path_length_rate,
            "synthetic/valid_mode_coverage": valid_mode_coverage,
            "synthetic/valid_mode_entropy_ratio": valid_entropy_ratio,
            "synthetic/mode_coverage": mode_coverage,
            "synthetic/mode_entropy_ratio": entropy_ratio,
        }
        for mode_index, count in per_mode.items():
            metrics[f"synthetic/mode_{mode_index}_count"] = count

        rollout_figure = plot_trajectories_2d(
            trajectories=trajectories,
            task_name=self.task_name,
            num_modes=self.num_modes,
            num_styles=self.num_styles,
            noise_std=self.noise_std,
        )
        metrics["synthetic/rollout_trajectories"] = figure_to_wandb_image(
            rollout_figure, dpi=150
        )
        plt.close(rollout_figure)

        trainer.logger.log_metrics(metrics, step=epoch)

    if was_training:
        policy.train()