synthetic_rollout

synthetic_rollout

Rollout and evaluation utilities for synthetic benchmark policies.

Provides open-loop and closed-loop rollout functions that load a trained policy checkpoint, generate position trajectories via the policy, and evaluate them against expert demonstrations using mode coverage and goal success metrics.

load_policy_from_checkpoint

load_policy_from_checkpoint(checkpoint_path, device='cuda', checkpoint_name='last.ckpt')

Load a trained policy and config from a checkpoint directory.

Follows the same loading pattern as the production inference clients: instantiate config, create policy, load weights via LightningPolicy.

Parameters:

Name	Type	Description	Default
checkpoint_path	str	Directory containing config.yaml and the .ckpt file.	required
device	str	Torch device string.	'cuda'
checkpoint_name	str	Checkpoint filename inside checkpoint_path.	'last.ckpt'

Returns:

Type	Description
tuple[Policy, MainConfig]	Tuple of (policy in eval mode, resolved MainConfig).

Source code in src/versatil/inference/synthetic_rollout.py

def load_policy_from_checkpoint(
    checkpoint_path: str,
    device: str = "cuda",
    checkpoint_name: str = "last.ckpt",
) -> tuple[Policy, MainConfig]:
    """Load a trained policy and config from a checkpoint directory.

    Follows the same loading pattern as the production inference clients:
    instantiate config, create policy, load weights via LightningPolicy.

    Args:
        checkpoint_path: Directory containing config.yaml and the .ckpt file.
        device: Torch device string.
        checkpoint_name: Checkpoint filename inside checkpoint_path.

    Returns:
        Tuple of (policy in eval mode, resolved MainConfig).
    """
    config_file = os.path.join(checkpoint_path, "config.yaml")
    if not os.path.exists(config_file):
        raise FileNotFoundError(f"Config not found at {config_file}")
    config: MainConfig = hydra.utils.instantiate(OmegaConf.load(config_file))
    policy: Policy = config.policy
    policy.to(device).eval()
    checkpoint_file = os.path.join(checkpoint_path, checkpoint_name)
    if not os.path.exists(checkpoint_file):
        raise FileNotFoundError(f"Checkpoint not found at {checkpoint_file}")
    checkpoint = torch.load(checkpoint_file, map_location=device, weights_only=False)
    lightning_module = LightningPolicy(policy=policy, training_config=config.training)
    lightning_module.load_state_dict(checkpoint["state_dict"], strict=False)
    return policy, config

run_rollouts

run_rollouts(policy, task_name, num_rollouts, image_size=DEFAULT_IMAGE_SIZE, context_mode=None, temporal_aggregation=False, exponential_decay=0.01, output_dir=None)

Re-render and re-predict at each timestep with temporal aggregation.

At each step, the policy predicts a full action trajectory. With temporal aggregation enabled, overlapping predictions for the current timestep are averaged with exponential weights favoring more recent queries. Without it, the policy executes the full predicted chunk before replanning.

Policy-input frames intentionally render without any goal marker. Per-mode expert goals are still drawn in the user-facing rollout visualization via layout.goals.

Parameters:

Name	Type	Description	Default
policy	Policy	Trained policy in eval mode.	required
task_name	str	SyntheticTaskName.value string.	required
num_rollouts	int	Number of independent rollouts.	required
image_size	int	Side length for rendered images.	DEFAULT_IMAGE_SIZE
context_mode	int | None	Context mode index for conditional tasks (None to omit).	None
temporal_aggregation	bool	Average overlapping action predictions with exponential weighting. Matches the production inference clients.	False
exponential_decay	float	Decay factor for temporal aggregation weights. Smaller values produce more uniform weighting across queries.	0.01
output_dir	str | None	Optional directory for saving PNG + GIF visualizations of the rollout trajectories. If None, no visualizations are saved.	None

Returns:

Type	Description
ndarray	Position trajectories, shape (num_rollouts, prediction_horizon + 1, 2).

Source code in src/versatil/inference/synthetic_rollout.py

