"""Transformers configuration for the portable Isaac-0.5 VLA repository.""" from __future__ import annotations import copy from collections.abc import Mapping from typing import Any from transformers import Qwen3_5MoeConfig _ISAAC05_ARCHITECTURES = ["Isaac05ForConditionalGeneration"] _ISAAC05_AUTO_MAP = { "AutoConfig": "configuration_isaac05.Isaac05Config", "AutoModelForCausalLM": "modeling_isaac05.Isaac05ForConditionalGeneration", "AutoProcessor": "processing_isaac05.Isaac05Processor", } _PRODUCTION_COORD_TOKENS = {"enabled": True, "offset": 248_320, "size": 1_001} _PRODUCTION_FAST_TOKENS = { "enabled": True, "tokenizer": "physical-intelligence/fast", "offset": 249_321, "size": 2_048, } _PRODUCTION_ARTIFACT = { "schema_version": 1, "artifact_kind": "trained_policy", "tensor_count": 1_616, "tensor_bytes": 142_894_027_456, } _PRODUCTION_VECTOR_ENCODER = { "type": "linear_silu_linear", "max_states": 128, "hidden_dim": 2_048, "output_dim": 2_048, "bias": False, } _PRODUCTION_ACTION_EXPERT = { "action_dim": 64, "action_horizon": 64, "num_layers": 36, "hidden_dim": 768, "num_heads": 8, "mlp_ratio": 4.0, "num_inference_steps": 10, "timestep_sampling_alpha": 1.5, "timestep_sampling_beta": 1.0, "timestep_sampling_scale": 0.999, "timestep_sampling_offset": 0.001, "train_samples_per_chunk": 8, "timestep_embed_dim": 256, "rtc_max_delay_steps": 12, "rtc_probability": 0.5, "rtc_delay_sampling": "poisson", "rtc_poisson_mean": 5.0, "mask_padded_action_rows": True, "drop_action_dim_overflow": False, "ffn_multiple_of": 256, "qk_norm": True, "qk_norm_eps": 1e-6, "rope": True, "context_layer_norm": True, "causal_attn": False, "k_batched_cross_attn": True, "k_batched_cross_attn_backend": "flash_gqa", "schema_version": 1, "type": "dit", } _PRODUCTION_MTP = { "present": False, "physical_layers": 0, "rollout_steps": 0, "action_runtime": "exclude", } def _copy_mapping(value: Mapping[str, Any] | None, *, name: str) -> dict[str, Any]: if not isinstance(value, Mapping): raise ValueError(f"Isaac05Config {name} must be a JSON object.") return copy.deepcopy(dict(value)) def _require_exact(value: Mapping[str, Any], expected: Mapping[str, Any], *, name: str) -> None: if dict(value) != dict(expected): raise ValueError(f"Isaac05Config {name} does not match the portable artifact contract.") def _reserved_token_range(value: Mapping[str, Any], *, name: str) -> range: if value.get("enabled") is not True: raise ValueError(f"Isaac05Config {name}.enabled must be true.") offset = value.get("offset") size = value.get("size") if not isinstance(offset, int) or isinstance(offset, bool) or offset < 0: raise ValueError(f"Isaac05Config {name}.offset must be a non-negative integer.") if not isinstance(size, int) or isinstance(size, bool) or size <= 0: raise ValueError(f"Isaac05Config {name}.size must be a positive integer.") return range(offset, offset + size) class Isaac05Config(Qwen3_5MoeConfig): """Portable Isaac-0.5 configuration for the published checkpoint.""" model_type = "isaac_0_5" has_no_defaults_at_init = True def __init__( self, *, isaac05_artifact: Mapping[str, Any] | None = None, isaac05_coord_tokens: Mapping[str, Any] | None = None, isaac05_fast_tokens: Mapping[str, Any] | None = None, isaac05_vla: Mapping[str, Any] | None = None, storage_dtype: str = "float32", runtime_dtype: str = "bfloat16", max_sequence_length: int = 262_144, vision_token: str = "<|image_pad|>", vision_rescale_factor: float = 1 / 255, isaac05_test_only_reduced_geometry: bool = False, **kwargs: Any, ) -> None: architectures = kwargs.pop("architectures", _ISAAC05_ARCHITECTURES) auto_map = kwargs.pop("auto_map", _ISAAC05_AUTO_MAP) if architectures != _ISAAC05_ARCHITECTURES: raise ValueError(f"Isaac05Config architectures must be {_ISAAC05_ARCHITECTURES!r}.") if auto_map != _ISAAC05_AUTO_MAP: raise ValueError("Isaac05Config auto_map does not match the portable repository API.") artifact = _copy_mapping(isaac05_artifact, name="isaac05_artifact") coord_tokens = _copy_mapping(isaac05_coord_tokens, name="isaac05_coord_tokens") fast_tokens = _copy_mapping(isaac05_fast_tokens, name="isaac05_fast_tokens") vla = _copy_mapping(isaac05_vla, name="isaac05_vla") if storage_dtype != "float32": raise ValueError("Isaac05Config storage_dtype must be 'float32'.") if runtime_dtype != "bfloat16": raise ValueError("Isaac05Config runtime_dtype must be 'bfloat16'.") if max_sequence_length <= 0: raise ValueError("Isaac05Config max_sequence_length must be positive.") if not vision_token: raise ValueError("Isaac05Config vision_token must not be empty.") if vision_rescale_factor <= 0: raise ValueError("Isaac05Config vision_rescale_factor must be positive.") coord_range = _reserved_token_range(coord_tokens, name="isaac05_coord_tokens") fast_range = _reserved_token_range(fast_tokens, name="isaac05_fast_tokens") if coord_range.start < fast_range.stop and fast_range.start < coord_range.stop: raise ValueError("Isaac05Config reserved token ranges overlap.") if not isaac05_test_only_reduced_geometry: _require_exact(artifact, _PRODUCTION_ARTIFACT, name="isaac05_artifact") _require_exact(coord_tokens, _PRODUCTION_COORD_TOKENS, name="isaac05_coord_tokens") _require_exact(fast_tokens, _PRODUCTION_FAST_TOKENS, name="isaac05_fast_tokens") _require_exact( _copy_mapping(vla.get("vector_encoder"), name="isaac05_vla.vector_encoder"), _PRODUCTION_VECTOR_ENCODER, name="isaac05_vla.vector_encoder", ) _require_exact( _copy_mapping(vla.get("action_expert"), name="isaac05_vla.action_expert"), _PRODUCTION_ACTION_EXPERT, name="isaac05_vla.action_expert", ) _require_exact( _copy_mapping(vla.get("mtp"), name="isaac05_vla.mtp"), _PRODUCTION_MTP, name="isaac05_vla.mtp", ) if vla.get("schema_version") != 1: raise ValueError("Isaac05Config isaac05_vla.schema_version must be 1.") if vla.get("state_dict_schema") != "pr3154_v1": raise ValueError("Isaac05Config isaac05_vla.state_dict_schema must be 'pr3154_v1'.") if vla.get("rmsnorm_weight_convention") != "zero_centered_1_plus_weight": raise ValueError( "Isaac05Config isaac05_vla.rmsnorm_weight_convention must be 'zero_centered_1_plus_weight'." ) super().__init__(architectures=architectures, auto_map=auto_map, **kwargs) self.isaac05_artifact = artifact self.isaac05_coord_tokens = coord_tokens self.isaac05_fast_tokens = fast_tokens self.isaac05_vla = vla vector_encoder = _copy_mapping(vla.get("vector_encoder"), name="isaac05_vla.vector_encoder") action_expert = _copy_mapping(vla.get("action_expert"), name="isaac05_vla.action_expert") self.vector_max_states = int(vector_encoder["max_states"]) self.action_expert = action_expert self.storage_dtype = storage_dtype self.runtime_dtype = runtime_dtype self.max_sequence_length = int(max_sequence_length) self.vision_token = vision_token self.vision_rescale_factor = float(vision_rescale_factor) self.isaac05_test_only_reduced_geometry = bool(isaac05_test_only_reduced_geometry)