File size: 5,719 Bytes
4bbff32 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 | """Dual-timescale matrix state with denoise and commit lifetimes.
This module is independent of the decoder and commit policy. It provides the
reference state transition used when replacing the old history and slot paths.
"""
from __future__ import annotations
from dataclasses import dataclass
import torch
from torch import nn
from .gdn2_memory import GDN2Memory
@dataclass(frozen=True)
class GDN2TrajectoryState:
cells: torch.Tensor
row: torch.Tensor
persistent: torch.Tensor
seen: torch.Tensor
def detach(self) -> "GDN2TrajectoryState":
return GDN2TrajectoryState(
self.cells.detach(), self.row.detach(), self.persistent.detach(),
self.seen.detach(),
)
def clear_refill(self, lengths: torch.Tensor, head: torch.Tensor) -> "GDN2TrajectoryState":
batch, canvas = self.seen.shape
if lengths.shape != (batch,) or head.shape != (batch,):
raise ValueError("Commit lengths and canvas heads must be per row.")
physical = torch.arange(canvas, device=self.seen.device)[None, :]
overwritten = ((physical - head[:, None]) % canvas) < lengths[:, None]
return GDN2TrajectoryState(
self.cells.masked_fill(overwritten[..., None, None, None], 0.0),
self.row, self.persistent, self.seen & ~overwritten,
)
def shift(self, lengths: torch.Tensor) -> "GDN2TrajectoryState":
"""Shift a non-ring canvas and zero its newly filled tail."""
batch, canvas = self.seen.shape
if lengths.shape != (batch,):
raise ValueError("Commit lengths must be per row.")
physical = torch.arange(canvas, device=self.seen.device)[None, :] + lengths[:, None]
kept = physical < canvas
selected = physical.clamp_max(canvas - 1)
cells = self.cells.gather(
1, selected[..., None, None, None].expand_as(self.cells)
)
seen = self.seen.gather(1, selected)
return GDN2TrajectoryState(
cells.masked_fill(~kept[..., None, None, None], 0.0),
self.row, self.persistent, seen & kept,
)
class GDN2TrajectoryMemory(nn.Module):
def __init__(self, width: int, *, probes: int = 4,
working_heads: int = 16, working_key: int = 64,
working_value: int = 64, persistent_heads: int = 16,
persistent_key: int = 128, persistent_value: int = 128,
persistent_observation_dim: int | None = None) -> None:
super().__init__()
if probes <= 0:
raise ValueError("Probe count must be positive.")
self.probes = probes
self.cell = GDN2Memory(width, working_heads, working_key, working_value)
self.row = GDN2Memory(width, working_heads, working_key, working_value)
self.persistent = GDN2Memory(width, persistent_heads, persistent_key, persistent_value,
observation_dim=persistent_observation_dim)
self.probe_embed = nn.Embedding(probes, width)
nn.init.normal_(self.probe_embed.weight, std=0.02)
def empty(self, batch: int, canvas: int, *, device: torch.device | None = None
) -> GDN2TrajectoryState:
return GDN2TrajectoryState(
self.cell.empty(batch, canvas, device=device),
self.row.empty(batch, device=device),
self.persistent.empty(batch, device=device),
torch.zeros(batch, canvas, device=device, dtype=torch.bool),
)
def read(self, state: GDN2TrajectoryState, query: torch.Tensor) -> torch.Tensor:
result = (self.cell.read(state.cells, query)
+ self.row.read_shared(state.row, query)
+ self.persistent.read_shared(state.persistent, query))
return torch.where(state.seen[..., None], result, torch.zeros_like(result))
def observe(self, state: GDN2TrajectoryState, observation: torch.Tensor,
live: torch.Tensor, head: torch.Tensor) -> GDN2TrajectoryState:
batch, canvas, width = observation.shape
if live.shape != (batch, canvas) or head.shape != (batch,):
raise ValueError("Observation mask and head have incorrect shapes.")
cells = self.cell.transition(state.cells, observation, live)
seen = state.seen | live
logical_idx = (head[:, None] + torch.arange(canvas, device=head.device)[None, :]) % canvas
logical = observation.gather(1, logical_idx[..., None].expand(-1, -1, width))
logical_live = live.gather(1, logical_idx)
row_state = state.row
for probe in range(self.probes):
lo = canvas * probe // self.probes
hi = canvas * (probe + 1) // self.probes
selected = logical_live[:, lo:hi]
count = selected.sum(dim=1, keepdim=True)
pooled = (logical[:, lo:hi].float() * selected[..., None]).sum(dim=1)
pooled = (pooled / count.clamp_min(1)).to(observation.dtype)
pooled = pooled + self.probe_embed.weight[probe].to(observation.dtype)
row_state = self.row.transition(row_state, pooled, count[:, 0] > 0)
return GDN2TrajectoryState(cells, row_state, state.persistent, seen)
def commit(self, state: GDN2TrajectoryState, experience: torch.Tensor,
valid: torch.Tensor) -> GDN2TrajectoryState:
if experience.shape[:2] != valid.shape:
raise ValueError("Commit experience and validity shapes differ.")
long_state = self.persistent.write_sequence(state.persistent, experience, valid)
return GDN2TrajectoryState(state.cells, state.row, long_state, state.seen)
__all__ = ["GDN2TrajectoryMemory", "GDN2TrajectoryState"]
|