def run_rollouts(
    policy: Policy,
    task_name: str,
    num_rollouts: int,
    image_size: int = DEFAULT_IMAGE_SIZE,
    context_mode: int | None = None,
    temporal_aggregation: bool = False,
    exponential_decay: float = 0.01,
    output_dir: str | None = None,
) -> np.ndarray:
    """Re-render and re-predict at each timestep with temporal aggregation.

    At each step, the policy predicts a full action trajectory. With
    temporal aggregation enabled, overlapping predictions for the current
    timestep are averaged with exponential weights favoring more recent
    queries. Without it, the policy executes the full predicted chunk
    before replanning.

    Policy-input frames intentionally render without any goal marker.
    Per-mode expert goals are still drawn in the user-facing rollout
    visualization via ``layout.goals``.

    Args:
        policy: Trained policy in eval mode.
        task_name: SyntheticTaskName.value string.
        num_rollouts: Number of independent rollouts.
        image_size: Side length for rendered images.
        context_mode: Context mode index for conditional tasks (None to omit).
        temporal_aggregation: Average overlapping action predictions with
            exponential weighting. Matches the production inference clients.
        exponential_decay: Decay factor for temporal aggregation weights.
            Smaller values produce more uniform weighting across queries.
        output_dir: Optional directory for saving PNG + GIF visualizations
            of the rollout trajectories. If None, no visualizations are saved.

    Returns:
        Position trajectories, shape (num_rollouts, prediction_horizon + 1, 2).
    """
    layout = get_task_layout(task_name=task_name)
    start = layout.start
    obstacles = layout.obstacles
    num_modes = layout.num_modes
    prediction_horizon = policy.prediction_horizon
    observation_keys = set(policy.observation_space.observations_metadata.keys())
    action_key = ProprioKey.SYNTHETIC_POSITION_ACTION.value
    context_vector = None
    context_color = None
    if context_mode is not None:
        context_vector = np.zeros(num_modes, dtype=np.float32)
        context_vector[context_mode] = 1.0
        context_color = CIRCLE_CONTEXT_COLORS.get(context_mode)

    obs_horizon = policy.observation_horizon
    all_trajectories = np.zeros(
        (num_rollouts, prediction_horizon + 1, 2), dtype=np.float32
    )

    for rollout_index in range(num_rollouts):
        positions = all_trajectories[rollout_index]
        positions[0] = start.copy()

        if temporal_aggregation:
            action_buffer = np.zeros(
                (prediction_horizon, prediction_horizon, 2), dtype=np.float32
            )
            populated_mask = np.zeros(
                (prediction_horizon, prediction_horizon), dtype=bool
            )
            for step in range(prediction_horizon):
                trail = positions[: step + 1]
                available_history = step + 1
                if available_history < obs_horizon:
                    positions[step + 1] = positions[step].copy()
                    continue
                history_start = step + 1 - obs_horizon
                history = positions[history_start : step + 1]
                observation = _prepare_observation(
                    position_history=history,
                    obstacles=obstacles,
                    image_size=image_size,
                    observation_keys=observation_keys,
                    trail=trail,
                    context_vector=context_vector,
                    context_color=context_color,
                )
                with torch.no_grad():
                    actions = policy.predict_action(obs_dict=observation)
                action_deltas = actions[action_key].squeeze(0).cpu().numpy()
                remaining = min(prediction_horizon - step, prediction_horizon)
                action_buffer[step, step : step + remaining] = action_deltas[:remaining]
                populated_mask[step, step : step + remaining] = True
                valid_queries = populated_mask[:, step]
                candidate_actions = action_buffer[valid_queries, step]
                num_candidates = len(candidate_actions)
                indices = np.arange(num_candidates)[::-1]
                weights = np.exp(-exponential_decay * indices)
                weights = weights / weights.sum()
                selected_action = (candidate_actions * weights[:, np.newaxis]).sum(
                    axis=0
                )
                positions[step + 1] = np.clip(
                    positions[step] + selected_action, 0.0, 1.0
                )
        else:
            # Execute full chunk without replanning
            step = 0
            while step < prediction_horizon:
                available_history = step + 1
                if available_history < obs_horizon:
                    positions[step + 1] = positions[step].copy()
                    step += 1
                    continue
                trail = positions[: step + 1]
                history_start = step + 1 - obs_horizon
                history = positions[history_start : step + 1]
                observation = _prepare_observation(
                    position_history=history,
                    obstacles=obstacles,
                    image_size=image_size,
                    observation_keys=observation_keys,
                    trail=trail,
                    context_vector=context_vector,
                    context_color=context_color,
                )
                with torch.no_grad():
                    actions = policy.predict_action(obs_dict=observation)
                action_deltas = actions[action_key].squeeze(0).cpu().numpy()
                # Execute all remaining actions in this chunk
                remaining = min(prediction_horizon - step, prediction_horizon)
                for chunk_offset in range(remaining):
                    positions[step + 1] = np.clip(
                        positions[step] + action_deltas[chunk_offset], 0.0, 1.0
                    )
                    step += 1

    if output_dir is not None:
        name_prefix = "rollout_temporal_agg" if temporal_aggregation else "rollout"
        _save_rollout_visualizations(
            trajectories=all_trajectories,
            task_name=task_name,
            output_dir=output_dir,
            name_prefix=name_prefix,
            image_size=image_size,
        )
    return all_trajectories

