Training › Continual Learning¶

@abstractmethod
def observe(self, batch: dict, task_id: int) -> None:
    """Called every training step with the **current task** batch.

    Replay-based methods (ER, GEM) use this to grow/refresh memory.
    Regularization methods (EWC, SI) use it to accumulate importance
    statistics (Fisher information, path integrals, etc.).
    """

@abstractmethod
def sample(self, batch_size: int) -> Any:
    """Return the algorithm's auxiliary artifact for this step (or None/empty to skip).

    Return shape is algorithm-specific:
      * ER / GEM : `list[dict]`       — raw samples ready to be collated.
      * EWC / SI : `dict[str, Tensor]`— per-parameter regularization terms.
      * LwF      : `dict[str, Tensor]`— teacher logits on the current batch.

    The trainer dispatches on algorithm type to combine this with the
    current-task batch (e.g. mix-in ratio for replay, KL term for LwF).
    """

@abstractmethod
def on_task_end(self, task_id: int) -> None:
    """Hook invoked after a task finishes.

    Typical uses:
    * EWC: snapshot parameters, compute Fisher on current task's dataset.
    * LwF: snapshot the teacher model weights.
    * ER : no-op (reservoir sampling happens online).
    """

@abstractmethod
def state_dict(self) -> dict[str, Any]:
    """Return a JSON-serializable snapshot of the algorithm state.

    This is written alongside model checkpoints so CL state survives
    interruption/restart across tasks.
    """

@abstractmethod
def load_state_dict(self, state: dict[str, Any]) -> None:
    """Restore algorithm state from a dict produced by `state_dict()`."""

def populate_from_dataset(self, task_id: int, dataset: Dataset, num_samples: Optional[int] = None):
    """Store samples from a dataset into the buffer using reservoir sampling.

    Args:
        task_id: Identifier for the task.
        dataset: Dataset to sample from (must support __len__ and __getitem__).
        num_samples: Number of samples to store. Defaults to buffer_size_per_task.
    """
    if num_samples is None:
        num_samples = self.buffer_size_per_task

    n = len(dataset)
    k = min(num_samples, n)

    # Reservoir sampling: select k items from n uniformly at random
    indices = list(range(n))
    self.rng.shuffle(indices)
    selected_indices = sorted(indices[:k])

    samples = []
    for idx in selected_indices:
        try:
            sample = dataset[idx]
            samples.append(sample)
        except Exception as e:
            logger.warning(f"Failed to read sample {idx} for task {task_id}: {e}")
            continue

    # Update buffer
    if task_id in self._buffers:
        self._total_samples -= len(self._buffers[task_id])
    self._buffers[task_id] = samples
    self._total_samples += len(samples)

    logger.info(
        f"Replay buffer: stored {len(samples)} samples for task {task_id} "
        f"(total: {self._total_samples} across {self.num_tasks} tasks)"
    )

def sample(self, batch_size: int) -> List[dict]:
    """Sample a batch uniformly from all stored tasks.

    Args:
        batch_size: Number of samples to return.

    Returns:
        List of sample dicts. Empty list if buffer is empty.
    """
    if self.is_empty():
        return []

    # Collect all samples across tasks
    all_samples = []
    for task_samples in self._buffers.values():
        all_samples.extend(task_samples)

    # Sample without replacement when possible to maximize diversity
    if batch_size <= len(all_samples):
        return self.rng.sample(all_samples, k=batch_size)
    else:
        return self.rng.choices(all_samples, k=batch_size)

