Source code for unilab.base.curriculum

"""Curriculum learning for adaptive difficulty adjustment."""

from __future__ import annotations

import torch




class EpisodeLengthTracker:
    """Track a weighted moving average of episode lengths on one tensor."""



    def __init__(self, num_envs: int, window_size: int = 1000, *, device=None):
        self.num_envs = num_envs
        self.window_size = max(1, int(window_size * num_envs / 4096))
        device = torch.device("cpu") if device is None else torch.device(device)
        self._average = torch.zeros((), dtype=torch.float64, device=device)




    def update(self, episode_lengths: torch.Tensor) -> None:
        """Recursively update the average without synchronizing each batch."""
        count = int(episode_lengths.numel())
        if count == 0:
            return
        lengths = episode_lengths.to(dtype=torch.float64, device=self._average.device)
        batch_average = lengths.sum() / count
        weight = torch.as_tensor(
            min(count / self.window_size, 1.0),
            dtype=torch.float64,
            device=self._average.device,
        )
        self._average.add_((batch_average - self._average) * weight)




    def update_masked(self, episode_lengths: torch.Tensor, mask: torch.Tensor) -> None:
        """Update from a fixed-width episode-length carrier and bool mask.

        This boundary avoids materializing reset
        rows; false values are excluded with fixed-shape ``where`` operations.
        """
        if mask.shape != episode_lengths.shape:
            raise ValueError("episode lengths and update mask must share one shape")
        selected_mask = mask.to(device=self._average.device, dtype=torch.bool)
        count = selected_mask.sum()
        if bool(count == 0):
            return
        lengths = episode_lengths.to(dtype=torch.float64, device=self._average.device)
        selected = torch.where(
            selected_mask,
            lengths,
            torch.zeros_like(lengths),
        )
        batch_average = selected.sum() / count.clamp(min=1)
        weight = (count.to(dtype=torch.float64) / float(self.window_size)).clamp(max=1.0)
        self._average.add_((batch_average - self._average) * weight)


    @property
    def average(self) -> torch.Tensor:
        """Recursive average as a device-resident scalar tensor."""
        return self._average