evaluate_rollouts

evaluate_rollouts(rollout_trajectories, task_name, num_expert_episodes=1000, expert_seed=42, trajectory_length=MULTIPATH_DEFAULT_TRAJECTORY_LENGTH, noise_std=MULTIPATH_DEFAULT_NOISE_STD, num_modes=MULTIPATH_DEFAULT_NUM_MODES, num_styles=CORRIDOR_DEFAULT_NUM_STYLES, image_size=DEFAULT_IMAGE_SIZE)

Evaluate rollout trajectories against regenerated expert demonstrations.

Generates expert data with the given parameters, then computes mode coverage and obstacle-aware success rate. The "reach" threshold used by the synthetic success metrics is derived from the expert endpoint spread (mean plus five standard deviations of distances from the per-mode endpoint mean), so it scales with the task's intrinsic noise.

Parameters:

Name	Type	Description	Default
rollout_trajectories	ndarray	Generated position trajectories. Shape (num_rollouts, num_timesteps, 2).	required
task_name	str	SyntheticTaskName.value string.	required
num_expert_episodes	int	Number of expert episodes for reference.	1000
expert_seed	int	Seed for reproducible expert generation.	42
trajectory_length	int	Expert trajectory length.	MULTIPATH_DEFAULT_TRAJECTORY_LENGTH
noise_std	float	Expert noise standard deviation.	MULTIPATH_DEFAULT_NOISE_STD
num_modes	int	Number of modes for expert generation.	MULTIPATH_DEFAULT_NUM_MODES
num_styles	int	Number of styles (trajectory_style task only).	CORRIDOR_DEFAULT_NUM_STYLES
image_size	int	Image size for expert generation.	DEFAULT_IMAGE_SIZE

Returns:

Type	Description
dict[str, float | dict[int, int]]	Dict with raw mode coverage metrics, valid mode coverage metrics,
dict[str, float | dict[int, int]]	success_rate, collision_rate, endpoint_reach_rate, and path_length_rate.

Source code in src/versatil/inference/synthetic_rollout.py

