Base¶

def backbone_forward(
    self,
    embeds: torch.Tensor,
    token_type: torch.Tensor,
    cache: dict[str, Any] | None = None,
    use_cache: bool = False,
    cache_position: torch.Tensor | None = None,
    **kwargs: Any,
) -> tuple[torch.Tensor, dict[str, Any] | None]:
    """Run backbone; return (hidden states [B, T, D], cache dict or None)."""
    raise NotImplementedError("Subclasses must implement backbone_forward.")

def head(self, h: torch.Tensor, batch_size: tuple[int, int]) -> TensorDict:
    """Run all enabled heads on pooled step representations.

    Args:
        h: Step representations ``[B, S, D]``.
        batch_size: ``(B, S)`` used to set the TensorDict batch dimensions.

    Returns:
        TensorDict ``[B, S]`` with a key for each enabled head.
        Logit heads (``sp``, ``dqn``, ``dqn_target``, ``sv``) have shape
        ``[B, S, A]``; vector heads (``vec_dqn``, ``vec_dqn_target``) have
        shape ``[B, S, A, D]``.  Disabled heads are absent.
    """
    tensors: dict[str, torch.Tensor] = {}
    if self.sp_head is not None:
        tensors["sp"] = self.sp_head(h)
    if self.dqn_head is not None:
        tensors["dqn"] = self.dqn_head(h)
        tensors["dqn_target"] = self.dqn_head.target_forward(h)
    if self.vec_dqn_head is not None:
        tensors["vec_dqn"] = self.vec_dqn_head(h)
        tensors["vec_dqn_target"] = self.vec_dqn_head.target_forward(h)
    if self.sv_head is not None:
        tensors["sv"] = self.sv_head(h)
    return TensorDict(tensors, batch_size=batch_size)

def forward(
    self,
    step_stream: TensorDict,
    cache: dict[str, Any] | None = None,
    use_cache: bool = False,
    cache_position: torch.Tensor | None = None,
) -> tuple[TensorDict, dict[str, Any] | None]:
    """Run a full forward pass over a batch of step sequences.

    Args:
        step_stream: TensorDict ``[B, S]`` of step records (observations,
            actions, rewards, etc. as configured by the embedder).
        cache: KV-cache dict from a previous call, or ``None`` for a full
            prefill.  Only meaningful when ``use_cache=True``.
        use_cache: If ``True`` the backbone returns an updated cache that
            can be passed back on the next call for incremental decoding.
        cache_position: Token position indices ``[T]`` for incremental
            decoding; leave ``None`` for full prefill.

    Returns:
        out: TensorDict ``[B, S]`` with one key per enabled head.
             Logit heads — ``sp``, ``dqn``, ``dqn_target``, ``sv`` —
             have shape ``[B, S, A]``.  Vector heads — ``vec_dqn``,
             ``vec_dqn_target`` — have shape ``[B, S, A, D]``.
             Use ``get_action`` or ``vec_dqn_scores`` for the vector heads.
        cache: Updated KV-cache dict, or ``None`` when ``use_cache=False``.
    """
    B, S = int(step_stream.batch_size[0]), int(step_stream.batch_size[1])

    embeds, token_type = self.embedder(step_stream)
    h, new_cache = self.backbone_forward(
        embeds=embeds,
        token_type=token_type,
        cache=cache,
        use_cache=use_cache,
        cache_position=cache_position,
    )

    # Take the last token per step as the step representation
    T = self.embedder.tokens_per_step
    h_step = h.view(B, S, T, self.hidden_dim)[:, :, -1, :]  # [B, S, D]

    return self.head(h_step.float(), batch_size=(B, S)), new_cache

def polyak_update(self, dqn_tau: float = 0.0, vec_dqn_tau: float = 0.0) -> None:
    """Soft-update all target heads toward their online counterparts."""
    if self.dqn_head is not None:
        self.dqn_head.polyak_update(tau=dqn_tau)
    if self.vec_dqn_head is not None:
        self.vec_dqn_head.polyak_update(tau=vec_dqn_tau)

def get_action(
    self,
    out: TensorDict,
    temperature: float = 1.0,
    num_actions: int | None = None,
) -> torch.Tensor:
    """Select an action from model output at the last step position.

    Uses ``self.action_head``, which is set at construction time (or
    auto-detected from enabled heads).

    Args:
        out: Model output TensorDict ``[B, S, ...]``.
        temperature: Sampling temperature. ``0.0`` → greedy argmax;
                     ``> 0`` → softmax sampling.
        num_actions: If given, trim scores to the first ``num_actions``
                     columns before sampling (useful when the environment
                     has fewer actions than the model's maximum).

    Returns:
        ``[B]`` int64 tensor of selected actions.
    """
    raw = out[self.action_head][:, -1]  # [B, A] or [B, A, D] for vec_dqn
    scores: torch.Tensor = vec_dqn_scores(raw) if self.action_head == "vec_dqn" else raw  # [B, A]
    if num_actions is not None:
        scores = scores[:, :num_actions]
    if temperature == 0.0:
        return scores.argmax(dim=-1)
    scores = scores - scores.max(dim=-1, keepdim=True).values  # numerical stability
    probs = F.softmax(scores / temperature, dim=-1)
    return torch.multinomial(probs, num_samples=1).squeeze(-1)