| """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: |
| |
| |
| 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") |
|
|