def evaluate_rollouts(
    rollout_trajectories: np.ndarray,
    task_name: str,
    num_expert_episodes: int = 1000,
    expert_seed: int = 42,
    trajectory_length: int = MULTIPATH_DEFAULT_TRAJECTORY_LENGTH,
    noise_std: float = MULTIPATH_DEFAULT_NOISE_STD,
    num_modes: int = MULTIPATH_DEFAULT_NUM_MODES,
    num_styles: int = CORRIDOR_DEFAULT_NUM_STYLES,
    image_size: int = DEFAULT_IMAGE_SIZE,
) -> dict[str, float | dict[int, int]]:
    """Evaluate rollout trajectories against regenerated expert demonstrations.

    Generates expert data with the given parameters, then computes
    mode coverage and obstacle-aware success rate. The "reach" threshold
    used by the synthetic success metrics is derived from the expert endpoint
    spread (mean plus five standard deviations of distances from the per-mode
    endpoint mean), so it scales with the task's intrinsic noise.

    Args:
        rollout_trajectories: Generated position trajectories.
            Shape (num_rollouts, num_timesteps, 2).
        task_name: SyntheticTaskName.value string.
        num_expert_episodes: Number of expert episodes for reference.
        expert_seed: Seed for reproducible expert generation.
        trajectory_length: Expert trajectory length.
        noise_std: Expert noise standard deviation.
        num_modes: Number of modes for expert generation.
        num_styles: Number of styles (trajectory_style task only).
        image_size: Image size for expert generation.

    Returns:
        Dict with raw mode coverage metrics, valid mode coverage metrics,
        success_rate, collision_rate, endpoint_reach_rate, and path_length_rate.
    """
    expert_episodes = generate_task_episodes(
        task_name=task_name,
        num_episodes=num_expert_episodes,
        seed=expert_seed,
        image_size=image_size,
        num_modes=num_modes,
        trajectory_length=trajectory_length,
        noise_std=noise_std,
        num_styles=num_styles,
    )

    expert_trajectories = np.array([episode["position"] for episode in expert_episodes])
    expert_mode_ids = np.array(
        [int(episode["mode_id"][0, 0]) for episode in expert_episodes]
    )
    layout = get_task_layout(
        task_name=task_name,
        num_modes=num_modes,
        num_styles=num_styles,
        noise_std=noise_std,
    )
    rollout_length = rollout_trajectories.shape[1]
    expert_length = expert_trajectories.shape[1]
    comparison_length = min(rollout_length, expert_length)
    rollout_comparison_trajectories = rollout_trajectories[:, :comparison_length, :]
    expert_comparison_trajectories = expert_trajectories[:, :comparison_length, :]
    mode_endpoints = compute_mode_endpoints(
        expert_trajectories=expert_trajectories,
        expert_mode_ids=expert_mode_ids,
        num_modes=layout.num_modes,
    )
    reach_threshold = _expert_endpoint_reach_threshold(
        expert_trajectories=expert_trajectories,
        expert_mode_ids=expert_mode_ids,
        mode_endpoints=mode_endpoints,
    )
    min_path_length = 0.5 * _expert_mean_path_length(
        expert_trajectories=expert_trajectories
    )
    success_masks = compute_success_masks(
        generated_trajectories=rollout_trajectories,
        obstacles=layout.obstacles,
        mode_endpoints=mode_endpoints,
        goal_threshold=reach_threshold,
        min_path_length=min_path_length,
    )
    collision_free_mask = ~success_masks["collision_mask"]
    coverage_results = compute_mode_coverage(
        generated_trajectories=rollout_comparison_trajectories,
        expert_trajectories=expert_comparison_trajectories,
        expert_mode_ids=expert_mode_ids,
        num_modes=layout.num_modes,
        valid_mask=collision_free_mask,
    )
    valid_coverage_results = compute_mode_coverage(
        generated_trajectories=rollout_comparison_trajectories,
        expert_trajectories=expert_comparison_trajectories,
        expert_mode_ids=expert_mode_ids,
        num_modes=layout.num_modes,
        valid_mask=success_masks["success_mask"],
    )
    success_stats = compute_success_rates_from_masks(success_masks=success_masks)
    results = dict(coverage_results)
    results["valid_mode_coverage"] = valid_coverage_results["mode_coverage"]
    results["valid_mode_entropy_ratio"] = valid_coverage_results["mode_entropy_ratio"]
    results.update(success_stats)
    return results