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