File size: 3,778 Bytes
026f63e | 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 | """Task plugin interface: the only place where MiniJEV touches a domain.
A task turns its native states into (a) compact arrays for storage and
(b) model inputs for one of MiniJEV's state encoders. Everything else (model,
training, calibration, export, runtime) is task-agnostic.
Plugins live in `environments/<name>/task.py` and expose `make_task(cfg)`.
"""
from __future__ import annotations
import importlib
from typing import Any, Sequence
import numpy as np
class Task:
name: str = "task"
#: which MiniJEV state encoder consumes this task's inputs ("grid", "text", ...)
state_encoder: str = "grid"
default_question: str = ""
default_choices: tuple[str, ...] = ()
def __init__(self, cfg: dict | None = None):
self.cfg = cfg or {}
# labels -----------------------------------------------------------------
#: stored label -> label shown to the model (e.g. Snake "UP" -> "move up"); empty = identity
@property
def choice_alias(self) -> dict[str, str]:
return {}
def model_label(self, label: str) -> str:
return self.choice_alias.get(label, label)
def action_of(self, label: str) -> str:
"""Model-facing label -> native action (inverse of choice_alias)."""
inv = {v: k for k, v in self.choice_alias.items()}
return inv.get(label, label)
# storage format --------------------------------------------------------
def states_to_arrays(self, states: Sequence[Any]) -> dict[str, np.ndarray]:
raise NotImplementedError
def encode_arrays(self, arrays: dict[str, np.ndarray], idx: np.ndarray) -> dict[str, np.ndarray]:
"""Model inputs for rows `idx` of stored arrays (vectorised numpy)."""
raise NotImplementedError
# convenience ------------------------------------------------------------
def encode_states(self, states: Sequence[Any]) -> dict[str, np.ndarray]:
arrays = self.states_to_arrays(states)
return self.encode_arrays(arrays, np.arange(len(states)))
def encoder_class(self):
"""Optional task-specific StateEncoder class (must keep the generic encoder's parameters)."""
return None
def encoder_config(self) -> dict:
"""Hyper-parameters the state encoder needs from the task (channels, ...)."""
return {}
def example_states(self, n: int = 4, seed: int = 0) -> list[Any]:
"""A few valid states (used for ONNX export tracing and parity checks)."""
raise NotImplementedError
BUILTIN = {"text": "minijev.text_task", "log": "minijev.text_task"}
def get_task(name: str, cfg: dict | None = None) -> Task:
mod = importlib.import_module(BUILTIN.get(name) or f"environments.{name.replace('-', '_')}.task")
return mod.make_task(cfg or {})
# ----------------------------------------------------------------- ragged utils
def ragged_pack(seqs: Sequence[Sequence[int]] | Sequence[np.ndarray], dtype=np.int32) -> tuple[np.ndarray, np.ndarray]:
lengths = np.fromiter((len(s) for s in seqs), dtype=np.int64, count=len(seqs))
offsets = np.zeros(len(seqs) + 1, dtype=np.int64)
np.cumsum(lengths, out=offsets[1:])
flat = np.concatenate([np.asarray(s, dtype=dtype) for s in seqs]) if len(seqs) else np.zeros(0, dtype)
return flat.astype(dtype, copy=False), offsets
def ragged_take(flat: np.ndarray, offsets: np.ndarray, idx: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
"""Gather rows `idx` of a ragged array; returns (flat, lengths)."""
starts = offsets[idx]
lengths = offsets[idx + 1] - starts
total = int(lengths.sum())
if total == 0:
return flat[:0], lengths
rel = np.arange(total) - np.repeat(np.cumsum(lengths) - lengths, lengths)
return flat[np.repeat(starts, lengths) + rel], lengths
|