"""Transformers loading and structured-state inference for Faynt. The encoder, Transformer, controller head, codec, and rolling cache are the benchmark-pinned native implementation shipped alongside this adapter. """ from dataclasses import dataclass from typing import Any import torch from transformers import PreTrainedModel from transformers.utils import ModelOutput from .configuration_faynt import FayntConfig from .controller_codec import ControllerLabels, ControllerState from .faynt_native import ( ControllerHeadOutput, KVCache, MeleePolicy, _autocast_context, _controller_labels, ) from .tensor_batch import ( ButtonsBatch, ControllerBatch, FoDPlatformsBatch, GameStateBatch, ItemsBatch, NanaBatch, PlayerBatch, RandallBatch, StickBatch, ) @dataclass class FayntModelOutput(ModelOutput): """Autoregressive controller predictions and optional updated KV cache. ``logits`` contains ``buttons`` and ``main_stick`` tensors. Main-stick logits condition on the predicted, sampled, or supplied next buttons. ``controller`` is the native decoded ControllerState. Its ``as_packed_tensor()`` method returns logical controller values in the documented native codec order. ``labels`` contains the two categories. """ loss: torch.Tensor | None = None logits: dict[str, torch.Tensor] | None = None labels: ControllerLabels | None = None controller: ControllerState | None = None hidden_states: torch.Tensor | None = None cache: KVCache | None = None class FayntModel(PreTrainedModel): """A Melee policy accepting structured game states and current controllers. ``forward`` follows the native deterministic/teacher-forced interface. ``sample`` samples autoregressive commands at the chosen temperature. ``step`` consumes a single frame and mutates a bounded rolling KV cache. The complete native policy remains available as ``model.policy``. """ config_class = FayntConfig base_model_prefix = "policy" main_input_name = "game_state_t" _supports_sdpa = True _no_split_modules = ["TransformerBlock"] def __init__(self, config: FayntConfig): super().__init__(config) self.policy = MeleePolicy(config.to_native_config()) self.post_init() def _init_weights(self, module: torch.nn.Module) -> None: # Native constructors own the architecture's initialization. Loaded # checkpoints overwrite every trainable tensor without remapping it. pass @property def codec(self): return self.policy.codec @staticmethod def _output( value: ControllerHeadOutput, cache: KVCache | None, loss: torch.Tensor | None = None, ) -> FayntModelOutput: return FayntModelOutput( loss=loss, logits=value.logits, labels=value.labels, controller=value.controller_state, hidden_states=value.hidden, cache=cache, ) def forward( self, game_state_t: Any, controller_t: Any, reset_mask: torch.Tensor | None = None, padding_mask: torch.Tensor | None = None, cache: KVCache | None = None, use_cache: bool = False, controller_t_plus_1: Any | None = None, player_name: Any | None = None, valid_position_mask: torch.Tensor | None = None, return_dict: bool | None = None, ) -> FayntModelOutput | tuple: """Process ``[batch, time, ...]`` inputs with already aligned targets. Supplying ``controller_t_plus_1`` enables teacher forcing. Supplying ``valid_position_mask`` as well computes summed controller-head NLL. Game boundaries belong in ``reset_mask``; real frames are True in ``padding_mask``. The native model never shifts targets internally. """ value, next_cache = self.policy( game_state_t=game_state_t, controller_t=controller_t, reset_mask=reset_mask, padding_mask=padding_mask, cache=cache, use_cache=use_cache, controller_t_plus_1=controller_t_plus_1, player_name=player_name, ) loss = None if valid_position_mask is not None: if controller_t_plus_1 is None: raise ValueError("Loss requires aligned controller_t_plus_1 targets") loss, _ = self.policy.loss(value, controller_t_plus_1, valid_position_mask) result = self._output(value, next_cache, loss) use_dict = self.config.use_return_dict if return_dict is None else return_dict return result if use_dict else result.to_tuple() def init_cache( self, batch_size: int, *, device: torch.device | str | None = None, dtype: torch.dtype | None = None, ) -> KVCache: """Allocate the native rolling cache with this model's context length.""" return self.policy.backbone.init_cache(batch_size, device=device, dtype=dtype) def sample( self, game_state_t: Any, controller_t: Any, *, temperature: float = 1.0, generator: torch.Generator | None = None, reset_mask: torch.Tensor | None = None, padding_mask: torch.Tensor | None = None, cache: KVCache | None = None, use_cache: bool = False, player_name: Any | None = None, ) -> FayntModelOutput: """Sample next controllers for a sequence of observed states. When a whole sequence is supplied, each position conditions on its supplied current controller. Live play should feed the actual executed controller back into the following frame through ``step``. """ labels = _controller_labels(self.codec, controller_t) with _autocast_context(labels.buttons.device, self.policy.config.torch_compute_dtype): encoded = self.policy.encoder(game_state_t, labels, player_name=player_name) hidden, next_cache = self.policy.backbone( encoded, reset_mask=reset_mask, padding_mask=padding_mask, cache=cache, use_cache=use_cache, ) value = self.policy.controller_head.sample( hidden, labels, temperature=temperature, generator=generator, ) return self._output(value, next_cache) def step( self, game_state_t: Any, controller_t: Any, cache: KVCache | None = None, reset_mask: torch.Tensor | None = None, *, sample: bool = True, temperature: float = 1.0, generator: torch.Generator | None = None, player_name: Any | None = None, ) -> FayntModelOutput: """Consume one frame with tensor prefixes ``[batch]``. Set ``reset_mask[b]`` when a new game begins for slot b. Other batch slots keep their histories. Call with eval() and inference_mode() for deployment. The returned cache is the same object updated in place. """ labels = _controller_labels(self.codec, controller_t) if labels.buttons.ndim != 1: raise ValueError("step expects controller tensors shaped [batch]") if cache is None: cache = self.init_cache(labels.buttons.shape[0], device=labels.buttons.device) with _autocast_context(labels.buttons.device, self.policy.config.torch_compute_dtype): encoded = self.policy.encoder(game_state_t, labels, player_name=player_name) hidden, cache = self.policy.backbone.step(encoded, cache, reset_mask=reset_mask) if sample: value = self.policy.controller_head.sample( hidden, labels, temperature=temperature, generator=generator, ) else: value = self.policy.controller_head(hidden, labels) return self._output(value, cache) def example_inputs( self, batch_size: int = 1, sequence_length: int | None = 1, *, device: torch.device | str | None = None, ) -> dict[str, Any]: """Create a synthetic input fixture for loading and inference checks. The fixture depicts two neutral Fox states on Final Destination. Use sequence_length=None for a single-frame step input shaped [batch]. Real gameplay requires the complete parsed live state using the field semantics in tensor_batch.py, with p0 as self and p1 as opponent. """ if batch_size <= 0 or (sequence_length is not None and sequence_length <= 0): raise ValueError("batch_size and sequence_length must be positive") shape = (batch_size,) if sequence_length is None else (batch_size, sequence_length) target_device = self.device if device is None else torch.device(device) def full(value, dtype=torch.float32, tensor_shape=shape): return torch.full(tensor_shape, value, dtype=dtype, device=target_device) def controller(): return ControllerBatch( main_stick=StickBatch(full(0.5), full(0.5)), c_stick=StickBatch(full(0.5), full(0.5)), shoulder=full(0.0), buttons=ButtonsBatch(**{ name: full(False, torch.bool) for name in ButtonsBatch.__dataclass_fields__ }), ) def player(x, facing): fields = { "percent": full(0, torch.long), "facing": full(facing, torch.bool), "x": full(x), "y": full(0.0), "action": full(14, torch.long), "invulnerable": full(False, torch.bool), "character": full(1, torch.long), "jumps_left": full(2, torch.long), "shield_strength": full(60.0), "on_ground": full(True, torch.bool), } follower = {key: torch.zeros_like(value) for key, value in fields.items()} return PlayerBatch( **fields, controller=controller(), nana=NanaBatch(exists=full(False, torch.bool), **follower), ) item_shape = (*shape, 15) game = GameStateBatch( p0=player(-20.0, True), p1=player(20.0, False), stage=full(25, torch.long), randall=RandallBatch(full(0.0), full(0.0)), fod_platforms=FoDPlatformsBatch(full(0.0), full(0.0)), items=ItemsBatch( exists=full(False, torch.bool, item_shape), type=full(0, torch.long, item_shape), state=full(0, torch.long, item_shape), x=full(0.0, tensor_shape=item_shape), y=full(0.0, tensor_shape=item_shape), ), ) return {"game_state_t": game, "controller_t": controller()} FayntModel.register_for_auto_class("AutoModel")