temporal_aggregation

temporal_aggregation

Temporal aggregation for action sequences.

TemporalAggregator

TemporalAggregator(device, action_keys_to_dimensions, prediction_horizon, max_timesteps=10000, exponential_decay=0.01, favor_more_recent=True)

Exponential-weighted temporal averaging of overlapping action predictions.

Accumulates action chunks over time and returns a weighted average for the current step. Works with any set of action keys and dimensions.

Initialize temporal aggregator.

Parameters:

Name	Type	Description	Default
device	device	Torch device for tensors.	required
action_keys_to_dimensions	dict[str, int]	Mapping from action key to dimension.	required
prediction_horizon	int	Number of future steps predicted per inference.	required
max_timesteps	int	Maximum episode length.	10000
exponential_decay	float	Decay factor for exponential weighting.	0.01
favor_more_recent	bool	Whether to weight newer predictions more heavily.	True

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

def __init__(
    self,
    device: torch.device,
    action_keys_to_dimensions: dict[str, int],
    prediction_horizon: int,
    max_timesteps: int = 10000,
    exponential_decay: float = 0.01,
    favor_more_recent: bool = True,
):
    """Initialize temporal aggregator.

    Args:
        device: Torch device for tensors.
        action_keys_to_dimensions: Mapping from action key to dimension.
        prediction_horizon: Number of future steps predicted per inference.
        max_timesteps: Maximum episode length.
        exponential_decay: Decay factor for exponential weighting.
        favor_more_recent: Whether to weight newer predictions more heavily.
    """
    self.device = device
    self.action_keys_to_dimensions = action_keys_to_dimensions
    self.prediction_horizon = prediction_horizon
    self.max_timesteps = max_timesteps
    self.exponential_decay = exponential_decay
    self.favor_more_recent = favor_more_recent
    self.timestep = 0
    total_length = self.max_timesteps + self.prediction_horizon
    self.populated_mask = torch.zeros(
        [self.max_timesteps, total_length],
        dtype=torch.bool,
    ).to(self.device)
    self.action_histories: dict[str, torch.Tensor] = {}
    for key, dimension in self.action_keys_to_dimensions.items():
        self.action_histories[key] = torch.zeros(
            [self.max_timesteps, total_length, dimension]
        ).to(self.device)

store_and_average

store_and_average(current_predictions)

Store current predictions and return averaged action for this timestep.

Parameters:

Name	Type	Description	Default
current_predictions	dict[str, Tensor]	Dict mapping action key to predicted tensor of shape (prediction_horizon, dimension).	required

Returns:

Type	Description
dict[str, Tensor]	Dict mapping action key to averaged tensor of shape (dimension,).

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

def store_and_average(
    self, current_predictions: dict[str, torch.Tensor]
) -> dict[str, torch.Tensor]:
    """Store current predictions and return averaged action for this timestep.

    Args:
        current_predictions: Dict mapping action key to predicted tensor
            of shape (prediction_horizon, dimension).

    Returns:
        Dict mapping action key to averaged tensor of shape (dimension,).
    """
    if self.timestep >= self.max_timesteps:
        raise RuntimeError(
            f"TemporalAggregator exceeded max_timesteps={self.max_timesteps} "
            f"at timestep {self.timestep}. Increase max_timesteps or call reset() "
            f"between episodes."
        )
    horizon_slice = slice(self.timestep, self.timestep + self.prediction_horizon)
    self.populated_mask[[self.timestep], horizon_slice] = True
    for key, predictions in current_predictions.items():
        self.action_histories[key][[self.timestep], horizon_slice] = (
            predictions.float()
        )

    actions_populated = self.populated_mask[:, self.timestep]
    num_populated = int(actions_populated.sum().item())
    exponential_weights = self._compute_exponential_weights(num_populated)
    averaged = {}
    for key in current_predictions:
        actions_for_step = self.action_histories[key][:, self.timestep][
            actions_populated
        ]
        averaged[key] = (actions_for_step * exponential_weights).sum(dim=0)

    self.timestep += 1
    return averaged

reset

reset()

Reset all state for a new episode.

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

def reset(self) -> None:
    """Reset all state for a new episode."""
    self.timestep = 0
    self.populated_mask.zero_()
    for tensor in self.action_histories.values():
        tensor.zero_()