src/forge_lab/models/pytorch_model.py

from __future__ import annotations

import torch
from torch import nn

from .numpy_runtime import DEFAULT_CENTER, DEFAULT_SCALE, WINDOW_SIZE


class TelemetryFeatureExtractor(nn.Module):
    """Export-friendly implementation of the six physical window features."""

    def __init__(self) -> None:
        super().__init__()
        x = torch.arange(WINDOW_SIZE, dtype=torch.float32)
        x_centered = x - x.mean()
        self.register_buffer("x_centered", x_centered)
        self.register_buffer("slope_denominator", torch.square(x_centered).sum())

    def forward(self, telemetry: torch.Tensor) -> torch.Tensor:
        force = telemetry[:, :, 0]
        deviation = telemetry[:, :, 1]
        slope = (force * self.x_centered).sum(dim=1) / self.slope_denominator
        shift = force[:, -10:].mean(dim=1) - force[:, :10].mean(dim=1)
        std = force.std(dim=1, correction=0)
        max_deviation = deviation.amax(dim=1)
        last_deviation = deviation[:, -1]
        force_range = force.amax(dim=1) - force.amin(dim=1)
        return torch.stack(
            [slope, shift, std, max_deviation, last_deviation, force_range],
            dim=1,
        )


class TinyDriftNet(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.features = TelemetryFeatureExtractor()
        self.register_buffer("feature_center", torch.tensor(DEFAULT_CENTER.copy()))
        self.register_buffer("feature_scale", torch.tensor(DEFAULT_SCALE.copy()))
        self.classifier = nn.Linear(6, 1)

    def forward(self, telemetry: torch.Tensor) -> torch.Tensor:
        features = self.features(telemetry)
        normalized = (features - self.feature_center) / self.feature_scale
        return torch.sigmoid(self.classifier(normalized)).squeeze(1)