| """FSR / tactile sensor pack — a projector head into the canonical FE2 text space.
|
|
|
| A force-sensitive-resistor (FSR) tactile glove reports a low-dimensional pressure map,
|
| not an optical image. It is the same signal family as the Tremor inertial pack, so it uses
|
| the Tier-3 PROJECTOR-HEAD pattern rather than the vision/adapter path: a small trainable
|
| encoder turns a short window of 32x32 pressure frames into a 512-d tactile feature, and a
|
| trained projector maps that feature into the frozen Qwen base's 2048-d text space.
|
|
|
| The head is a plain ``callable(pressure) -> torch.Tensor[.., 2048]`` returning a RAW (un-
|
| normalized, un-whitened) base-space feature, matching the projector-head contract consumed
|
| by :meth:`fusion_embedding.unified.UnifiedEmbedder._project`. The unified read-out then
|
| L2-normalizes; non-text modalities are never whitened. Wire it in exactly like the released
|
| Tremor head::
|
|
|
| from fusion_embedding.tactile import TactileEncoder, build_fsr_head, build_projector
|
| enc, proj = TactileEncoder(), build_projector()
|
| head = build_fsr_head(enc, proj)
|
| ue = UnifiedEmbedder.from_pretrained(FE2_REPO, projector_heads={"tactile": head})
|
| v = ue.embed_tactile(pressure_window) # [2048], L2-normalized
|
|
|
| Pressure convention: raw frames are uint8 0-255 (548 active sensors on a 32x32 grid);
|
| :func:`preprocess_pressure` scales them to float [0, 1]. The encoder aggregates the F frames
|
| of a window into one tactile event. Three temporal modes are available (``temporal=`` arg):
|
| ``"conv"`` (default) concatenates the per-frame SPATIAL feature maps along channels and fuses
|
| them with a 1x1 convolution, replicating STAG's own frame-fusion; ``"attn"`` runs a small
|
| CLS-token Transformer over the per-frame tokens; ``"maxpool"`` is the earlier frame max-pool.
|
| Two per-frame trunks are available (``depth=`` arg): ``"resnet"`` (default) is a compact
|
| ResNet-style trunk (3x3 stem, three stages of two BasicBlocks) and ``"small"`` is the earlier
|
| three-conv net. A ``flow=`` arg (default True) feeds the trunk a second "tactile flow" channel
|
| (the frame-to-frame pressure difference) alongside the raw frame, following the dual raw+flow
|
| input of ROTConvPCE-mv. Every combination preserves the external contract (input shapes, 512-d
|
| feature), so the modes are directly A/B-able.
|
| """
|
|
|
| from __future__ import annotations
|
|
|
| from typing import Callable, Union
|
|
|
| import numpy as np
|
| import torch
|
| import torch.nn as nn
|
|
|
| GRID = 32
|
| FEAT_DIM = 512
|
| OUT_DIM = 2048
|
| TEMPORAL_MODES = ("conv", "attn", "maxpool")
|
| DEPTHS = ("resnet", "resnet18", "small")
|
|
|
|
|
|
|
|
|
|
|
| def preprocess_pressure(raw) -> torch.Tensor:
|
| """Raw pressure -> float32 tensor ``[F, 32, 32]`` scaled to [0, 1].
|
|
|
| Accepts a numpy array, a nested list, or a torch tensor of shape ``[32, 32]`` (a single
|
| frame) or ``[F, 32, 32]`` (a window). uint8 0-255 input is scaled by 1/255; float input
|
| already in [0, 1] is passed through, and out-of-range float (a raw 0-255 float dump) is
|
| also rescaled. A single frame is promoted to a one-frame window ``[1, 32, 32]``."""
|
| if torch.is_tensor(raw):
|
| arr = raw.detach().cpu().numpy()
|
| else:
|
| arr = np.asarray(raw)
|
| was_integer = np.issubdtype(arr.dtype, np.integer)
|
| t = torch.as_tensor(np.ascontiguousarray(arr), dtype=torch.float32)
|
| if t.dim() == 2:
|
| t = t.unsqueeze(0)
|
| elif t.dim() != 3:
|
| raise ValueError(f"expected pressure of shape [32,32] or [F,32,32], got {list(t.shape)}")
|
| if t.shape[-2:] != (GRID, GRID):
|
| raise ValueError(f"expected a {GRID}x{GRID} pressure grid, got {list(t.shape[-2:])}")
|
|
|
| if was_integer or float(t.max().item()) > 1.5:
|
| t = t / 255.0
|
| return t
|
|
|
|
|
|
|
|
|
|
|
| class _BasicBlock(nn.Module):
|
| """Compact ResNet BasicBlock (two 3x3 convs + identity/1x1-projection shortcut).
|
|
|
| Written inline so the tactile pack keeps its torch-only dependency footprint (no
|
| torchvision). Pre-activation is NOT used; this is the standard post-activation block."""
|
|
|
| def __init__(self, c_in: int, c_out: int, stride: int = 1):
|
| super().__init__()
|
| self.conv1 = nn.Conv2d(c_in, c_out, 3, stride=stride, padding=1, bias=False)
|
| self.bn1 = nn.BatchNorm2d(c_out)
|
| self.conv2 = nn.Conv2d(c_out, c_out, 3, stride=1, padding=1, bias=False)
|
| self.bn2 = nn.BatchNorm2d(c_out)
|
| self.relu = nn.ReLU(inplace=True)
|
| self.down = None
|
| if stride != 1 or c_in != c_out:
|
| self.down = nn.Sequential(nn.Conv2d(c_in, c_out, 1, stride=stride, bias=False),
|
| nn.BatchNorm2d(c_out))
|
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor:
|
| idt = x if self.down is None else self.down(x)
|
| y = self.relu(self.bn1(self.conv1(x)))
|
| y = self.bn2(self.conv2(y))
|
| return self.relu(y + idt)
|
|
|
|
|
| def _resnet_trunk(token_dim: int = 128, in_ch: int = 1) -> nn.Sequential:
|
| """ResNet-style per-frame trunk for 32x32 pressure -> ``[N,token_dim,8,8]``.
|
|
|
| A 3x3 stem (no 7x7 stem, no initial max-pool: a 32x32 taxel grid cannot afford an early
|
| /4 downsample) followed by three stages of two :class:`_BasicBlock` each at 32 / 64 /
|
| ``token_dim`` channels, stride 2 between stages: 32x32 -> 32x32 -> 16x16 -> 8x8. The 8x8
|
| output grid matches the ``"small"`` trunk's, so the temporal stage is interchangeable.
|
| ``in_ch`` is 1 for the raw pressure frame alone and 2 when the tactile-flow channel is
|
| stacked on it; only the stem conv's shape changes, parameter NAMES are identical."""
|
| c1, c2, c3 = 32, 64, token_dim
|
| return nn.Sequential(
|
| nn.Conv2d(in_ch, c1, 3, padding=1, bias=False), nn.BatchNorm2d(c1), nn.ReLU(inplace=True),
|
| _BasicBlock(c1, c1), _BasicBlock(c1, c1),
|
| _BasicBlock(c1, c2, stride=2), _BasicBlock(c2, c2),
|
| _BasicBlock(c2, c3, stride=2), _BasicBlock(c3, c3),
|
| )
|
|
|
|
|
| def _resnet18_trunk(token_dim: int = 512, in_ch: int = 1) -> nn.Sequential:
|
| """Full ResNet-18-width per-frame trunk for 32x32 pressure -> ``[N,token_dim,4,4]``.
|
|
|
| STAG's own classification network uses a stock ResNet-18 (four stages, 64/128/256/512) per
|
| frame; the compact :func:`_resnet_trunk` above is roughly 12x smaller, and in our ablations
|
| encoder capacity, not the readout, was what limited accuracy. This restores the
|
| reference width: 3x3 stem (still no 7x7 / no initial max-pool, since a 32x32 taxel grid cannot
|
| afford an early /4 downsample), then four stages of two :class:`_BasicBlock` each,
|
| 32x32 -> 32x32 -> 16x16 -> 8x8 -> 4x4."""
|
| c1, c2, c3, c4 = 64, 128, 256, token_dim
|
| return nn.Sequential(
|
| nn.Conv2d(in_ch, c1, 3, padding=1, bias=False), nn.BatchNorm2d(c1), nn.ReLU(inplace=True),
|
| _BasicBlock(c1, c1), _BasicBlock(c1, c1),
|
| _BasicBlock(c1, c2, stride=2), _BasicBlock(c2, c2),
|
| _BasicBlock(c2, c3, stride=2), _BasicBlock(c3, c3),
|
| _BasicBlock(c3, c4, stride=2), _BasicBlock(c4, c4),
|
| )
|
|
|
|
|
| class TactileEncoder(nn.Module):
|
| """Per-frame 2D CNN + temporal aggregation -> a 512-d tactile feature.
|
|
|
| Input is a window of 32x32 pressure frames. Shapes accepted: ``[32,32]`` (single frame)
|
| and ``[F,32,32]`` (a window) both return an unbatched ``[512]`` feature; ``[B,F,32,32]``
|
| returns ``[B,512]``. Each frame passes through the trunk independently to a per-frame
|
| spatial feature map; the F maps are then aggregated into one window descriptor and
|
| projected to 512-d.
|
|
|
| Per-frame trunk (``depth=``):
|
| * ``"resnet"`` (default) — :func:`_resnet_trunk`, a compact ResNet (3x3 stem, three
|
| stages of two BasicBlocks, 32/64/``token_dim`` channels), matching STAG's per-frame
|
| ResNet-18 idiom at a size that suits a 32x32 grid.
|
| * ``"small"`` — the earlier three-conv net, kept for A/B (and for loading the
|
| checkpoints trained with it: the ``cnn.*`` parameter names are unchanged).
|
|
|
| Temporal aggregation (``temporal=``):
|
| * ``"conv"`` (default) — STAG's frame fusion: concatenate the F per-frame spatial maps
|
| along the CHANNEL axis and fuse with a single ``Conv2d(frames*token_dim, token_dim,
|
| kernel_size=1)``, then global-average-pool. The fusion conv is defined for the
|
| CONFIGURED ``frames`` (=K) count, so a window whose actual F differs is fitted to K
|
| first: F < K repeat-pads the last frame, F > K takes ``K`` evenly spaced frames
|
| (a stride subsample, not a head truncation, so a long recording-level window is still
|
| covered end to end).
|
| * ``"attn"`` — prepend a learned CLS token to the F frame tokens, add a learned
|
| per-position embedding, run ``n_layers`` ``nn.TransformerEncoderLayer``
|
| (``batch_first``, ``n_heads`` heads, GELU, ``dropout``) over the ``F+1`` tokens, and
|
| read the CLS output. Works for ``F==1`` (a length-2 sequence with CLS).
|
| * ``"maxpool"`` — the earlier max-pool over the F frame tokens, kept for A/B.
|
|
|
| For ``"attn"``, frame position is embedded up to ``max_frames`` (default 64, covering the
|
| recording-level eval cap); a longer window is truncated to ``max_frames`` frames.
|
|
|
| Tactile flow (``flow=``, default True):
|
| ROTConvPCE-mv (Cao et al.) feeds its tactile network BOTH the raw frames and a "tactile
|
| flow" field, on the argument that a per-frame trunk otherwise throws away the grasp
|
| DYNAMICS: how the contact pattern moves and builds between frames. Here the flow channel
|
| is the first-order temporal difference ``frame[i] - frame[i-1]`` (zeros for ``i == 0``,
|
| which has no preceding frame), stacked as a second input channel so the trunk sees
|
| ``[N,2,32,32]`` instead of ``[N,1,32,32]``. The difference is taken on the [0,1]-scaled
|
| frames and is NOT renormalized, so it lives in [-1, 1] and its magnitude stays comparable
|
| to the raw channel. Only the stem conv's ``in_channels`` changes; ``flow=False`` restores
|
| the exact single-channel behavior with identical parameter names and shapes, so
|
| checkpoints trained before this option still load.
|
| """
|
|
|
| def __init__(self, temporal: str = "conv", token_dim: int = 128,
|
| n_layers: int = 2, n_heads: int = 4, max_frames: int = 64,
|
| dropout: float = 0.1, depth: str = "resnet", frames: int = 8,
|
| flow: bool = True):
|
| super().__init__()
|
| if temporal not in TEMPORAL_MODES:
|
| raise ValueError(f"temporal must be one of {TEMPORAL_MODES}, got {temporal!r}")
|
| if depth not in DEPTHS:
|
| raise ValueError(f"depth must be one of {DEPTHS}, got {depth!r}")
|
| if int(frames) < 1:
|
| raise ValueError(f"frames must be >= 1, got {frames!r}")
|
| self.temporal = temporal
|
| self.depth = depth
|
| self.frames = int(frames)
|
| self.token_dim = token_dim
|
| self.max_frames = max_frames
|
| self.flow = bool(flow)
|
| in_ch = 2 if self.flow else 1
|
| self.in_ch = in_ch
|
| if depth == "resnet18":
|
|
|
|
|
| self.token_dim = token_dim = max(int(token_dim), 512)
|
| self.trunk = _resnet18_trunk(token_dim, in_ch=in_ch)
|
| elif depth == "resnet":
|
| self.trunk = _resnet_trunk(token_dim, in_ch=in_ch)
|
| else:
|
| self.cnn = nn.Sequential(
|
| nn.Conv2d(in_ch, 32, 3, padding=1), nn.BatchNorm2d(32), nn.ReLU(inplace=True),
|
| nn.MaxPool2d(2),
|
| nn.Conv2d(32, 64, 3, padding=1), nn.BatchNorm2d(64), nn.ReLU(inplace=True),
|
| nn.MaxPool2d(2),
|
| nn.Conv2d(64, token_dim, 3, padding=1), nn.BatchNorm2d(token_dim),
|
| nn.ReLU(inplace=True),
|
| nn.AdaptiveAvgPool2d(1),
|
| )
|
| if temporal == "conv":
|
| self.fuse = nn.Conv2d(self.frames * token_dim, token_dim, kernel_size=1)
|
| if temporal == "attn":
|
| self.cls = nn.Parameter(torch.zeros(1, 1, token_dim))
|
| self.pos = nn.Parameter(torch.zeros(1, max_frames + 1, token_dim))
|
| nn.init.trunc_normal_(self.cls, std=0.02)
|
| nn.init.trunc_normal_(self.pos, std=0.02)
|
| layer = nn.TransformerEncoderLayer(
|
| d_model=token_dim, nhead=n_heads, dim_feedforward=token_dim * 4,
|
| dropout=dropout, activation="gelu", batch_first=True, norm_first=True)
|
| self.transformer = nn.TransformerEncoder(
|
| layer, num_layers=n_layers, enable_nested_tensor=False)
|
| self.fc = nn.Linear(token_dim, FEAT_DIM)
|
| self.act = nn.GELU()
|
|
|
| def _stack_flow(self, x: torch.Tensor) -> torch.Tensor:
|
| """``[B,F,H,W]`` -> per-frame trunk input ``[B*F,in_ch,H,W]``.
|
|
|
| With ``flow=False`` this is just the raw frame as one channel. With ``flow=True`` a
|
| second channel carries the tactile flow ``frame[i] - frame[i-1]``, and frame 0 (which
|
| has no predecessor) gets an all-zero flow channel. The difference is taken on the
|
| already-[0,1]-scaled frames and is deliberately left unnormalized."""
|
| B, Fr, H, W = x.shape
|
| if not self.flow:
|
| return x.reshape(B * Fr, 1, H, W)
|
| if Fr > 1:
|
| diff = torch.cat([torch.zeros_like(x[:, :1]), x[:, 1:] - x[:, :-1]], dim=1)
|
| else:
|
| diff = torch.zeros_like(x)
|
| return torch.stack([x, diff], dim=2).reshape(B * Fr, 2, H, W)
|
|
|
| def _spatial(self, flat: torch.Tensor) -> torch.Tensor:
|
| """Per-frame trunk on a flattened ``[N,in_ch,32,32]`` batch -> spatial map ``[N,C,h,w]``.
|
|
|
| For ``depth="small"`` this is the three-conv net WITHOUT its trailing
|
| ``AdaptiveAvgPool2d(1)``, so the pooled token path stays numerically identical to the
|
| earlier encoder while the conv-fusion path gets the full 8x8 map."""
|
| if self.depth in ("resnet", "resnet18"):
|
| return self.trunk(flat)
|
| for m in list(self.cnn)[:-1]:
|
| flat = m(flat)
|
| return flat
|
|
|
| def _fit_frames(self, maps: torch.Tensor) -> torch.Tensor:
|
| """Fit the frame axis of ``[B,F,C,h,w]`` to the configured ``frames`` (=K) count.
|
|
|
| F == K passes through. F < K repeat-pads the last frame. F > K takes K evenly spaced
|
| frames, so a long window (recording-level eval feeds up to RECORDING_CAP frames) is
|
| still sampled across its whole span rather than truncated to its head."""
|
| Fr = maps.shape[1]
|
| K = self.frames
|
| if Fr == K:
|
| return maps
|
| if Fr < K:
|
| pad = maps[:, -1:].expand(-1, K - Fr, -1, -1, -1)
|
| return torch.cat([maps, pad], dim=1)
|
| idx = torch.linspace(0, Fr - 1, K, device=maps.device).round().long()
|
| return maps.index_select(1, idx)
|
|
|
| def forward(self, x) -> torch.Tensor:
|
| if not torch.is_tensor(x):
|
| x = torch.as_tensor(x, dtype=torch.float32)
|
| x = x.float()
|
| squeeze = False
|
| if x.dim() == 2:
|
| x = x.unsqueeze(0).unsqueeze(0)
|
| squeeze = True
|
| elif x.dim() == 3:
|
| x = x.unsqueeze(0)
|
| squeeze = True
|
| elif x.dim() != 4:
|
| raise ValueError(
|
| f"expected pressure [32,32], [F,32,32] or [B,F,32,32], got {list(x.shape)}")
|
| B, Fr, H, W = x.shape
|
| maps = self._spatial(self._stack_flow(x))
|
| C, h, w = maps.shape[1], maps.shape[2], maps.shape[3]
|
| if self.temporal == "conv":
|
| maps = self._fit_frames(maps.reshape(B, Fr, C, h, w))
|
| fused = self.fuse(maps.reshape(B, self.frames * C, h, w))
|
| z = nn.functional.adaptive_avg_pool2d(fused, 1).reshape(B, C)
|
| z = self.act(self.fc(z))
|
| return z.squeeze(0) if squeeze else z
|
| z = nn.functional.adaptive_avg_pool2d(maps, 1)
|
| z = z.reshape(B, Fr, self.token_dim)
|
| if self.temporal == "maxpool":
|
| z = z.max(dim=1).values
|
| else:
|
| if Fr > self.max_frames:
|
| z = z[:, :self.max_frames]
|
| Fr = self.max_frames
|
| cls = self.cls.expand(B, -1, -1)
|
| seq = torch.cat([cls, z], dim=1)
|
| seq = seq + self.pos[:, :Fr + 1]
|
| seq = self.transformer(seq)
|
| z = seq[:, 0]
|
| z = self.act(self.fc(z))
|
| return z.squeeze(0) if squeeze else z
|
|
|
|
|
| def build_projector(in_dim: int = FEAT_DIM, out_dim: int = OUT_DIM) -> nn.Sequential:
|
| """The tactile projector head (same arch idiom as the Tremor 'ln' projector):
|
| LayerNorm(512) -> Linear(512,1024) -> GELU -> Dropout(0.3) -> Linear(1024,2048)."""
|
| return nn.Sequential(
|
| nn.LayerNorm(in_dim), nn.Linear(in_dim, 1024), nn.GELU(),
|
| nn.Dropout(0.3), nn.Linear(1024, out_dim),
|
| )
|
|
|
|
|
| def build_fsr_head(encoder: TactileEncoder, projector: nn.Module) -> Callable:
|
| """Wrap ``encoder`` + ``projector`` into a projector-head callable.
|
|
|
| Returns ``head(pressure) -> [.., 2048]`` a RAW (un-normalized) base-space feature, ready
|
| to pass as ``projector_heads={"tactile": head}`` to ``UnifiedEmbedder.from_pretrained``
|
| (or ``projectors={"tactile": head}`` to ``UnifiedEmbedder(...)``). Non-tensor input (a
|
| numpy array / list) is run through :func:`preprocess_pressure` first; a tensor is used
|
| as-is (training passes a pre-scaled ``[B,1,32,32]`` batch straight through)."""
|
|
|
| def head(pressure: Union[torch.Tensor, np.ndarray, list]) -> torch.Tensor:
|
| x = pressure if torch.is_tensor(pressure) else preprocess_pressure(pressure)
|
| p = next(projector.parameters(), None)
|
| if p is not None:
|
| x = x.to(p.device)
|
| return projector(encoder(x))
|
|
|
| return head
|
|
|
