Embedding¶

def __init__(
    self,
    hidden_dim: int,
    max_num_actions: int,
    max_num_obs_continuous: int,
    max_num_obs_discrete: int,
    max_num_obs_image: int,
    max_num_time_steps: int,
    include_action_token: bool,
    include_done_token: bool,
    include_reward_token: bool,
    include_obs_continuous: bool,
    include_obs_discrete: bool,
    include_obs_image: bool,
    include_time_token: bool,
    include_type_token: bool,
    token_data_len: int,
    num_compute_tokens: int = 0,
    concat_modalities: bool = False,
    fourier_in_min: float = -1.0,
    fourier_in_max: float = 1.0,
    std: float = 0.02,
) -> None:
    super().__init__()

    _size_checks = [
        ("include_action_token", include_action_token, "max_num_actions", max_num_actions),
        ("include_obs_continuous", include_obs_continuous, "max_num_obs_continuous", max_num_obs_continuous),
        ("include_obs_discrete", include_obs_discrete, "max_num_obs_discrete", max_num_obs_discrete),
        ("include_obs_image", include_obs_image, "max_num_obs_image", max_num_obs_image),
        ("include_time_token", include_time_token, "max_num_time_steps", max_num_time_steps),
    ]
    for inc_name, inc_val, size_name, size_val in _size_checks:
        if inc_val and int(size_val) <= 0:
            raise ValueError(f"{inc_name} is True but {size_name} is {size_val} (must be > 0).")

    if int(num_compute_tokens) < 0:
        raise ValueError(f"num_compute_tokens must be >= 0, got {num_compute_tokens}.")

    self.hidden_dim = int(hidden_dim)
    self.include_action_token = bool(include_action_token)
    self.include_time_token = bool(include_time_token)
    self.include_done_token = bool(include_done_token)
    self.include_reward_token = bool(include_reward_token)
    self.include_obs_continuous = bool(include_obs_continuous)
    self.include_obs_discrete = bool(include_obs_discrete)
    self.include_obs_image = bool(include_obs_image)
    self.include_type_token = bool(include_type_token)
    self.num_compute_tokens = int(num_compute_tokens)
    self.concat_modalities = bool(concat_modalities)
    self.token_data_len = int(token_data_len)

    # Count active data modalities (determines tokens_per_step in concat mode).
    self._num_data_modalities: int = sum([
        include_time_token,
        include_action_token,
        include_obs_continuous,
        include_obs_discrete,
        include_obs_image,
        include_reward_token,
        include_done_token,
    ])

    # Compute tokens_per_step.
    T = self.token_data_len
    K = self.num_compute_tokens
    if concat_modalities:
        data_slots = self._num_data_modalities * T
    else:
        data_slots = T
    self.tokens_per_step: int = data_slots + K

    # Shared type embedding (only used internally, not in returned token_types).
    self.type_embedder = TypeEmbedder(hidden_dim=hidden_dim, token_data_len=T, embedding_std=std)

    # Action (optional)
    self.action_embedder = (
        ActionEmbedder(hidden_dim=hidden_dim, token_data_len=T, max_num_actions=int(max_num_actions), embedding_std=std)
        if include_action_token else None
    )

    # Time (optional)
    self.time_embedder = (
        TimeEmbedder(hidden_dim=hidden_dim, token_data_len=T, max_num_time_steps=int(max_num_time_steps), embedding_std=std)
        if include_time_token else None
    )

    # Done (optional)
    self.done_embedder = (
        DoneEmbedder(hidden_dim=hidden_dim, token_data_len=T, embedding_std=std)
        if include_done_token else None
    )

    # Reward (optional)
    self.reward_embedder = (
        RewardEmbedder(hidden_dim=hidden_dim, token_data_len=T, in_min=fourier_in_min, in_max=fourier_in_max, embedding_std=std)
        if include_reward_token else None
    )

    # Continuous obs (optional)
    if include_obs_continuous:
        self.obs_continuous_embedder = ObsContinuousEmbedder(
            hidden_dim=hidden_dim, max_num_obs=int(max_num_obs_continuous), token_data_len=T,
            in_min=fourier_in_min, in_max=fourier_in_max, embedding_std=std,
        )
    else:
        self.obs_continuous_embedder = None

    # Discrete obs (optional)
    self.obs_discrete_embedder = (
        ObsDiscreteEmbedder(
            hidden_dim=hidden_dim, max_num_obs=int(max_num_obs_discrete), token_data_len=T, embedding_std=std
        )
        if include_obs_discrete else None
    )

    # Image obs (optional)
    self.obs_image_embedder = (
        ObsImageEmbedder(
            hidden_dim=hidden_dim, max_num_obs=int(max_num_obs_image), token_data_len=T, embedding_std=std
        )
        if include_obs_image else None
    )

    # Compute tokens — one shared learned embedding per compute slot, broadcast over (B, S).
    if K > 0:
        self.compute_embed = nn.Parameter(torch.empty(K, int(hidden_dim)))
        nn.init.normal_(self.compute_embed, std=std)
    else:
        self.compute_embed = None  # type: ignore[assignment]