def sample_balanced(self, batch_size: int) -> List[dict]:
    """Sample a batch with equal representation from each stored task.

    Args:
        batch_size: Number of samples to return.

    Returns:
        List of sample dicts. Empty list if buffer is empty.
    """
    if self.is_empty():
        return []

    samples_per_task = max(1, batch_size // self.num_tasks)
    result = []

    for task_samples in self._buffers.values():
        k = min(samples_per_task, len(task_samples))
        result.extend(self.rng.choices(task_samples, k=k))

    # If we need more samples to reach batch_size, sample randomly
    while len(result) < batch_size:
        task_id = self.rng.choice(list(self._buffers.keys()))
        result.append(self.rng.choice(self._buffers[task_id]))

    return result[:batch_size]

def get_task_ids(self) -> List[int]:
    """Return list of task IDs stored in the buffer."""
    return sorted(self._buffers.keys())

def get_task_size(self, task_id: int) -> int:
    """Return number of samples stored for a specific task."""
    return len(self._buffers.get(task_id, []))

def clear(self):
    """Clear all stored samples."""
    self._buffers.clear()
    self._total_samples = 0

def state_dict(self) -> dict:
    """Return serializable state for checkpointing.

    Note: only metadata is serialized — the actual sample tensors are not
    saved (they can be large).  On resume, callers must re-populate the
    buffer by iterating each task's dataset again.
    """
    return {
        "algorithm": self.name,
        "buffer_size_per_task": self.buffer_size_per_task,
        "seed": self.seed,
        "num_tasks": self.num_tasks,
        "total_samples": self._total_samples,
        "task_sizes": {k: len(v) for k, v in self._buffers.items()},
    }

def observe(self, batch: dict, task_id: int) -> None:
    """No-op: ER populates from the full dataset at task-end, not per step.

    See :meth:`populate_from_dataset` for the actual memory update, which
    the trainer invokes from its task-end handler.
    """
    return None

def on_task_end(self, task_id: int) -> None:
    """No-op: the trainer calls :meth:`populate_from_dataset` directly.

    (ER needs the full task dataset object, which a generic no-arg hook
    cannot provide — so the trainer orchestrates the population explicitly.)
    """
    return None

def load_state_dict(self, state: Dict[str, Any]) -> None:
    """Restore metadata from a snapshot produced by :meth:`state_dict`.

    Only hyperparameters are restored (buffer size, seed).  Actual samples
    must be re-populated from the task datasets on resume.
    """
    self.buffer_size_per_task = state.get(
        "buffer_size_per_task", self.buffer_size_per_task
    )
    self.seed = state.get("seed", self.seed)
    self.rng = random.Random(self.seed)
    self._buffers = {}
    self._total_samples = 0

def prepare_training(self):
    """Initialize training state (checkpoints, freezing, distributed setup)."""
    rank = dist.get_rank() if dist.is_initialized() else 0
    seed = self.config.seed + rank if hasattr(self.config, "seed") else rank + 3047
    set_seed(seed)

    self._init_checkpointing()

    freeze_modules = (
        self.config.trainer.freeze_modules
        if hasattr(self.config.trainer, "freeze_modules")
        else None
    )
    self.model = self.freeze_backbones(self.model, freeze_modules=freeze_modules)
    self.print_trainable_parameters(self.model)

    # NOTE: we prepare model and optimizer here, dataloaders are prepared per-task.
    # DeepSpeed requires train_micro_batch_size_per_gpu when no dataloader is passed
    # to accelerator.prepare(), so set it explicitly from config.
    if hasattr(self.accelerator.state, "deepspeed_plugin") and self.accelerator.state.deepspeed_plugin is not None:
        ds_cfg = self.accelerator.state.deepspeed_plugin.deepspeed_config
        micro_bs = self.config.datasets.vla_data.per_device_batch_size
        if ds_cfg.get("train_micro_batch_size_per_gpu") == "auto":
            ds_cfg["train_micro_batch_size_per_gpu"] = micro_bs
        if ds_cfg.get("gradient_accumulation_steps") == "auto":
            ds_cfg["gradient_accumulation_steps"] = getattr(self.config.trainer, "gradient_accumulation_steps", 1)
        if ds_cfg.get("train_batch_size") == "auto":
            grad_acc = ds_cfg.get("gradient_accumulation_steps", 1)
            ds_cfg["train_batch_size"] = micro_bs * self.accelerator.num_processes * grad_acc
    self.model, self.optimizer = self.setup_distributed_training(
        self.accelerator, self.model, self.optimizer
    )

    self._init_wandb()

def train(self, full_dataset, episode_task_map):
    """Execute the continual learning training loop."""
    seq_cfg = get_task_sequence(self.cl_cfg.task_sequence)
    num_tasks = seq_cfg["num_tasks"]
    task_order = seq_cfg.get("task_order", list(range(num_tasks)))
    steps_per_task = self.cl_cfg.steps_per_task
    save_per_task = self.cl_cfg.get("save_checkpoint_per_task", True)

    # Determine which tasks to skip based on completed steps
    start_task_idx = self.completed_steps // steps_per_task
    if start_task_idx > 0:
        logger.info(
            f"Resuming from step {self.completed_steps}: "
            f"skipping {start_task_idx} completed tasks"
        )
        # Rebuild replay buffer from completed tasks
        if self.replay_enabled:
            for skip_idx in range(start_task_idx):
                skip_task_id = task_order[skip_idx]
                _, skip_dataset = build_task_dataloader(
                    full_dataset, skip_task_id, episode_task_map, self.config
                )
                self.replay_buffer.populate_from_dataset(
                    task_id=skip_task_id, dataset=skip_dataset,
                )
            logger.info(
                f"Rebuilt replay buffer for {start_task_idx} completed tasks: "
                f"{self.replay_buffer}"
            )

    self._log_cl_config(num_tasks, steps_per_task)

    # Outer tqdm bar over the CL task sequence (1..num_tasks).
    # disable on non-main ranks to avoid duplicate bars under DeepSpeed.
    task_pbar = tqdm(
        enumerate(task_order),
        total=num_tasks,
        desc="CL tasks",
        disable=not self.accelerator.is_local_main_process,
        initial=start_task_idx,
    )
    for task_idx_in_seq, task_id in task_pbar:
        # Skip already completed tasks
        if task_idx_in_seq < start_task_idx:
            task_pbar.write(
                f"Skipping Task {task_idx_in_seq + 1}/{num_tasks} "
                f"(task_index={task_id}) — already completed"
            )
            continue

        task_pbar.set_postfix(task_id=task_id, step=self.completed_steps)
        task_pbar.write(f"{'='*60}")
        task_pbar.write(
            f"Starting Task {task_idx_in_seq + 1}/{num_tasks} "
            f"(task_index={task_id})"
        )
        task_pbar.write(f"{'='*60}")

        # Build per-task dataloader
        task_dataloader, task_dataset = build_task_dataloader(
            full_dataset, task_id, episode_task_map, self.config
        )

        # Prepare dataloader for distributed training
        self.accelerator.dataloader_config.dispatch_batches = False
        task_dataloader = self.accelerator.prepare(task_dataloader)
        dist.barrier()

        # Reset LR scheduler for new task
        # Unwrap through AcceleratedOptimizer → DeepSpeedZeroOptimizer → base AdamW
        base_optimizer = self.optimizer
        while hasattr(base_optimizer, "optimizer"):
            base_optimizer = base_optimizer.optimizer
        task_lr_scheduler = get_scheduler(
            name=self.config.trainer.lr_scheduler_type,
            optimizer=base_optimizer,
            num_warmup_steps=self.config.trainer.num_warmup_steps,
            num_training_steps=steps_per_task,
            scheduler_specific_kwargs=self.config.trainer.scheduler_specific_kwargs,
        )

        # Train on current task
        self._train_single_task(
            task_id=task_id,
            task_idx_in_seq=task_idx_in_seq,
            num_tasks=num_tasks,
            task_dataloader=task_dataloader,
            lr_scheduler=task_lr_scheduler,
            steps_per_task=steps_per_task,
        )

        # Post-task: populate replay buffer
        if self.replay_enabled:
            logger.info(f"Populating replay buffer for task {task_id}...")
            self.replay_buffer.populate_from_dataset(
                task_id=task_id,
                dataset=task_dataset,
            )
            logger.info(f"Replay buffer state: {self.replay_buffer}")

        # Save checkpoint after task
        if save_per_task:
            self._save_task_checkpoint(task_id, task_idx_in_seq)

        dist.barrier()

    self._finalize_training()

@abstractmethod
def observe(self, batch: dict, task_id: int) -> None:
    """Called every training step with the **current task** batch.

    Replay-based methods (ER, GEM) use this to grow/refresh memory.
    Regularization methods (EWC, SI) use it to accumulate importance
    statistics (Fisher information, path integrals, etc.).
    """

@abstractmethod
def sample(self, batch_size: int) -> Any:
    """Return the algorithm's auxiliary artifact for this step (or None/empty to skip).

    Return shape is algorithm-specific:
      * ER / GEM : `list[dict]`       — raw samples ready to be collated.
      * EWC / SI : `dict[str, Tensor]`— per-parameter regularization terms.
      * LwF      : `dict[str, Tensor]`— teacher logits on the current batch.

    The trainer dispatches on algorithm type to combine this with the
    current-task batch (e.g. mix-in ratio for replay, KL term for LwF).
    """

@abstractmethod
def on_task_end(self, task_id: int) -> None:
    """Hook invoked after a task finishes.

    Typical uses:
    * EWC: snapshot parameters, compute Fisher on current task's dataset.
    * LwF: snapshot the teacher model weights.
    * ER : no-op (reservoir sampling happens online).
    """

@abstractmethod
def state_dict(self) -> dict[str, Any]:
    """Return a JSON-serializable snapshot of the algorithm state.

    This is written alongside model checkpoints so CL state survives
    interruption/restart across tasks.
    """

@abstractmethod
def load_state_dict(self, state: dict[str, Any]) -> None:
    """Restore algorithm state from a dict produced by `state_dict()`."""

def populate_from_dataset(self, task_id: int, dataset: Dataset, num_samples: Optional[int] = None):
    """Store samples from a dataset into the buffer using reservoir sampling.

    Args:
        task_id: Identifier for the task.
        dataset: Dataset to sample from (must support __len__ and __getitem__).
        num_samples: Number of samples to store. Defaults to buffer_size_per_task.
    """
    if num_samples is None:
        num_samples = self.buffer_size_per_task

    n = len(dataset)
    k = min(num_samples, n)

    # Reservoir sampling: select k items from n uniformly at random
    indices = list(range(n))
    self.rng.shuffle(indices)
    selected_indices = sorted(indices[:k])

    samples = []
    for idx in selected_indices:
        try:
            sample = dataset[idx]
            samples.append(sample)
        except Exception as e:
            logger.warning(f"Failed to read sample {idx} for task {task_id}: {e}")
            continue

    # Update buffer
    if task_id in self._buffers:
        self._total_samples -= len(self._buffers[task_id])
    self._buffers[task_id] = samples
    self._total_samples += len(samples)

    logger.info(
        f"Replay buffer: stored {len(samples)} samples for task {task_id} "
        f"(total: {self._total_samples} across {self.num_tasks} tasks)"
    )

def sample(self, batch_size: int) -> List[dict]:
    """Sample a batch uniformly from all stored tasks.

    Args:
        batch_size: Number of samples to return.

    Returns:
        List of sample dicts. Empty list if buffer is empty.
    """
    if self.is_empty():
        return []

    # Collect all samples across tasks
    all_samples = []
    for task_samples in self._buffers.values():
        all_samples.extend(task_samples)

    # Sample without replacement when possible to maximize diversity
    if batch_size <= len(all_samples):
        return self.rng.sample(all_samples, k=batch_size)
    else:
        return self.rng.choices(all_samples, k=batch_size)

def sample_balanced(self, batch_size: int) -> List[dict]:
    """Sample a batch with equal representation from each stored task.

    Args:
        batch_size: Number of samples to return.

    Returns:
        List of sample dicts. Empty list if buffer is empty.
    """
    if self.is_empty():
        return []

    samples_per_task = max(1, batch_size // self.num_tasks)
    result = []

    for task_samples in self._buffers.values():
        k = min(samples_per_task, len(task_samples))
        result.extend(self.rng.choices(task_samples, k=k))

    # If we need more samples to reach batch_size, sample randomly
    while len(result) < batch_size:
        task_id = self.rng.choice(list(self._buffers.keys()))
        result.append(self.rng.choice(self._buffers[task_id]))

    return result[:batch_size]

def get_task_ids(self) -> List[int]:
    """Return list of task IDs stored in the buffer."""
    return sorted(self._buffers.keys())

def get_task_size(self, task_id: int) -> int:
    """Return number of samples stored for a specific task."""
    return len(self._buffers.get(task_id, []))

def clear(self):
    """Clear all stored samples."""
    self._buffers.clear()
    self._total_samples = 0

def state_dict(self) -> dict:
    """Return serializable state for checkpointing.

    Note: only metadata is serialized — the actual sample tensors are not
    saved (they can be large).  On resume, callers must re-populate the
    buffer by iterating each task's dataset again.
    """
    return {
        "algorithm": self.name,
        "buffer_size_per_task": self.buffer_size_per_task,
        "seed": self.seed,
        "num_tasks": self.num_tasks,
        "total_samples": self._total_samples,
        "task_sizes": {k: len(v) for k, v in self._buffers.items()},
    }

def observe(self, batch: dict, task_id: int) -> None:
    """No-op: ER populates from the full dataset at task-end, not per step.

    See :meth:`populate_from_dataset` for the actual memory update, which
    the trainer invokes from its task-end handler.
    """
    return None

def on_task_end(self, task_id: int) -> None:
    """No-op: the trainer calls :meth:`populate_from_dataset` directly.

    (ER needs the full task dataset object, which a generic no-arg hook
    cannot provide — so the trainer orchestrates the population explicitly.)
    """
    return None

def load_state_dict(self, state: Dict[str, Any]) -> None:
    """Restore metadata from a snapshot produced by :meth:`state_dict`.

    Only hyperparameters are restored (buffer size, seed).  Actual samples
    must be re-populated from the task datasets on resume.
    """
    self.buffer_size_per_task = state.get(
        "buffer_size_per_task", self.buffer_size_per_task
    )
    self.seed = state.get("seed", self.seed)
    self.rng = random.Random(self.seed)
    self._buffers = {}
    self._total_samples = 0

@abstractmethod
def observe(self, batch: dict, task_id: int) -> None:
    """Called every training step with the **current task** batch.

    Replay-based methods (ER, GEM) use this to grow/refresh memory.
    Regularization methods (EWC, SI) use it to accumulate importance
    statistics (Fisher information, path integrals, etc.).
    """

@abstractmethod
def sample(self, batch_size: int) -> Any:
    """Return the algorithm's auxiliary artifact for this step (or None/empty to skip).

    Return shape is algorithm-specific:
      * ER / GEM : `list[dict]`       — raw samples ready to be collated.
      * EWC / SI : `dict[str, Tensor]`— per-parameter regularization terms.
      * LwF      : `dict[str, Tensor]`— teacher logits on the current batch.

    The trainer dispatches on algorithm type to combine this with the
    current-task batch (e.g. mix-in ratio for replay, KL term for LwF).
    """

@abstractmethod
def on_task_end(self, task_id: int) -> None:
    """Hook invoked after a task finishes.

    Typical uses:
    * EWC: snapshot parameters, compute Fisher on current task's dataset.
    * LwF: snapshot the teacher model weights.
    * ER : no-op (reservoir sampling happens online).
    """

@abstractmethod
def state_dict(self) -> dict[str, Any]:
    """Return a JSON-serializable snapshot of the algorithm state.

    This is written alongside model checkpoints so CL state survives
    interruption/restart across tasks.
    """

@abstractmethod
def load_state_dict(self, state: dict[str, Any]) -> None:
    """Restore algorithm state from a dict produced by `state_dict()`."""

def populate_from_dataset(self, task_id: int, dataset: Dataset, num_samples: Optional[int] = None):
    """Store samples from a dataset into the buffer using reservoir sampling.

    Args:
        task_id: Identifier for the task.
        dataset: Dataset to sample from (must support __len__ and __getitem__).
        num_samples: Number of samples to store. Defaults to buffer_size_per_task.
    """
    if num_samples is None:
        num_samples = self.buffer_size_per_task

    n = len(dataset)
    k = min(num_samples, n)

    # Reservoir sampling: select k items from n uniformly at random
    indices = list(range(n))
    self.rng.shuffle(indices)
    selected_indices = sorted(indices[:k])

    samples = []
    for idx in selected_indices:
        try:
            sample = dataset[idx]
            samples.append(sample)
        except Exception as e:
            logger.warning(f"Failed to read sample {idx} for task {task_id}: {e}")
            continue

    # Update buffer
    if task_id in self._buffers:
        self._total_samples -= len(self._buffers[task_id])
    self._buffers[task_id] = samples
    self._total_samples += len(samples)

    logger.info(
        f"Replay buffer: stored {len(samples)} samples for task {task_id} "
        f"(total: {self._total_samples} across {self.num_tasks} tasks)"
    )