diff --git a/.github/workflows/axis_mira_v7_research.yml b/.github/workflows/axis_mira_v7_research.yml new file mode 100644 index 000000000..3628f7767 --- /dev/null +++ b/.github/workflows/axis_mira_v7_research.yml @@ -0,0 +1,55 @@ +name: axis-mira-v7-research + +on: + pull_request: + paths: + - 'models/axis/rig_program_model_v7.py' + - 'models/mira/skin_fsq_codec_v7.py' + - 'models/mira/skin_attention_codec_v7.py' + - 'compiler/realsas_compiler_core/mechanics_interaction_v1.py' + - 'tests/models/test_axis_mira_v7_research.py' + - 'canonical/AXIS_MIRA_V7_CORPUS370_PREREG_V1.json' + - 'docs/research/AXIS_MIRA_V7_SKINTOKENS_PLUSPLUS_20261106.md' + - '.github/workflows/axis_mira_v7_research.yml' + workflow_dispatch: + +permissions: + contents: read + +concurrency: + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }} + cancel-in-progress: ${{ github.event_name == 'pull_request' }} + +jobs: + v7-contract-courts: + runs-on: ubuntu-latest + timeout-minutes: 25 + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: '3.11' + - name: Install CPU research requirements + run: | + set -euo pipefail + python -m pip install --no-input --progress-bar off -r requirements/mainline-ci.txt + python -m pip install --no-input --progress-bar off -r requirements/torch-cpu.txt + - name: Syntax + run: | + python -m py_compile models/axis/rig_program_model_v7.py models/mira/skin_fsq_codec_v7.py models/mira/skin_attention_codec_v7.py compiler/realsas_compiler_core/mechanics_interaction_v1.py + - name: Source and Mechanic court + run: | + pytest -q tests/models/test_axis_mira_v7_research.py + - name: Verify no V7 training authority was minted + run: | + python - <<'PY' + import json + from pathlib import Path + d = json.loads(Path('canonical/AXIS_MIRA_V7_CORPUS370_PREREG_V1.json').read_text()) + assert d['training']['allowed_now'] is False + assert d['training']['checkpoint_product_authority'] is False + assert d['owner_contract']['graphics'] == 'SEPARATE_IRIS_LINKED_GRAPHICS_COMPILER' + assert d['skin']['vocab_size'] == 64000 + assert d['skin']['source_vae_attention_parity_claimed'] is False + print('AXIS_MIRA_V7_RESEARCH_AUTHORITY_SEALED') + PY diff --git a/canonical/AXIS_MIRA_V7_CORPUS370_PREREG_V1.json b/canonical/AXIS_MIRA_V7_CORPUS370_PREREG_V1.json new file mode 100644 index 000000000..9fdc034ba --- /dev/null +++ b/canonical/AXIS_MIRA_V7_CORPUS370_PREREG_V1.json @@ -0,0 +1,60 @@ +{ + "schema": "RealSaS.AxisMiraV7Corpus370Prereg.v1", + "date_utc": "2026-11-06", + "status": "RESEARCH_SOURCE_ONLY__NO_TRAINING_AUTHORITY", + "branch": "research/axis-mira-v7-skintokens-control-370", + "source_cohort": { + "candidate_count": 370, + "admitted_for_training": null, + "observation_pass_not_equivalent_to_truth_admission": true, + "upstream_work": "PR #327" + }, + "owner_contract": { + "axis": "LEARNED_MECHANICS_RIG_CONTROL_AND_SKIN_CODE", + "mira": "LEARNED_MECHANICAL_SKIN_FIELD_CODEC", + "mechanics": "EXECUTABLE_CONSTRAINTS_IK_DEFORMATION_INTERACTION", + "graphics": "SEPARATE_IRIS_LINKED_GRAPHICS_COMPILER", + "motion": "CONTROL_SPACE_MOTION" + }, + "rig": { + "architecture_id": "RealSaS.AXIS.UnifiedRigControlSkin.CausalCrossAttention.v7_0", + "scene_dim": 47, + "hidden_dim": 768, + "layers": 16, + "heads": 12, + "max_sequence_length": 4096, + "skin_vocab_size": 64000 + }, + "skin": { + "architecture_id": "RealSaS.MIRA.CarrierNativeFSQ.STGradient.ActiveField.v7_0", + "tokens_per_joint": 4, + "fsq_levels": [ + 8, + 8, + 8, + 5, + 5, + 5 + ], + "vocab_size": 64000, + "source_vae_attention_parity_claimed": false + }, + "required_receipts": [ + "CORPUS370_SOURCE_FAMILY_QUALIFICATION", + "CORPUS370_MECHANICS_INTERVENTION_TRUTH", + "CORPUS370_CANONICAL_MESH_SHARED_SCENE_BINDINGS", + "CORPUS370_DENSE_SKIN_TARGET_AUTHORITY", + "CORPUS370_LEAKAGE_CONNECTED_TRAIN_UNSEEN_SPLIT", + "MIRA_V7_CODEC_FROZEN_TRAIN_ONLY", + "AXIS_V7_CAUSAL_EVENT_TEACHER_SEAL", + "MECHANICS_INTERACTABILITY_BEHAVIOR_COURT" + ], + "training": { + "allowed_now": false, + "checkpoint_product_authority": false, + "unseen_tuning_allowed": false, + "source_name_student_truth_allowed": false, + "compiler_control_synthesis_allowed": false, + "graphics_train_target_in_axis_mira": false + } +} diff --git a/compiler/realsas_compiler_core/mechanics_interaction_v1.py b/compiler/realsas_compiler_core/mechanics_interaction_v1.py new file mode 100644 index 000000000..c5a347ca0 --- /dev/null +++ b/compiler/realsas_compiler_core/mechanics_interaction_v1.py @@ -0,0 +1,127 @@ +from __future__ import annotations + +"""Mechanics-only implicit control selection for a picked carrier triangle. + +This is a read-only runtime selection of ALREADY DECLARED controls; it never +synthesizes controls, animates graphics, or mints CanonicalRig authority. +The caller supplies the hit triangle and barycentrics from a presentation hit +test. Returned control IDs are opaque; only Mechanics interprets their action. +""" + +from dataclasses import dataclass +import math + +from .canonical_rig_ir_v2 import CanonicalRigIRV2 +from .canonical_skin_ir_v2 import CanonicalSkinIRV2 +from .types import QualificationError + + +@dataclass(frozen=True) +class MechanicsPickCandidateV1: + canonical_control_id: str + control_kind: str + influence_score: float + supporting_joint_ids: tuple[str, ...] + + +@dataclass(frozen=True) +class MechanicsSurfacePickV1: + triangle_vertex_ids: tuple[str, str, str] + barycentric: tuple[float, float, float] + candidates: tuple[MechanicsPickCandidateV1, ...] + rig_lineage_hash: str + skin_lineage_hash: str + status: str # CANDIDATES | ABSTAIN_NO_DECLARED_CONTROL + + +def select_mechanics_controls_from_surface_v1( + rig: CanonicalRigIRV2, + skin: CanonicalSkinIRV2, + *, + triangle_vertex_ids: tuple[str, str, str], + barycentric: tuple[float, float, float], + maximum_candidates: int = 8, + minimum_support: float = 1e-6, +) -> MechanicsSurfacePickV1: + if rig.mesh_lineage_hash != skin.mesh_lineage_hash: + raise QualificationError("MECHANICS_PICK_MESH_LINEAGE_DRIFT") + if rig.deform_structure_hash != skin.deform_structure_hash: + raise QualificationError("MECHANICS_PICK_DEFORM_LINEAGE_DRIFT") + if maximum_candidates < 1 or not (0.0 <= minimum_support <= 1.0): + raise QualificationError("MECHANICS_PICK_POLICY_INVALID") + if len(triangle_vertex_ids) != 3 or len(set(triangle_vertex_ids)) != 3: + raise QualificationError("MECHANICS_PICK_TRIANGLE_INVALID") + if len(barycentric) != 3 or not all(math.isfinite(float(x)) and float(x) >= 0.0 for x in barycentric): + raise QualificationError("MECHANICS_PICK_BARYCENTRIC_INVALID") + mass = sum(barycentric) + if abs(mass - 1.0) > 1e-5: + raise QualificationError("MECHANICS_PICK_BARYCENTRIC_NOT_SIMPLEX") + + by_vertex = {row.canonical_mesh_vertex_id: row for row in skin.rows} + if len(by_vertex) != len(skin.rows): + raise QualificationError("MECHANICS_PICK_DUPLICATE_SKIN_VERTEX") + if any(vid not in by_vertex for vid in triangle_vertex_ids): + raise QualificationError("MECHANICS_PICK_SKIN_VERTEX_UNKNOWN") + + influence: dict[str, float] = {} + for vid, bary in zip(triangle_vertex_ids, barycentric): + for joint_id, weight in by_vertex[vid].influences: + value = float(weight) + if not math.isfinite(value) or value < 0.0: + raise QualificationError("MECHANICS_PICK_SKIN_WEIGHT_INVALID") + influence[joint_id] = influence.get(joint_id, 0.0) + float(bary) * value + + joint_ids = {j.canonical_joint_id for j in rig.joints} + if not set(influence).issubset(joint_ids): + raise QualificationError("MECHANICS_PICK_JOINT_UNKNOWN") + chains = {row.canonical_chain_id: row for row in rig.chains} + # IK effectors are recognized only through their DECLARED constraint inputs. + ik_targets: dict[str, set[str]] = {} + for constraint in rig.constraints: + if constraint.operator_id != "IK" or constraint.canonical_chain_id is None: + continue + chain = chains.get(constraint.canonical_chain_id) + if chain is None or not chain.canonical_joint_ids: + raise QualificationError("MECHANICS_PICK_CHAIN_UNKNOWN") + tip = chain.canonical_joint_ids[-1] + for control_id in constraint.input_canonical_control_ids[:1]: + ik_targets.setdefault(control_id, set()).add(tip) + + candidates: list[MechanicsPickCandidateV1] = [] + seen: set[str] = set() + for control in rig.controls: + cid = control.canonical_control_id + if cid in seen: + raise QualificationError("MECHANICS_PICK_DUPLICATE_CONTROL") + seen.add(cid) + if control.control_kind not in {"IK_EFFECTOR", "FK", "ROOT", "SPLINE", "AIM_TARGET", "SOCKET_CONTROL"}: + continue + targets = set(control.target_canonical_joint_ids) + if control.control_kind == "IK_EFFECTOR": + targets.update(ik_targets.get(cid, set())) + if not targets or not targets.issubset(joint_ids): + continue + support = sum(influence.get(jid, 0.0) for jid in targets) + if support < minimum_support: + continue + candidates.append(MechanicsPickCandidateV1( + canonical_control_id=cid, + control_kind=control.control_kind, + influence_score=float(support), + supporting_joint_ids=tuple(sorted(targets)), + )) + candidates.sort(key=lambda x: (-x.influence_score, x.canonical_control_id)) + return MechanicsSurfacePickV1( + triangle_vertex_ids=triangle_vertex_ids, + barycentric=tuple(map(float, barycentric)), + candidates=tuple(candidates[:maximum_candidates]), + rig_lineage_hash=rig.rig_lineage_hash, + skin_lineage_hash=skin.skin_lineage_hash, + status="CANDIDATES" if candidates else "ABSTAIN_NO_DECLARED_CONTROL", + ) + + +__all__ = [ + "MechanicsPickCandidateV1", "MechanicsSurfacePickV1", + "select_mechanics_controls_from_surface_v1", +] diff --git a/docs/research/AXIS_MIRA_V7_SKINTOKENS_PLUSPLUS_20261106.md b/docs/research/AXIS_MIRA_V7_SKINTOKENS_PLUSPLUS_20261106.md new file mode 100644 index 000000000..791e1c59e --- /dev/null +++ b/docs/research/AXIS_MIRA_V7_SKINTOKENS_PLUSPLUS_20261106.md @@ -0,0 +1,89 @@ +# AXIS–MIRA V7 / SkinTokens++ — Mechanics-first research contract + +**Date:** 2026-11-06 +**Status:** `RESEARCH_CODE_ONLY__NO_TRAINING_AUTHORITY__NO_PRODUCT_MODELRELEASE` +**Parent:** V6 product branch `axis-v6-rig-program-product`; isolated PR stack. +**Goal:** compiled puppet is natively controllable ("You can play with your puppet"). + +## Ownership is non-negotiable + +| Domain | Authority | What it owns | +|---|---|---| +| Shape / geometry | IRIS + TESSA / qualified Mesh | shared 3D evidence and mechanical carrier geometry | +| **Rig/Skin learning** | **AXIS–MIRA V7** | skeleton and rest geometry; frame/DOF; controls and relations; compact skin influence fields | +| **Mechanical execution** | **Mechanics Compiler/runtime** | constraint solving, IK/FK, skin deformation, contact/limits, attachment transforms, typed residuals | +| **Graphics** | **IRIS-linked Graphics domain + Graphics Compiler** | CAA/appearance, layer decomposition, occlusion, draw order, visibility, 2D presentation | +| Motion | Control-space motion | editable control trajectories executed by Mechanics | + +The Mechanics result may be consumed by Graphics, but neither AXIS–MIRA nor Mechanics owns graphics decisions. No AXIS model target for alpha/layer/RGB/draw order is permitted. + +## SkinTokens++ — upstream baseline and what *plus* must mean + +Researched the **official source**, not just the paper. Reference repository: https://github.com/VAST-AI-Research/SkinTokens. +Source files reviewed on 2026-11-06: + +- `src/model/tokenrig.py` — GitHub blob `ab35e3fb8dd850737cba7d1782f03c75846971be`: Qwen3 0.6B-style autoregressive model; learned Michelangelo mesh conditioning; vocabulary switch grammar; frozen skin VAE in TokenRig. +- `src/model/skin_vae/autoencoders/skin_fsq_cvae_model.py` — blob `6b66d9e9486d261f0a5d67f795bbe60499fb26e3`: separate skin and conditional attention encoders, query sampling, latent quantization, learned decoder with sampled conditioning. +- `src/model/skin_vae/autoencoders/FSQ.py` — blob `9f7bcf0e670071a04a19495ea8efa7d3ead0a1fe`: differentiable round-STE finite scalar quantization, integer-code conversions. +- `src/model/skin_vae_model.py` — blob `a2b4e88b6997883ce5861333bb498dc9593ff80d`: token compression / decompression, per-bone examples, sampling and training wrapper. +- `src/tokenizer/spec.py` — blob `d4cc653519937bb98d6059d765686e2ea53b89dc`: structural skeleton sequence and parent-link tokenization. + +Their published checkpoint naming includes `skin_vae_2_10_32768`. **Our 64,000 vocabulary is a new local configuration, not evidence of checkpoint or tokenizer compatibility**. Never suggest that official weights load directly. Respect repository/model licenses and provenance independently. + +### Mechanistic comparison + +| Capability | Upstream TokenRig / SkinTokens | RealSaS V7 initial implementation | Next step | +|---|---|---|---| +| Unified skeleton + skin sequence | Yes | Yes, via typed V6 wire reused by V7 | new event grammar for interaction-affordance targets only after teacher seal | +| FSQ + straight-through reconstruction | Yes | V7 MIRA uses true gradient-carrying code path | verify with reconstruction-only backward court | +| Spatial attention-based VAE | Yes | Two V7 options: baseline uses weighted pooling; new **learned-query local attention candidate** encodes active sampled support. Not exact upstream CVAE parity | compare baseline and new attention candidate under equal train/validation budgets; add geometry-conditioned decoder attention only if measured improvements justify it | +| Rig conditioning | Learned mesh encoder | Current IRIS/TESSA-adaptable 47-D surface evidence seam | first-class versioned shared-scene adapter, no hidden oracle truth | +| Skeleton generation | Yes | V6 RigProgram language extended by large V7 decoder | require true free-running unseen tests | +| Control rig / interaction affordances | Not the core TokenRig target | V6 control/constraint program + runtime implicit pick from declared controls | learned/evidence-grounded affordances; local DOF/limit solver | +| Real-time deformation | Downstream use after rigging | V6 Mechanics solver; new pick selection bridge | batched/C++ runtime; contact and joint-limit courts | +| 2D drawing/occlusion | Not a TokenRig output | **Not owned here** | downstream Graphics consumers only | + +## Implementation in this V7 research branch + +- `models/mira/skin_fsq_codec_v7.py`: 64k code vocabulary from (8,8,8,5,5,5), 4 tokens/joint, large carrier ceiling, sparse-aware reconstruction, teacher-relative edge gradient loss, reconstruction-only STE gradient. **Initial codec baseline, not full official architecture.** +- `models/mira/skin_attention_codec_v7.py`: source-inspired but independently implemented learned-query per-joint attention with bounded active-support top-k sampling, inherited straight-through FSQ decode. Not weight compatible with upstream.\n- `models/axis/rig_program_model_v7.py`: 16 x 768, 12-head conditional causal model; typed RigProgram + SkinCodes, tied embedding/output projection, sparse skin logits at SKIN_CODE positions, V6 wire contract and variable morpologies. +- `compiler/realsas_compiler_core/mechanics_interaction_v1.py`: generic barycentric hit-to-declared-control scoring over qualified skin/rig. This is runtime interpretation of already generated structure; it **does not** synthesize control rig, choose new rig parents, or make Graphics decisions. +- `tests/models/test_axis_mira_v7_research.py`: reconstruction-only gradient, hard-code parity, sparse CE/backprop, read-only pick, fail-closed lineage. + +## Corpus370 policy + +**370 is a candidate acquisition universe, NOT 370 qualified training examples.** Existing PR #327 retains 352 observation PASS / 18 original mask FAIL while further family courts, source-locus ownership evidence, mechanical target qualification and lineage-safe splits remain open. Do not train V7 simply because source observations pass. + +V7 requires separated versioned receipts: + +1. (Mesh, scene, Rig, Skin) source-to-canonical mapping with evidence and no privileged source-name leakage. +2. Mechanically intervention-tested joint/frame/DOF labels; absence/ambiguous labels typed as masked/abstain, not guessed. +3. Exact MIRA V7 codec TRAIN split targets; no cross-family leakage or unseen access. +4. Frozen MIRA codec **before** AXIS SkinCode teacher targets. +5. AXIS V7 frozen TRAIN/internal validation then sealed one-time UNSEEN scope. +6. Free-running **behavioral** tests: reachable handle IK, hand/weapon attachment, contact consistency, joint-limit compliance, deformation distortion, sparse influence, presentation-agnostic under motion. +7. Knight M+C AA regression independent of unseen selection. No product certification based on a single sample. + +No V7 training or ProductRevision promotion is authorized by this research branch. + +## Milestones (large increments, no model-size theater) + +**A — engineering research candidate (this PR):** higher-capacity coupled Rig/Skin model, improved V7 baseline codec, typed Mechanics hit interaction, CI, exact upstream code references, versioned corpus gate. + +**B — architecture science:** benchmark new local-attention MIRA vs initial pooling baseline, then evaluate deeper conditional decoder/encoder parity against upstream; quantify active-support sampling on nonhumanoid corpus, gradient preservation, smoothness and reconstruction at 32k vs 64k, fixed budget. + +**C — mechanics semantics:** local DOF/limits, targets/attachments orientation, multi-constraint solves, influence-to-control ambiguity and abstention, spatial interaction tracking; no graphics authority creep. + +**D — data/model training:** qualified Corpus370 train/unseen split, frozen codec, coupled AXIS, checkpoint selection only on TRAIN-internal validation, independent unseen behavior court. + +**E — real-time product:** C++ cached solver/deform path, mouse interaction latency/p95 metrics, Graphics consuming Mechanics state, independent source fidelity and dynamic visual proofs. + +## Research falsification rules + +- Strong SkinCode MAE but bad contact, fold or vertex distortion is **FAIL** for playable puppet. +- A legal skeleton with no usable declared handle is **FAIL** for interactive suitability. +- One Knight oracle control-space parity is regression only, never proof of learned rigging. +- A Graphics fidelity problem belongs to Graphics and its evidence binding, not to an AXIS graphics head. +- Any missing ModelRelease, absent exact codec hash or unsealed subject family is fail-closed. +- Bigger architecture is not itself a gain; model capacity must be backed by training data, sampling, pretraining or demonstrable learning curves. + diff --git a/models/axis/rig_program_model_v7.py b/models/axis/rig_program_model_v7.py new file mode 100644 index 000000000..8c777acdc --- /dev/null +++ b/models/axis/rig_program_model_v7.py @@ -0,0 +1,335 @@ +from __future__ import annotations + +"""AXIS V7: large conditional causal rig+control+SkinCode model. + +Mechanics is the learning target. Graphics is NOT part of this architecture. +Uses the typed V6 program language as a lossless migration seam; a separate +V7 interactability language may be introduced only after truth/courts approve it. + +This research model outputs rig graph, rest frames, controls, constraints, +attachments and MIRA V7's 64k skin codes within ONE causal hypothesis. Skin +logits are computed sparsely only at actual SKIN_CODE positions to avoid the +B x sequence_length x 64000 activation tensor during teacher forcing. + +Research checkpoint ONLY; no trained weights, ModelRelease or ProductRevision. +""" + +from dataclasses import asdict, dataclass +import math + +import torch +from torch import nn, Tensor +import torch.nn.functional as F + +from models.axis.rig_program_model_v6 import ( + AXISSequenceV6, AXISRigProgramGrammarV6Product, + RigGrammarStateV6, RigProgramEventV6, RigSectionV6, RigTokenKindV6, + CONTROL_KINDS_V6, CONSTRAINT_OPERATORS_V6, DOF_LABELS_V6, +) +from models.mira.skin_fsq_codec_v7 import MIRACodecConfigV7 + + +AXIS_V7_ARCHITECTURE_ID = "RealSaS.AXIS.UnifiedRigControlSkin.CausalCrossAttention.v7_0" + + +@dataclass(frozen=True) +class AXISRigProgramConfigV7: + scene_dim: int = 47 + hidden_dim: int = 768 + layers: int = 16 + heads: int = 12 + feedforward_mult: int = 4 + max_sequence_length: int = 4096 + max_pointer_entities: int = 256 + coord_bins: int = 256 + coord_min: float = -1.0 + coord_max: float = 1.0 + skin_vocab_size: int = 64000 + tokens_per_joint: int = 4 + dropout: float = 0.0 + + def __post_init__(self): + if min(self.scene_dim, self.hidden_dim, self.layers, self.heads, + self.feedforward_mult, self.max_pointer_entities, + self.max_sequence_length, self.coord_bins, self.skin_vocab_size, + self.tokens_per_joint) < 1: + raise ValueError("AXIS_V7_CONFIG_NONPOSITIVE") + if self.hidden_dim % self.heads: + raise ValueError("AXIS_V7_HEAD_DIVISIBILITY") + if not self.coord_max > self.coord_min: + raise ValueError("AXIS_V7_COORD_RANGE_INVALID") + if self.max_sequence_length > 65536: + raise ValueError("AXIS_V7_SEQUENCE_POLICY_LIMIT") + + def to_dict(self) -> dict: + return asdict(self) + + +class AXISRigProgramModelV7(nn.Module): + architecture_id = AXIS_V7_ARCHITECTURE_ID + + def __init__(self, config: AXISRigProgramConfigV7 = AXISRigProgramConfigV7()): + super().__init__() + self.config = config + h = config.hidden_dim + self.kind_embedding = nn.Embedding(len(RigTokenKindV6), h) + self.coord_embedding = nn.Embedding(config.coord_bins + 1, h) + self.pointer_embedding = nn.Embedding(config.max_pointer_entities + 1, h) + self.aux_embedding = nn.Embedding(max(len(CONTROL_KINDS_V6), len(CONSTRAINT_OPERATORS_V6)) + 1, h) + self.skin_embedding = nn.Embedding(config.skin_vocab_size + 1, h) + self.anchor_parent_embedding = nn.Embedding(3, h) + self.position_embedding = nn.Embedding(config.max_sequence_length, h) + self.dof_proj = nn.Linear(len(DOF_LABELS_V6), h, bias=False) + self.frame_proj = nn.Linear(6, h, bias=False) + self.scalar_proj = nn.Linear(4, h, bias=False) + self.input_norm = nn.LayerNorm(h) + self.scene_projection = nn.Sequential(nn.Linear(config.scene_dim, h), + nn.LayerNorm(h), nn.GELU()) + self.scene_type_embedding = nn.Parameter(torch.zeros(1, 1, h)) + layer = nn.TransformerDecoderLayer( + d_model=h, nhead=config.heads, + dim_feedforward=h * config.feedforward_mult, + dropout=config.dropout, activation="gelu", + batch_first=True, norm_first=True, + ) + self.transformer = nn.TransformerDecoder(layer, num_layers=config.layers, + norm=nn.LayerNorm(h)) + self.kind_head = nn.Linear(h, len(RigTokenKindV6)) + self.coord_head = nn.Linear(h, 3 * config.coord_bins) + self.pointer_a_head = nn.Linear(h, config.max_pointer_entities + 1) + self.pointer_b_head = nn.Linear(h, config.max_pointer_entities + 1) + self.pointer_c_head = nn.Linear(h, config.max_pointer_entities + 1) + self.control_kind_head = nn.Linear(h, len(CONTROL_KINDS_V6)) + self.constraint_operator_head = nn.Linear(h, len(CONSTRAINT_OPERATORS_V6)) + self.dof_head = nn.Linear(h, len(DOF_LABELS_V6)) + self.frame_6d_head = nn.Linear(h, 6) + self.scalar_head = nn.Linear(h, 4) + self.anchor_parent_head = nn.Linear(h, 2) + self.skin_bias = nn.Parameter(torch.zeros(config.skin_vocab_size)) + + @staticmethod + def _shift(value: Tensor, maximum: int) -> Tensor: + return (value.to(torch.long) + 1).clamp(0, maximum) + + def _embed(self, seq: AXISSequenceV6) -> Tensor: + seq.validate(self.config) + b, length = seq.kind_ids.shape + positions = torch.arange(length, device=seq.kind_ids.device)[None].expand(b, -1) + x = self.kind_embedding(seq.kind_ids) + self.position_embedding(positions) + x = x + self.coord_embedding(self._shift(seq.coord_indices, self.config.coord_bins)).sum(dim=2) + for ptr in (seq.pointer_a, seq.pointer_b, seq.pointer_c): + x = x + self.pointer_embedding(self._shift(ptr, self.config.max_pointer_entities)) + x = x + self.aux_embedding(self._shift(seq.aux_type_ids, self.aux_embedding.num_embeddings - 1)) + x = x + self.skin_embedding(self._shift(seq.skin_code_ids, self.config.skin_vocab_size)) + x = x + self.anchor_parent_embedding(self._shift(seq.anchor_parent_kind, 2)) + x = x + self.dof_proj(seq.dof_masks.to(x.dtype)) + x = x + self.frame_proj(seq.frame_6d.to(x.dtype)) + x = x + self.scalar_proj(seq.scalar_params.to(x.dtype)) + return self.input_norm(x) + + def forward(self, scene_memory: Tensor, sequence: AXISSequenceV6) -> dict[str, Tensor]: + if scene_memory.ndim != 3 or scene_memory.shape[0] != sequence.kind_ids.shape[0]: + raise ValueError("AXIS_V7_SCENE_BATCH_OR_RANK_INVALID") + if scene_memory.shape[-1] != self.config.scene_dim or scene_memory.shape[1] < 1: + raise ValueError("AXIS_V7_SCENE_DIM_INVALID") + if not torch.isfinite(scene_memory).all(): + raise ValueError("AXIS_V7_SCENE_NONFINITE") + x = self._embed(sequence) + s = x.shape[1] + # True values mask forbidden future tokens; no attention to targets. + causal_mask = torch.triu( + torch.ones((s, s), device=x.device, dtype=torch.bool), diagonal=1, + ) + scene = self.scene_projection(scene_memory.to(x.dtype)) + self.scene_type_embedding + hidden = self.transformer( + x, scene, tgt_mask=causal_mask, + tgt_key_padding_mask=~sequence.attention_mask.to(torch.bool), + ) + return { + "hidden": hidden, + "kind_logits": self.kind_head(hidden), + "coord_logits": self.coord_head(hidden).reshape(*hidden.shape[:2], 3, self.config.coord_bins), + "pointer_a_logits": self.pointer_a_head(hidden), + "pointer_b_logits": self.pointer_b_head(hidden), + "pointer_c_logits": self.pointer_c_head(hidden), + "control_kind_logits": self.control_kind_head(hidden), + "constraint_operator_logits": self.constraint_operator_head(hidden), + "dof_logits": self.dof_head(hidden), + "frame_6d": self.frame_6d_head(hidden), + "scalar_params": self.scalar_head(hidden), + "anchor_parent_logits": self.anchor_parent_head(hidden), + } + + def skin_code_logits(self, hidden: Tensor) -> Tensor: + # Tied code embedding avoids a second 64000 x hidden_dim weight matrix. + return F.linear(hidden, self.skin_embedding.weight[1:], self.skin_bias) + + @staticmethod + def _pointer(logits: Tensor, count: int, allow_null: bool) -> int: + if count < 0 or count + 1 > logits.shape[-1]: + raise ValueError("AXIS_V7_POINTER_RESOURCE_GUARD") + values = logits[:count + 1].clone() + if not allow_null: + if count < 1: + raise ValueError("AXIS_V7_REQUIRED_POINTER_EMPTY") + values[0] = float("-inf") + return int(torch.argmax(values).item()) - 1 + + @torch.no_grad() + def generate_events(self, scene_memory: Tensor, *, max_steps: int | None = None + ) -> tuple[RigProgramEventV6, ...]: + """Exact V6 product event wire; autoregressive correctness baseline. + + Full-prefix decoding is intentionally not represented as a low-latency + product path. It requires KV-cache benchmarking before shipping. + """ + self.eval() + if scene_memory.ndim != 3 or scene_memory.shape[0] != 1: + raise ValueError("AXIS_V7_GENERATION_BATCH_ONE_REQUIRED") + limit = min(max_steps or self.config.max_sequence_length, self.config.max_sequence_length) + state = AXISRigProgramGrammarV6Product.advance( + RigGrammarStateV6(), RigTokenKindV6.BOS, + tokens_per_joint=self.config.tokens_per_joint, + ) + events = [RigProgramEventV6(RigTokenKindV6.BOS)] + for _ in range(limit - 1): + prefix = AXISSequenceV6.from_events(events, device=scene_memory.device) + out = self(scene_memory, prefix) + kind_logits = AXISRigProgramGrammarV6Product.mask( + out["kind_logits"][0, -1], state, + tokens_per_joint=self.config.tokens_per_joint, + ) + kind = RigTokenKindV6(int(torch.argmax(kind_logits).item())) + coord = (-1, -1, -1) + pa = pb = aux = code = anchor_kind = -1 + frame = (0.0,) * 6 + dof = (0.0,) * len(DOF_LABELS_V6) + params = (0.0,) * 4 + last = {key: value[0, -1] for key, value in out.items() if key != "hidden"} + ptr = lambda which, count, nullable=False: self._pointer( + last[which], count, nullable, + ) + if kind in (RigTokenKindV6.JOINT, RigTokenKindV6.CONTROL, RigTokenKindV6.ANCHOR): + coord = tuple(int(x) for x in torch.argmax(last["coord_logits"], dim=-1).tolist()) + if kind == RigTokenKindV6.JOINT: + pa = -1 if not state.joint_count else ptr("pointer_a_logits", state.joint_count, True) + elif kind == RigTokenKindV6.FRAME: + pa = ptr("pointer_a_logits", state.joint_count) + frame = tuple(float(x) for x in last["frame_6d"].tolist()) + dof = tuple(float(x) for x in (torch.sigmoid(last["dof_logits"]) >= 0.5).tolist()) + elif kind == RigTokenKindV6.CONTROL: + aux = int(torch.argmax(last["control_kind_logits"]).item()) + pa = ptr("pointer_a_logits", state.control_count, True) + elif kind == RigTokenKindV6.CONTROL_TARGET: + pa = ptr("pointer_a_logits", state.joint_count) + elif kind == RigTokenKindV6.CHAIN: + pa = ptr("pointer_a_logits", state.joint_count) + pb = ptr("pointer_b_logits", state.joint_count) + elif kind == RigTokenKindV6.CONSTRAINT: + aux = int(torch.argmax(last["constraint_operator_logits"]).item()) + pa = ptr("pointer_a_logits", state.chain_count, True) + params = tuple(float(x) for x in last["scalar_params"].tolist()) + elif kind == RigTokenKindV6.CONSTRAINT_INPUT: + pa = ptr("pointer_a_logits", state.control_count) + elif kind == RigTokenKindV6.CONSTRAINT_OUTPUT: + pa = ptr("pointer_a_logits", state.joint_count) + elif kind == RigTokenKindV6.ANCHOR: + anchor_kind = int(torch.argmax(last["anchor_parent_logits"]).item()) + if anchor_kind == 0: + pa = ptr("pointer_a_logits", state.joint_count) + else: + pb = ptr("pointer_b_logits", state.control_count) + elif kind == RigTokenKindV6.SKIN_CODE: + code = int(torch.argmax(self.skin_code_logits(out["hidden"][0, -1])).item()) + event = RigProgramEventV6( + kind=kind, coord_indices=coord, pointer_a=pa, pointer_b=pb, + aux_type_id=aux, dof_mask=dof, frame_6d=frame, + scalar_params=params, skin_code_id=code, + anchor_parent_kind=anchor_kind, + ) + events.append(event) + state = AXISRigProgramGrammarV6Product.advance( + state, kind, tokens_per_joint=self.config.tokens_per_joint, + ) + if state.section == RigSectionV6.DONE: + return tuple(events) + raise RuntimeError("AXIS_V7_GENERATION_RESOURCE_LIMIT_EXHAUSTED") + + +def axis_v7_teacher_forced_objective( + model: AXISRigProgramModelV7, scene_memory: Tensor, + inputs: AXISSequenceV6, targets: AXISSequenceV6, +) -> dict[str, Tensor]: + """All token-field losses; sparse code CE and inclusive nullable pointers.""" + out = model(scene_memory, inputs) + active = targets.attention_mask.to(torch.bool) + kinds = targets.kind_ids + + def ce(key: str, truth: Tensor, mask: Tensor) -> Tensor: + if not bool(mask.any()): + return out["hidden"].sum() * 0.0 + return F.cross_entropy(out[key][mask], truth[mask]) + + def selected(*event_kinds: RigTokenKindV6) -> Tensor: + mask = torch.zeros_like(active) + for kind in event_kinds: + mask = mask | (kinds == int(kind)) + return active & mask + + losses = {} + losses["kind"] = ce("kind_logits", kinds, active) + xyz = selected(RigTokenKindV6.JOINT, RigTokenKindV6.CONTROL, RigTokenKindV6.ANCHOR) + losses["coordinate"] = sum( + F.cross_entropy(out["coord_logits"][:, :, axis][xyz], + targets.coord_indices[:, :, axis][xyz]) + if bool(xyz.any()) else out["hidden"].sum() * 0.0 + for axis in range(3) + ) / 3.0 + ptr_a = selected(RigTokenKindV6.JOINT, RigTokenKindV6.FRAME, + RigTokenKindV6.CONTROL, RigTokenKindV6.CONTROL_TARGET, + RigTokenKindV6.CHAIN, RigTokenKindV6.CONSTRAINT, + RigTokenKindV6.CONSTRAINT_INPUT, RigTokenKindV6.CONSTRAINT_OUTPUT, + RigTokenKindV6.ANCHOR) + ptr_b = selected(RigTokenKindV6.CHAIN) | ( + selected(RigTokenKindV6.ANCHOR) & (targets.anchor_parent_kind == 1) + ) + losses["pointer"] = (ce("pointer_a_logits", (targets.pointer_a + 1).clamp_min(0), ptr_a) + + ce("pointer_b_logits", (targets.pointer_b + 1).clamp_min(0), ptr_b)) + control = selected(RigTokenKindV6.CONTROL) + constraint = selected(RigTokenKindV6.CONSTRAINT) + losses["control_kind"] = ce("control_kind_logits", targets.aux_type_ids.clamp_min(0), control) + losses["constraint_operator"] = ce("constraint_operator_logits", targets.aux_type_ids.clamp_min(0), constraint) + frame = selected(RigTokenKindV6.FRAME) + if bool(frame.any()): + losses["dof"] = F.binary_cross_entropy_with_logits(out["dof_logits"][frame], targets.dof_masks[frame]) + losses["frame"] = F.smooth_l1_loss(out["frame_6d"][frame], targets.frame_6d[frame]) + else: + losses["dof"] = losses["frame"] = out["hidden"].sum() * 0.0 + losses["constraint_scalar"] = ( + F.smooth_l1_loss(out["scalar_params"][constraint], targets.scalar_params[constraint]) + if bool(constraint.any()) else out["hidden"].sum() * 0.0 + ) + anchor = selected(RigTokenKindV6.ANCHOR) + losses["anchor_parent"] = ce("anchor_parent_logits", targets.anchor_parent_kind.clamp_min(0), anchor) + code_mask = selected(RigTokenKindV6.SKIN_CODE) + if bool(code_mask.any()): + logits = model.skin_code_logits(out["hidden"][code_mask]) + losses["skin_code"] = F.cross_entropy(logits, targets.skin_code_ids[code_mask]) + else: + losses["skin_code"] = out["hidden"].sum() * 0.0 + losses["total"] = sum(losses.values()) + return losses + + +def assert_axis_mira_v7_binding(axis: AXISRigProgramConfigV7, mira: MIRACodecConfigV7) -> None: + if axis.skin_vocab_size != mira.vocab_size: + raise ValueError("AXIS_MIRA_V7_VOCAB_BINDING_DRIFT") + if axis.tokens_per_joint != mira.tokens_per_joint: + raise ValueError("AXIS_MIRA_V7_TOKEN_CARDINALITY_DRIFT") + + +__all__ = [ + "AXIS_V7_ARCHITECTURE_ID", "AXISRigProgramConfigV7", + "AXISRigProgramModelV7", "axis_v7_teacher_forced_objective", + "assert_axis_mira_v7_binding", +] diff --git a/models/mira/skin_attention_codec_v7.py b/models/mira/skin_attention_codec_v7.py new file mode 100644 index 000000000..0b5472736 --- /dev/null +++ b/models/mira/skin_attention_codec_v7.py @@ -0,0 +1,123 @@ +from __future__ import annotations + +"""MIRA V7 learned-query FSQ-CVAE research candidate. + +Separately implemented after auditing official SkinTokens attention encoders: +https://github.com/VAST-AI-Research/SkinTokens/blob/main/src/model/skin_vae/autoencoders/skin_fsq_cvae_model.py + +- Active-support sampled per-joint surface context (TRAIN target only). +- Learned query attention compresses local influence distributions into T tokens. +- Straight-through FSQ and conditional carrier-native decoding inherited from V7. +- No source body-part naming, no segmentation labels, no graphics target. +- Memory-bounded sampling: B x J x K x H, not B x J x N x H. +This is NOT a port or checkpoint-compatible clone of upstream SkinTokens. +""" + +from dataclasses import dataclass, asdict + +import torch +from torch import nn, Tensor +import torch.nn.functional as F + +from .skin_fsq_codec_v7 import ( + MIRASkinCodecV7, MIRACodecConfigV7, + MIRA_V7_ARCHITECTURE_ID, +) + +MIRA_V7_ATTENTION_ARCHITECTURE_ID = ( + "RealSaS.MIRA.AttentionCarrierCVAE.FSQ.AuthorizedMechanics.v7_0" +) + + +@dataclass(frozen=True) +class MIRAAttentionConfigV7(MIRACodecConfigV7): + attention_heads: int = 8 + attention_layers: int = 4 + active_samples_per_joint: int = 128 + attention_dropout: float = 0.0 + + def __post_init__(self): + super().__post_init__() + if (self.attention_heads < 1 or self.attention_layers < 1 + or self.active_samples_per_joint < 1 + or self.hidden_dim % self.attention_heads): + raise ValueError("MIRA_V7_ATTENTION_CONFIG_INVALID") + if not 0.0 <= self.attention_dropout < 1.0: + raise ValueError("MIRA_V7_ATTENTION_DROPOUT_INVALID") + + +class _JointLocalAttentionLayer(nn.Module): + def __init__(self, width: int, heads: int, dropout: float): + super().__init__() + self.norm_q = nn.LayerNorm(width) + self.norm_kv = nn.LayerNorm(width) + self.attn = nn.MultiheadAttention(width, heads, dropout=dropout, batch_first=True) + self.norm_ffn = nn.LayerNorm(width) + self.ffn = nn.Sequential( + nn.Linear(width, 4 * width), nn.GELU(), + nn.Dropout(dropout), nn.Linear(4 * width, width), + ) + + def forward(self, queries: Tensor, sampled_context: Tensor) -> Tensor: + q = self.norm_q(queries) + kv = self.norm_kv(sampled_context) + value, _ = self.attn(q, kv, kv, need_weights=False) + out = queries + value + return out + self.ffn(self.norm_ffn(out)) + + +class MIRAAttentionSkinCodecV7(MIRASkinCodecV7): + architecture_id = MIRA_V7_ATTENTION_ARCHITECTURE_ID + + def __init__(self, config: MIRAAttentionConfigV7 = MIRAAttentionConfigV7()): + super().__init__(config) + h, t = config.hidden_dim, config.tokens_per_joint + self.learned_queries = nn.Parameter(torch.empty(t, h)) + nn.init.normal_(self.learned_queries, std=0.02) + self.local_weight_embedding = nn.Sequential(nn.Linear(1, h), nn.GELU(), nn.Linear(h, h)) + self.attention = nn.ModuleList([ + _JointLocalAttentionLayer(h, config.attention_heads, config.attention_dropout) + for _ in range(config.attention_layers) + ]) + self.attention_post_norm = nn.LayerNorm(h) + self.attention_latent = nn.Linear(h, config.latent_dim) + # Note: prior V7 regional latent encoder remains a registered baseline, + # but not on this attention candidate's active forward path. + + def _active_samples(self, target: Tensor) -> tuple[Tensor, Tensor]: + # Deterministic max-influence subset per joint. This is teacher-only + # sampling; downstream decoder does NOT need teacher weights. + _, n, _ = target.shape + count = min(n, self.config.active_samples_per_joint) + selected_weights, indices = torch.topk( + target.transpose(1, 2).contiguous(), k=count, + dim=-1, largest=True, sorted=True, + ) + return indices, selected_weights + + def _latent(self, target: Tensor, vertices: Tensor, joints: Tensor) -> Tensor: + b, n, j = self._shape(vertices, joints) + sampled_ids, weights = self._active_samples(target) + h, t = self.config.hidden_dim, self.config.tokens_per_joint + vertex_embed = self.vertex_encoder(vertices) + joint_embed = self.joint_encoder(joints) + selected = torch.gather( + vertex_embed[:, None, :, :].expand(b, j, n, h), + 2, sampled_ids[..., None].expand(b, j, sampled_ids.shape[-1], h), + ) + context = selected + self.local_weight_embedding(weights[..., None]) + q = (self.learned_queries[None, None, :, :] + + joint_embed[:, :, None, :]) + query_batch = q.reshape(b * j, t, h) + sampled_context = context.reshape(b * j, sampled_ids.shape[-1], h) + for layer in self.attention: + query_batch = layer(query_batch, sampled_context) + latent = self.attention_latent(self.attention_post_norm(query_batch)) + return latent.reshape(b, j, t, self.config.latent_dim) + + +__all__ = [ + "MIRA_V7_ATTENTION_ARCHITECTURE_ID", + "MIRAAttentionConfigV7", + "MIRAAttentionSkinCodecV7", +] diff --git a/models/mira/skin_fsq_codec_v7.py b/models/mira/skin_fsq_codec_v7.py new file mode 100644 index 000000000..b7fc87210 --- /dev/null +++ b/models/mira/skin_fsq_codec_v7.py @@ -0,0 +1,216 @@ +from __future__ import annotations + +"""RealSaS MIRA V7: large-capacity, carrier-native discrete skin field codec. + +Research implementation; NOT a trained ModelRelease. No SkinTokens code is copied. +The finite-scalar representation is independently implemented using the +existing RealSaS quantizer. The critical V6 reconstruction-gradient severance is +fixed: reconstruction decodes straight-through continuous quantized latents, +while inference decodes exactly the same hard codes from integer token IDs. + +Complexity of encoding is O(B*(N+J)*H + B*N*J), NOT O(B*N*J*H). +This matters for dense Corpus370 carrier meshes and nonhumanoid skeletons. +""" + +from dataclasses import asdict, dataclass +from math import prod, sqrt + +import torch +from torch import Tensor, nn +import torch.nn.functional as F + +from models.mira.skin_fsq_codec_v6 import FiniteScalarQuantizer + + +MIRA_V7_ARCHITECTURE_ID = "RealSaS.MIRA.CarrierNativeFSQ.STGradient.ActiveField.v7_0" + + +@dataclass(frozen=True) +class MIRACodecConfigV7: + vertex_feature_dim: int = 7 + joint_feature_dim: int = 25 + hidden_dim: int = 384 + tokens_per_joint: int = 4 + levels: tuple[int, ...] = (8, 8, 8, 5, 5, 5) + max_vertices: int = 131072 + max_joints: int = 256 + importance_strength: float = 4.0 + active_threshold: float = 0.02 + eps: float = 1e-8 + + def __post_init__(self): + if min(self.vertex_feature_dim, self.joint_feature_dim, self.hidden_dim, + self.tokens_per_joint, self.max_vertices, self.max_joints) < 1: + raise ValueError("MIRA_V7_CONFIG_NONPOSITIVE") + if not self.levels or any(x < 2 for x in self.levels): + raise ValueError("MIRA_V7_FSQ_LEVELS_INVALID") + if not (0 <= self.active_threshold < 1) or self.importance_strength < 0: + raise ValueError("MIRA_V7_IMPORTANCE_CONFIG_INVALID") + + @property + def vocab_size(self) -> int: + return prod(self.levels) + + @property + def latent_dim(self) -> int: + return len(self.levels) + + def to_dict(self) -> dict: + return asdict(self) + + +class MIRASkinCodecV7(nn.Module): + architecture_id = MIRA_V7_ARCHITECTURE_ID + + def __init__(self, config: MIRACodecConfigV7 = MIRACodecConfigV7()): + super().__init__() + self.config = config + h = config.hidden_dim + self.quantizer = FiniteScalarQuantizer(config.levels) + self.vertex_encoder = nn.Sequential( + nn.Linear(config.vertex_feature_dim, h), nn.GELU(), + nn.Linear(h, h), nn.GELU(), nn.LayerNorm(h), + ) + self.joint_encoder = nn.Sequential( + nn.Linear(config.joint_feature_dim, h), nn.GELU(), + nn.Linear(h, h), nn.GELU(), nn.LayerNorm(h), + ) + self.latent_encoder = nn.Sequential( + nn.Linear(h * 3, h * 2), nn.GELU(), + nn.Linear(h * 2, h), + nn.GELU(), nn.Linear(h, config.tokens_per_joint * config.latent_dim), + ) + self.vertex_decoder = nn.Sequential( + nn.Linear(config.vertex_feature_dim, h), nn.GELU(), + nn.Linear(h, h), nn.LayerNorm(h), + ) + self.joint_decoder = nn.Sequential( + nn.Linear(config.joint_feature_dim, h), nn.GELU(), + nn.Linear(h, h), nn.LayerNorm(h), + ) + self.code_decoder = nn.Sequential( + nn.Linear(config.tokens_per_joint * config.latent_dim, h * 2), + nn.GELU(), nn.Linear(h * 2, h), + ) + self.joint_bias = nn.Linear(h, 1) + + def _shape(self, vertex_features: Tensor, joint_features: Tensor) -> tuple[int, int, int]: + if vertex_features.ndim != 3 or joint_features.ndim != 3: + raise ValueError("MIRA_V7_FEATURE_RANK") + b, n, fv = vertex_features.shape + bj, j, fj = joint_features.shape + if b != bj or fv != self.config.vertex_feature_dim or fj != self.config.joint_feature_dim: + raise ValueError("MIRA_V7_FEATURE_SHAPE") + if b < 1 or n < 1 or j < 1 or n > self.config.max_vertices or j > self.config.max_joints: + raise ValueError("MIRA_V7_CARRIER_RESOURCE_BOUND") + if not torch.isfinite(vertex_features).all() or not torch.isfinite(joint_features).all(): + raise ValueError("MIRA_V7_FEATURE_NONFINITE") + return b, n, j + + def _target(self, weights: Tensor, b: int, n: int, j: int) -> Tensor: + if weights.shape != (b, n, j) or not torch.isfinite(weights).all(): + raise ValueError("MIRA_V7_WEIGHT_SHAPE_OR_NONFINITE") + if torch.any(weights < 0): + raise ValueError("MIRA_V7_WEIGHT_NEGATIVE") + mass = weights.sum(dim=-1, keepdim=True) + if torch.any(mass <= self.config.eps): + raise ValueError("MIRA_V7_WEIGHT_UNCOVERED_VERTEX") + return weights / mass + + def _latent(self, target: Tensor, vertices: Tensor, joints: Tensor) -> Tensor: + # No B*J*N*H tensor: only B*N*H, B*J*H and B*N*J. + encoded_v = self.vertex_encoder(vertices) + encoded_j = self.joint_encoder(joints) + support_mass = target.sum(dim=1).clamp_min(self.config.eps) # [B,J] + regional = torch.einsum("bnj,bnh->bjh", target, encoded_v) + regional = regional / support_mass.unsqueeze(-1) + global_state = encoded_v.mean(dim=1, keepdim=True).expand_as(encoded_j) + latents = self.latent_encoder(torch.cat((regional, encoded_j, global_state), dim=-1)) + return latents.reshape( + joints.shape[0], joints.shape[1], + self.config.tokens_per_joint, self.config.latent_dim, + ) + + def _hard_and_straight_through(self, latent: Tensor) -> tuple[Tensor, Tensor, Tensor]: + # Straight-through reconstruction gradient must reach the latent encoder. + bounded = torch.tanh(latent) + indices = self.quantizer._quantize_indices(latent) + hard = self.quantizer.indices_to_codes(indices).to(dtype=latent.dtype) + straight_through = bounded + (hard - bounded).detach() + ids = self.quantizer.indices_to_token_ids(indices) + return straight_through, hard, ids + + def _decode_codes(self, codes: Tensor, vertices: Tensor, joints: Tensor) -> Tensor: + b, n, j = self._shape(vertices, joints) + if codes.shape != (b, j, self.config.tokens_per_joint, self.config.latent_dim): + raise ValueError("MIRA_V7_CODE_SHAPE") + latent = codes.flatten(start_dim=2) + vertex_state = self.vertex_decoder(vertices) + joint_state = self.joint_decoder(joints) + self.code_decoder(latent) + logits = torch.einsum("bnh,bjh->bnj", vertex_state, joint_state) + logits = logits / sqrt(float(self.config.hidden_dim)) + logits = logits + self.joint_bias(joint_state).transpose(1, 2) + return torch.softmax(logits, dim=-1) + + def encode_tokens(self, weights: Tensor, vertex_features: Tensor, + joint_features: Tensor) -> Tensor: + b, n, j = self._shape(vertex_features, joint_features) + target = self._target(weights, b, n, j) + _, _, ids = self._hard_and_straight_through( + self._latent(target, vertex_features, joint_features) + ) + return ids + + def decode_tokens(self, token_ids: Tensor, vertex_features: Tensor, + joint_features: Tensor) -> Tensor: + b, _, j = self._shape(vertex_features, joint_features) + if token_ids.shape != (b, j, self.config.tokens_per_joint): + raise ValueError("MIRA_V7_TOKEN_SHAPE") + hard = self.quantizer.token_ids_to_codes(token_ids).to(dtype=vertex_features.dtype) + return self._decode_codes(hard, vertex_features, joint_features) + + def forward(self, weights: Tensor, vertex_features: Tensor, + joint_features: Tensor, *, mesh_edge_index: Tensor | None = None) -> dict[str, Tensor]: + b, n, j = self._shape(vertex_features, joint_features) + target = self._target(weights, b, n, j) + latent = self._latent(target, vertex_features, joint_features) + straight_through, hard, ids = self._hard_and_straight_through(latent) + decoded = self._decode_codes(straight_through, vertex_features, joint_features) + abs_error = (decoded - target).abs() + active = (target > self.config.active_threshold).to(abs_error.dtype) + weighted = 1.0 + self.config.importance_strength * active + reconstruction = (abs_error * weighted).sum() / weighted.sum() + quantization = F.mse_loss(torch.tanh(latent), hard.detach()) + # Encourage sparse local influence without rewarding uncovered vertices; + # per-vertex simplex is enforced by softmax. + sparsity = -(decoded.clamp_min(self.config.eps) * + decoded.clamp_min(self.config.eps).log()).sum(dim=-1).mean() + smoothness = decoded.sum() * 0.0 + if mesh_edge_index is not None: + edges = mesh_edge_index.to(device=decoded.device, dtype=torch.long) + if edges.ndim != 2 or edges.shape[0] != 2: + raise ValueError("MIRA_V7_EDGE_SHAPE") + if edges.numel() and (torch.any(edges < 0) or torch.any(edges >= n)): + raise ValueError("MIRA_V7_EDGE_OUT_OF_RANGE") + if edges.shape[1] > 0: + src, dst = edges[0], edges[1] + predicted_delta = decoded[:, src] - decoded[:, dst] + truth_delta = target[:, src] - target[:, dst] + # Teacher-relative gradient consistency, not blind smoothing across seams. + smoothness = F.smooth_l1_loss(predicted_delta, truth_delta) + loss = reconstruction + 0.05 * quantization + 0.001 * sparsity + 0.1 * smoothness + return { + "loss": loss, + "reconstruction_loss": reconstruction, + "quantization_loss": quantization, + "sparsity_loss": sparsity, + "edge_gradient_loss": smoothness, + "mae": abs_error.mean(), + "max_abs": abs_error.amax(), + "simplex_max_abs": (decoded.sum(dim=-1) - 1.0).abs().amax(), + "token_ids": ids, + "reconstructed_weights": decoded, + } + + +__all__ = ["MIRA_V7_ARCHITECTURE_ID", "MIRACodecConfigV7", "MIRASkinCodecV7"] diff --git a/tests/models/test_axis_mira_v7_research.py b/tests/models/test_axis_mira_v7_research.py new file mode 100644 index 000000000..f7689c7cf --- /dev/null +++ b/tests/models/test_axis_mira_v7_research.py @@ -0,0 +1,180 @@ +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +from models.mira.skin_fsq_codec_v7 import ( + MIRACodecConfigV7, MIRASkinCodecV7, +) +from models.axis.rig_program_model_v7 import ( + AXISRigProgramConfigV7, AXISRigProgramModelV7, + assert_axis_mira_v7_binding, axis_v7_teacher_forced_objective, +) +from models.axis.rig_program_model_v6 import ( + AXISSequenceV6, RigProgramEventV6, RigTokenKindV6, +) +from compiler.realsas_compiler_core.mechanics_interaction_v1 import ( + select_mechanics_controls_from_surface_v1, +) +from compiler.realsas_compiler_core.types import QualificationError + + +def test_mira_v7_reconstruction_gradient_and_hard_token_parity(): + torch.manual_seed(701) + cfg = MIRACodecConfigV7( + vertex_feature_dim=7, joint_feature_dim=25, + hidden_dim=32, tokens_per_joint=2, levels=(3, 3), + max_vertices=128, max_joints=8, + ) + model = MIRASkinCodecV7(cfg) + v = torch.randn(2, 12, 7) + j = torch.randn(2, 3, 25) + w = torch.softmax(torch.randn(2, 12, 3), dim=-1) + result = model(w, v, j) + assert result["token_ids"].shape == (2, 3, 2) + assert torch.isfinite(result["loss"]) + assert result["simplex_max_abs"] < 1e-6 + # Critical regression: reconstruction alone MUST train the MIRA encoder. + result["reconstruction_loss"].backward() + grad = model.latent_encoder[-1].weight.grad + assert grad is not None and torch.isfinite(grad).all() + assert float(grad.abs().sum()) > 0.0 + + with torch.no_grad(): + decoded = model.decode_tokens(result["token_ids"], v, j) + assert torch.allclose(decoded, result["reconstructed_weights"], atol=1e-6, rtol=0.0) + + +def test_mira_v7_sparse_positive_weights_not_all_zero(): + model = MIRASkinCodecV7(MIRACodecConfigV7( + hidden_dim=24, tokens_per_joint=2, levels=(3, 3), + max_vertices=24, max_joints=4, + )) + vertices = torch.randn(1, 5, 7) + joints = torch.randn(1, 2, 25) + weights = torch.tensor([[[1., 0.], [0., 1.], [1., 0.], [0.5, 0.5], [0., 1.]]]) + result = model(weights, vertices, joints, mesh_edge_index=torch.tensor([[0, 1, 2], [1, 2, 3]])) + assert torch.isfinite(result["edge_gradient_loss"]) + with pytest.raises(ValueError, match="UNCOVERED"): + model(torch.zeros_like(weights), vertices, joints) + + +def test_axis_v7_large_vocab_contract_and_sparse_skin_loss(): + cfg = AXISRigProgramConfigV7( + scene_dim=47, hidden_dim=32, layers=1, heads=4, + feedforward_mult=2, max_sequence_length=20, + max_pointer_entities=6, skin_vocab_size=9, tokens_per_joint=2, + ) + assert_axis_mira_v7_binding(cfg, MIRACodecConfigV7( + hidden_dim=32, tokens_per_joint=2, levels=(3, 3) + )) + with pytest.raises(ValueError, match="VOCAB"): + assert_axis_mira_v7_binding(cfg, MIRACodecConfigV7( + hidden_dim=32, tokens_per_joint=2, levels=(5, 5) + )) + model = AXISRigProgramModelV7(cfg) + inputs = AXISSequenceV6.from_events(( + RigProgramEventV6(RigTokenKindV6.BOS), + RigProgramEventV6(RigTokenKindV6.JOINT, coord_indices=(1, 2, 3), pointer_a=-1), + ), device=torch.device("cpu")) + targets = AXISSequenceV6.from_events(( + RigProgramEventV6(RigTokenKindV6.JOINT, coord_indices=(1, 2, 3), pointer_a=-1), + RigProgramEventV6(RigTokenKindV6.SKIN_CODE, skin_code_id=4), + ), device=torch.device("cpu")) + result = axis_v7_teacher_forced_objective( + model, torch.randn(1, 5, 47), inputs, targets, + ) + assert torch.isfinite(result["total"]) + assert result["skin_code"] > 0 + result["total"].backward() + assert model.skin_embedding.weight.grad is not None + assert model.kind_head.weight.grad is not None + + +def _mechanics_fixture(): + joints = ( + SimpleNamespace(canonical_joint_id="J0"), + SimpleNamespace(canonical_joint_id="J1"), + ) + controls = ( + SimpleNamespace(canonical_control_id="FK1", control_kind="FK", + target_canonical_joint_ids=("J1",)), + SimpleNamespace(canonical_control_id="IK1", control_kind="IK_EFFECTOR", + target_canonical_joint_ids=()), + SimpleNamespace(canonical_control_id="GLOBAL", control_kind="GLOBAL", + target_canonical_joint_ids=()), + ) + chains = (SimpleNamespace(canonical_chain_id="CH", canonical_joint_ids=("J0", "J1")),) + constraints = ( + SimpleNamespace(operator_id="IK", canonical_chain_id="CH", + input_canonical_control_ids=("IK1",)), + ) + rig = SimpleNamespace( + mesh_lineage_hash="mesh", deform_structure_hash="deform", + rig_lineage_hash="rig", joints=joints, controls=controls, + chains=chains, constraints=constraints, + ) + rows = ( + SimpleNamespace(canonical_mesh_vertex_id="V0", influences=(("J0", 0.1), ("J1", 0.9))), + SimpleNamespace(canonical_mesh_vertex_id="V1", influences=(("J1", 1.0),)), + SimpleNamespace(canonical_mesh_vertex_id="V2", influences=(("J0", 0.2), ("J1", 0.8))), + ) + skin = SimpleNamespace( + mesh_lineage_hash="mesh", deform_structure_hash="deform", + skin_lineage_hash="skin", rows=rows, + ) + return rig, skin + + +def test_mechanics_pick_is_anonymous_readonly_and_declared_only(): + rig, skin = _mechanics_fixture() + out = select_mechanics_controls_from_surface_v1( + rig, skin, triangle_vertex_ids=("V0", "V1", "V2"), + barycentric=(0.2, 0.5, 0.3), + ) + assert out.status == "CANDIDATES" + assert tuple(c.canonical_control_id for c in out.candidates) == ("FK1", "IK1") + assert out.candidates[0].influence_score > 0.9 + assert rig.controls[1].target_canonical_joint_ids == () # never synthesized + with pytest.raises(QualificationError, match="SIMPLEX"): + select_mechanics_controls_from_surface_v1( + rig, skin, triangle_vertex_ids=("V0", "V1", "V2"), + barycentric=(0.5, 0.5, 0.5), + ) + with pytest.raises(QualificationError, match="LINEAGE"): + bad_skin = SimpleNamespace(**{**skin.__dict__, "deform_structure_hash": "wrong"}) + select_mechanics_controls_from_surface_v1( + rig, bad_skin, triangle_vertex_ids=("V0", "V1", "V2"), + barycentric=(0.2, 0.5, 0.3), + ) + + +def test_mira_v7_attention_cv_ae_active_sampling_gradient_and_token_parity(): + from models.mira.skin_attention_codec_v7 import ( + MIRAAttentionConfigV7, MIRAAttentionSkinCodecV7, + ) + torch.manual_seed(711) + cfg = MIRAAttentionConfigV7( + hidden_dim=32, attention_heads=4, attention_layers=2, + active_samples_per_joint=5, tokens_per_joint=2, + levels=(3, 3), max_vertices=128, max_joints=8, + ) + model = MIRAAttentionSkinCodecV7(cfg) + vertices = torch.randn(2, 19, 7) + joints = torch.randn(2, 4, 25) + truth = torch.softmax(torch.randn(2, 19, 4) * 2, dim=-1) + indices, weights = model._active_samples(truth) + assert indices.shape == (2, 4, 5) + assert torch.all(weights[:, :, :-1] >= weights[:, :, 1:]) + out = model(truth, vertices, joints) + assert out["token_ids"].shape == (2, 4, 2) + out["reconstruction_loss"].backward() + assert model.learned_queries.grad is not None + assert float(model.learned_queries.grad.abs().sum()) > 0.0 + assert model.attention_latent.weight.grad is not None + assert float(model.attention_latent.weight.grad.abs().sum()) > 0.0 + with torch.no_grad(): + decoded = model.decode_tokens(out["token_ids"], vertices, joints) + assert torch.allclose(out["reconstructed_weights"], decoded, atol=1e-6, rtol=0.0)