def __init__(self, hidden_dim: int, token_data_len: int, max_num_actions: int, embedding_std: float = 0.02) -> None:
    super().__init__()
    self.embed = ScaledEmbedding(
        num_embeddings=max_num_actions, embedding_dim=hidden_dim * token_data_len, scale=embedding_std
    )

def __init__(
    self,
    hidden_dim: int,
    token_data_len: int,
    in_min: float,
    in_max: float,
    embedding_std: float = 0.02,
) -> None:
    super().__init__()
    rff_scale = embedding_std / 0.5 ** 0.5
    self.rff = RandomFourierFeatures(
        num_features=hidden_dim * token_data_len, in_min=in_min, in_max=in_max, output_scale=rff_scale
    )

def __init__(self, hidden_dim: int, token_data_len: int, embedding_std: float = 0.02) -> None:
    super().__init__()
    self.embed = ScaledEmbedding(num_embeddings=3, embedding_dim=hidden_dim * token_data_len, scale=embedding_std)

def __init__(
    self, hidden_dim: int, token_data_len: int, max_num_time_steps: int, embedding_std: float = 0.02
) -> None:
    super().__init__()
    self.embed = ScaledEmbedding(
        num_embeddings=max_num_time_steps, embedding_dim=hidden_dim * token_data_len, scale=embedding_std
    )

def __init__(
    self,
    hidden_dim: int,
    max_num_obs: int,
    token_data_len: int,
    in_min: float,
    in_max: float,
    embedding_std: float = 0.02,
) -> None:
    super().__init__()
    self.max_num_obs = max_num_obs
    # cos has std ≈ 1/√2; sum over max_num_obs dims grows by √max_num_obs → divide both out
    rff_scale = embedding_std / (0.5 ** 0.5 * max_num_obs ** 0.5)
    self.rff = RandomFourierFeatures(
        num_features=hidden_dim * token_data_len, in_min=in_min, in_max=in_max,
        num_freq_sets=max_num_obs, output_scale=rff_scale,
    )

def __init__(
    self,
    hidden_dim: int,
    max_num_obs: int,
    token_data_len: int,
    input_std: float = 1.0,
    embedding_std: float = 0.02,
) -> None:
    super().__init__()
    self.max_num_obs = max_num_obs
    # Kaiming uniform for in_features=1 has std = 1/√3 (Uniform[-1,1]).
    # ScaledPosLinear multiplies those weights by scale, so per-dim output std =
    # scale × (1/√3) × input_std.  Divide scale by (1/√3) × √max_num_obs to
    # hit embedding_std after summing max_num_obs independent dims.
    _kaiming_std = 3.0 ** -0.5
    self.projs = ScaledPosLinear(
        num_positions=max_num_obs,
        in_features=1,
        out_features=hidden_dim * token_data_len,
        scale=embedding_std / (_kaiming_std * input_std),
    )

def __init__(
    self,
    hidden_dim: int,
    max_num_obs: int,
    token_data_len: int,
    embedding_std: float = 0.02,
) -> None:
    super().__init__()
    self.max_num_obs = max_num_obs
    # Summing max_num_obs independent N(0, scale) rows inflates std by √max_num_obs.
    # Note: if obs values are uniform integers in [0, max_num_obs-1], collisions further
    # inflate by ≈√(2·max_num_obs-1)/√max_num_obs; √max_num_obs is a reasonable approximation.
    self.embed = ScaledEmbedding(
        num_embeddings=max_num_obs, embedding_dim=hidden_dim * token_data_len,
        scale=embedding_std / max_num_obs ** 0.5,
    )

def __init__(
    self,
    hidden_dim: int,
    max_num_obs: int,
    token_data_len: int,
    embedding_std: float = 0.02,
) -> None:
    super().__init__()
    self.max_num_obs = max_num_obs
    # pixel_norm_std: std of NormalizedPixel output  ≈ std of Uniform[-1,1] = 1/√3
    # _kaiming_std:   std of Kaiming-uniform weights for in_features=1 = 1/√3
    # per-dim output std = scale × _kaiming_std × pixel_norm_std
    # Divide scale by that product × √max_num_obs to hit embedding_std after summing.
    pixel_norm_std = 3.0 ** -0.5
    _kaiming_std = 3.0 ** -0.5
    self.norm = NormalizedPixel()
    self.projs = ScaledPosLinear(
        num_positions=max_num_obs, in_features=1, out_features=hidden_dim * token_data_len,
        scale=embedding_std / (_kaiming_std * pixel_norm_std * max_num_obs ** 0.5),
    )

def __init__(self, hidden_dim: int, token_data_len: int, embedding_std: float = 0.02) -> None:
    super().__init__()
    self.embed = ScaledEmbedding(num_embeddings=8, embedding_dim=hidden_dim * token_data_len, scale=embedding_std)

def __init__(
    self,
    num_features: int,
    in_min: float = 1e-2,
    in_max: float = 1e2,
    num_freq_sets: int = 1,
    output_scale: float = 1.0,
    dtype: torch.dtype = torch.float32,
):
    super().__init__()
    if num_features <= 0:
        raise ValueError("num_features must be > 0")
    if in_min <= 0 or in_max <= 0:
        raise ValueError("in_min and in_max must be > 0")
    if in_min >= in_max:
        raise ValueError("in_min must be < in_max")
    if num_freq_sets < 1:
        raise ValueError("num_freq_sets must be >= 1")
    # ω = 1/x so that one cycle spans x units of input: ω_min = 1/in_max, ω_max = 1/in_min
    log_w_min = math.log(1.0 / in_max)
    log_w_max = math.log(1.0 / in_min)
    freqs = torch.empty(num_freq_sets, num_features, dtype=dtype).uniform_(log_w_min, log_w_max).exp()
    phases = torch.empty(num_freq_sets, num_features, dtype=dtype).uniform_(0.0, 2.0 * math.pi)
    self.register_buffer("freqs", freqs, persistent=True)
    self.register_buffer("phases", phases, persistent=True)
    self.weight = nn.Parameter(torch.full((num_freq_sets, num_features), float(output_scale), dtype=dtype))

def __init__(self, num_embeddings: int, embedding_dim: int, scale: float = 1.0, **kwargs) -> None:
    super().__init__(num_embeddings, embedding_dim, **kwargs)
    self.weight.data.mul_(scale)

def __init__(
    self,
    in_features: int,
    out_features: int,
    scale: float,
    bias: bool = True,
    device: torch.device | str | None = None,
    dtype: torch.dtype | None = None,
) -> None:
    super().__init__(in_features, out_features, bias=bias, device=device, dtype=dtype)
    sc = float(scale)
    if sc < 0.0:
        raise ValueError(f"scale must be >= 0, got {scale!r}.")
    with torch.no_grad():
        self.weight.mul_(sc)
        if self.bias is not None:
            self.bias.mul_(sc)

def __init__(
    self,
    num_positions: int,
    in_features: int,
    out_features: int,
    device: torch.device | str | None = None,
    dtype: torch.dtype | None = None,
) -> None:
    super().__init__()
    self.in_features = in_features
    self.out_features = out_features
    factory = {"device": device, "dtype": dtype}
    self.weight = nn.Parameter(torch.empty(num_positions, out_features, in_features, **factory))
    self.bias = nn.Parameter(torch.empty(num_positions, out_features, **factory))
    nn.init.kaiming_uniform_(self.weight.view(-1, in_features), a=5**0.5)
    nn.init.zeros_(self.bias)

def __init__(
    self,
    num_positions: int,
    in_features: int,
    out_features: int,
    scale: float = 1.0,
    device: torch.device | str | None = None,
    dtype: torch.dtype | None = None,
) -> None:
    super().__init__(
        num_positions=num_positions,
        in_features=in_features,
        out_features=out_features,
        device=device,
        dtype=dtype,
    )
    if scale <= 0.0:
        raise ValueError(f"scale must be > 0, got {scale!r}.")
    with torch.no_grad():
        self.weight.mul_(scale)
        if self.bias is not None:
            self.bias.mul_(scale)