diff --git a/agent_context/MAP.yaml b/agent_context/MAP.yaml index 548b67337c..39bed0f05d 100644 --- a/agent_context/MAP.yaml +++ b/agent_context/MAP.yaml @@ -273,7 +273,7 @@ topics: aliases: [neural dynamics, neural robot dynamics, NeRD, RigidFormer, 神经动力学, 动力学模型训练, NeRD 示例, RigidFormer 示例] keywords: [NeRDStateCodec, NeRDModel, DynamicsGenerationCfg, DynamicsShardWriter, DynamicsWindowDataset, DynamicsTrainer, embodichain_models, - learning/nerd, nerd reference assets, RigidFormerModel, + learning/nerd, nerd reference assets, RigidFormerModel, MOViSequenceSource, build_rigidformer_trainer, 官方训练精度对比] paths: [topics/neural-dynamics/neural-dynamics.md] source_of_truth: diff --git a/agent_context/topics/neural-dynamics/neural-dynamics.md b/agent_context/topics/neural-dynamics/neural-dynamics.md index 1a05799ba9..1fbe47dcd0 100644 --- a/agent_context/topics/neural-dynamics/neural-dynamics.md +++ b/agent_context/topics/neural-dynamics/neural-dynamics.md @@ -47,11 +47,18 @@ The native family lives in `model.py` owns geometry/configuration; `_native/` preserves licensed author code. Checkpoint import validates recorded architecture and named tensors. The community `rigidformer` package is not a runtime dependency. - -Integration dt=1 differs from physical seconds. Geometry prediction requires -CUDA and at least two valid objects; construction and weight inspection support -CPU. Attention backend selection is explicit. The model predicts geometry -without action inputs and does not qualify state/action gradients. +`data.py` reads verified packed scenes; `supervised.py` separates shifted +conditioning and anchor targets; `sl_recipe.py` constructs the shared trainer +with committed sampling, epoch scheduling and EMA participants. + +Integration dt=1 differs from physical seconds. Stored vertices are unscaled; +metadata sizes apply about each origin once. Geometry prediction requires CUDA +and at least two valid objects. Attention backend and source semantics belong +to recipe/recovery identity. Source-owned epoch sampling does not claim the +legacy research DataLoader's worker-RNG trajectory. + +Numerical agreement, recovery and full benchmark/retraining qualification +remain separate evidence requirements. ## Resource and physical identity diff --git a/agent_context/topics/supervised-learning/supervised-learning.md b/agent_context/topics/supervised-learning/supervised-learning.md index f42c6819e3..c12d16714d 100644 --- a/agent_context/topics/supervised-learning/supervised-learning.md +++ b/agent_context/topics/supervised-learning/supervised-learning.md @@ -42,8 +42,9 @@ Objectives own loss masks, normalization, integration and multi-step graph construction. The optimizer does not infer action semantics or detach histories. Under DDP, `LossResult.loss` is a local mean and `weight` is its effective valid supervision count; the learner forms the global weighted mean. Batch size is per -rank, the learning rate is unchanged, and rank-local sampling streams draw with -replacement, so samples may overlap. All ranks call fitting, saving and recovery; +rank and the learning rate is unchanged. Sampling belongs to the source: tensor +sources draw with replacement, while epoch-traversal sources may shard windows. +All ranks call fitting, saving and recovery; only rank 0 writes artifacts or runs validation through the unwrapped model. Schema, action/frame/timestep and numerical-policy identifiers belong in each recipe's checkpoint signature. Tensor sources additionally bind content identity. @@ -57,8 +58,10 @@ changes. Live buffers and delivery credits are not restored: such sources may train/save snapshots, but exact continuation is rejected. Acquisition, waits, refill and delivery debits stay source/caller-owned; encoding failures do not refund a delivered context. -Distributed checkpoints retain each rank's RNG, sampler and buffers and require -the same world size for exact recovery. Inside rank 0's validation callback only, +Distributed checkpoints retain each rank's RNG, sampler, buffers and participant +state and require the same world size for exact recovery. Rank-local snapshots +are transferred to rank 0 in bounded chunks to limit device memory. Inside rank +0's validation callback only, a save uses the previously gathered committed-update snapshot without a new collective; ordinary save calls require every rank to participate. @@ -80,3 +83,9 @@ qualification belongs to each model's acceptance recipe. The distributed tests exercise weighted global updates and fixed-world-size recovery with CPU/Gloo; sequence tests reuse real CPU source fixtures for masks, provenance, owned output, quota and recovery admission. Examples use launch smoke checks, not tutorial tests. + +Factories may register ordered `SupervisedTrainingState` participants for source +cursors, schedules and model averages. They commit after optimizer success, +produce owned rank-local snapshots, validate before live restoration and restore +after validation callbacks. Empty participants retain existing checkpoint +compatibility. Model-specific participant math stays in its model family. diff --git a/docs/source/api_reference/embodichain/embodichain.learning.sl.rst b/docs/source/api_reference/embodichain/embodichain.learning.sl.rst index 0d979cbfa1..555850f0f1 100644 --- a/docs/source/api_reference/embodichain/embodichain.learning.sl.rst +++ b/docs/source/api_reference/embodichain/embodichain.learning.sl.rst @@ -30,8 +30,8 @@ The example factory uses a small masked MLP regression corpus generated locally. For ``WORLD_SIZE > 1``, the workflow initializes the process group before invoking the recipe factory. It resolves ``trainer.device`` to ``cpu`` or ``cuda:LOCAL_RANK``. ``trainer.batch_size`` is per worker; the global number of -drawn examples is ``batch_size * world_size``. Each worker samples independently -with replacement, so examples can overlap. The learning rate is not scaled +drawn examples is ``batch_size * world_size``. Sampling semantics belong to each source: the tensor example samples independently +with replacement, while explicit offline adapters may shard an epoch traversal. The learning rate is not scaled automatically. Objectives return a local mean loss and its valid supervision count in @@ -99,6 +99,7 @@ embodichain.learning.sl LossResult SupervisedBatch SupervisedObjective + SupervisedTrainingState SupervisedLearner SupervisedSource TensorBatchSource @@ -116,6 +117,9 @@ embodichain.learning.sl .. autoclass:: SupervisedObjective :members: +.. autoclass:: SupervisedTrainingState + :members: + .. autoclass:: SupervisedLearner :members: @@ -234,6 +238,7 @@ embodichain.learning.sl.types SupervisedBatch LossResult SupervisedObjective + SupervisedTrainingState .. autoclass:: SupervisedBatch :members: @@ -244,6 +249,9 @@ embodichain.learning.sl.types .. autoclass:: SupervisedObjective :members: +.. autoclass:: SupervisedTrainingState + :members: + embodichain.learning.sl.workflow -------------------------------- @@ -265,3 +273,17 @@ Concrete models own their physical codecs, objectives and source adapters in the :doc:`official model library `. The NeRD recipe connects committed trajectory windows to this common trainer without loading simulation or requiring the example runner. + +Recoverable update participants +------------------------------- + +Factories may pass ordered ``training_states`` implementing +``SupervisedTrainingState``. Participants commit after optimizer success; their +owned state is saved per rank, validated before live recovery mutation, and +restored after validation callbacks. This supports data cursors and model +averages without model-specific trainer subclasses. Default empty participants +retain compatibility with existing checkpoints. + +Checkpoint formats 3/4 mark nonempty training participants or nondefault buffer +broadcasting, so older readers reject required state instead of discarding it. +Default recipes continue to write/read formats 1/2. diff --git a/docs/source/api_reference/model_rigidformer.rst b/docs/source/api_reference/model_rigidformer.rst index 61db20c7c6..803c9730b1 100644 --- a/docs/source/api_reference/model_rigidformer.rst +++ b/docs/source/api_reference/model_rigidformer.rst @@ -14,7 +14,7 @@ weight inspection can run on CPU. The geometry integration unit dt=1 is distinct from physical intervals step_size / 240 seconds. Attention uses explicit SDPA or an optional compatible Flash-Attention build. -Backend selection is explicit in model configuration. Inference rollout restores +Backend identity is part of recipe/recovery semantics. Inference rollout restores the caller's module modes and does not retain a gradient graph. Public model family @@ -24,10 +24,14 @@ Public model family .. autosummary:: - RIGIDFORMER_REVISION RigidFormerCfg RigidFormerModel load_rigidformer_checkpoint + export_rigidformer_checkpoint + MOViSequenceSource + RigidFormerObjective + rigidformer_batch + build_rigidformer_trainer .. autodata:: RIGIDFORMER_REVISION @@ -39,6 +43,18 @@ Public model family .. autofunction:: load_rigidformer_checkpoint +.. autofunction:: export_rigidformer_checkpoint + +.. autoclass:: MOViSequenceSource + :members: + +.. autoclass:: RigidFormerObjective + :members: + +.. autofunction:: rigidformer_batch + +.. autofunction:: build_rigidformer_trainer + Architecture and geometry contract ---------------------------------- @@ -46,7 +62,6 @@ Architecture and geometry contract .. autosummary:: - RIGIDFORMER_REVISION RigidFormerCfg RigidFormerModel @@ -58,13 +73,78 @@ Architecture and geometry contract .. autoclass:: RigidFormerModel :members: -Official inference weights --------------------------- +Official weights and native SL export +------------------------------------- .. currentmodule:: embodichain_models.dynamics.rigidformer.checkpoint .. autosummary:: load_rigidformer_checkpoint + export_rigidformer_checkpoint .. autofunction:: load_rigidformer_checkpoint + +.. autofunction:: export_rigidformer_checkpoint + +Prepared packed scenes +---------------------- + +.. currentmodule:: embodichain_models.dynamics.rigidformer.data + +.. autosummary:: + + MOViSequenceSource + +.. autoclass:: MOViSequenceSource + :members: + +Shifted supervision and anchor losses +------------------------------------- + +.. currentmodule:: embodichain_models.dynamics.rigidformer.supervised + +.. autosummary:: + + RigidFormerObjective + rigidformer_batch + +.. autoclass:: RigidFormerObjective + :members: + +.. autofunction:: rigidformer_batch + +Recoverable EMA state +--------------------- + +.. currentmodule:: embodichain_models.dynamics.rigidformer.ema + +.. autosummary:: + + RigidFormerEMA + +.. autoclass:: RigidFormerEMA + :members: + +Shared SL training recipe +------------------------- + +.. currentmodule:: embodichain_models.dynamics.rigidformer.sl_recipe + +.. autosummary:: + + build_rigidformer_trainer + +.. autofunction:: build_rigidformer_trainer + +Recovery and experiment boundaries +---------------------------------- + +The source is explicitly registered as a training-state participant so its +epoch/cursor commits after optimizer success. Prepared files are content-checked; +source identity also binds splits, scaling, sampling and seed. EMA shadows remain +FP32 and retain the author's warmup policy. Official weight import is distinct +from native optimizer recovery; export selects EMA or raw weights explicitly. + +Applications own acquisition and independent reference experiments. Short +numerical comparisons do not establish published accuracy. diff --git a/embodichain/learning/sl/__init__.py b/embodichain/learning/sl/__init__.py index 468b0a89ec..365f7095a8 100644 --- a/embodichain/learning/sl/__init__.py +++ b/embodichain/learning/sl/__init__.py @@ -21,7 +21,12 @@ from importlib import import_module from typing import TYPE_CHECKING, Any -from .types import LossResult, SupervisedBatch, SupervisedObjective +from .types import ( + LossResult, + SupervisedBatch, + SupervisedObjective, + SupervisedTrainingState, +) from .distributed import distributed_run from .learner import SupervisedLearner from .sources import SupervisedSource, TensorBatchSource @@ -34,6 +39,7 @@ "LossResult", "SupervisedBatch", "SupervisedObjective", + "SupervisedTrainingState", "SupervisedLearner", "SupervisedSource", "TensorBatchSource", diff --git a/embodichain/learning/sl/distributed.py b/embodichain/learning/sl/distributed.py index 07c28dc773..00d0f1ff5a 100644 --- a/embodichain/learning/sl/distributed.py +++ b/embodichain/learning/sl/distributed.py @@ -21,8 +21,10 @@ from collections.abc import Callable, Iterator from contextlib import contextmanager from datetime import timedelta +import io import math import os +import pickle from typing import Any, TypeVar import torch @@ -159,8 +161,53 @@ def agree(self, value: Any, stage: str) -> None: def gather(self, value: Any) -> list[Any] | None: if not self.enabled: return [value] - # Keep the small control-state path on all-gather collectives, which - # are supported by both NCCL and Gloo. Non-root copies are temporary. - values: list[Any] = [None] * self.world_size - dist.all_gather_object(values, value) - return values if self.rank == 0 else None + + # Participant snapshots can contain a full model EMA. Broadcasting one + # bounded chunk at a time keeps NCCL's device allocation independent of + # both the snapshot size and world size; only rank zero retains copies. + def serialize() -> io.BytesIO | None: + if self.rank == 0: + return None + stream = io.BytesIO() + pickle.dump(value, stream, protocol=pickle.HIGHEST_PROTOCOL) + return stream + + payload = self.run(serialize, "checkpoint state serialization") + sizes: list[int] = [0] * self.world_size + dist.all_gather_object(sizes, payload.tell() if payload is not None else 0) + chunk_size = min(max(sizes), 16 * 1024 * 1024) + buffer = self.run( + lambda: torch.empty( + chunk_size, dtype=torch.uint8, device=self.control_device + ), + "checkpoint transfer buffer", + ) + values: list[Any] | None = [value] if self.rank == 0 else None + sender = ( + torch.frombuffer(payload.getbuffer(), dtype=torch.uint8) + if payload is not None + else None + ) + for source in range(1, self.world_size): + received = io.BytesIO() if self.rank == 0 else None + failure = None + for offset in range(0, sizes[source], chunk_size): + count = min(chunk_size, sizes[source] - offset) + if self.rank == source: + assert sender is not None + buffer[:count].copy_(sender[offset : offset + count]) + dist.broadcast(buffer[:count], src=source) + if received is not None and failure is None: + try: + received.write(buffer[:count].cpu().numpy().tobytes()) + except Exception as error: + failure = error + if received is not None and failure is None: + try: + received.seek(0) + assert values is not None + values.append(pickle.load(received)) + except Exception as error: + failure = error + self.check(failure, "checkpoint state collection") + return values diff --git a/embodichain/learning/sl/learner.py b/embodichain/learning/sl/learner.py index 0c8c7e898c..31ff28f914 100644 --- a/embodichain/learning/sl/learner.py +++ b/embodichain/learning/sl/learner.py @@ -63,6 +63,9 @@ class SupervisedLearner: numerical_context: Context factory spanning forward, backward and clipping. ddp_find_unused_parameters: Enable graph traversal for objectives that omit trainable parameters. Applies only with an initialized group. + ddp_broadcast_buffers: Broadcast rank-zero buffers before forwards. + Defaults to the existing rank-local policy; explicit reference + recipes can request the PyTorch DDP broadcast policy. Sampling, normalizer updates, learning-rate schedules and model publication belong to callers. An update is committed only after optimizer.step succeeds. @@ -81,6 +84,7 @@ def __init__( ) = None, numerical_context: Callable[[], AbstractContextManager[Any]] | None = None, ddp_find_unused_parameters: bool = False, + ddp_broadcast_buffers: bool = False, ) -> None: if not math.isfinite(max_grad_norm) or max_grad_norm < 0: raise ValueError("max_grad_norm must be finite and nonnegative") @@ -92,6 +96,7 @@ def __init__( self.numerical_context = numerical_context or nullcontext self.num_updates = 0 self.ddp_find_unused_parameters = ddp_find_unused_parameters + self.ddp_broadcast_buffers = ddp_broadcast_buffers self._ddp: DistributedDataParallel | None = None self._distributed_metrics: tuple[str, ...] | None = None @@ -103,7 +108,7 @@ def _configure_distributed(self) -> _DistributedContext: device_ids=( [context.device.index] if context.device.type == "cuda" else None ), - broadcast_buffers=False, + broadcast_buffers=self.ddp_broadcast_buffers, find_unused_parameters=self.ddp_find_unused_parameters, ) return context diff --git a/embodichain/learning/sl/trainer.py b/embodichain/learning/sl/trainer.py index ede0a9f169..67764a79ef 100644 --- a/embodichain/learning/sl/trainer.py +++ b/embodichain/learning/sl/trainer.py @@ -31,6 +31,7 @@ from .distributed import _DistributedContext from .learner import SupervisedLearner from .sources import SupervisedSource +from .types import SupervisedTrainingState __all__ = ["SupervisedTrainerCfg", "SupervisedTrainer"] @@ -70,6 +71,8 @@ class SupervisedTrainer: signature: Recipe semantics, including feature/action units, frames, model configuration and numerical policy where applicable. scheduler: Optional scheduler stepped after each successful update. + training_states: Ordered update/recovery participants. States commit + after optimizer success and are checkpointed separately per rank. evaluate: Optional callback receiving the committed learner. It owns validation metrics and must not select a model on a test split. """ @@ -83,6 +86,7 @@ def __init__( signature: Mapping[str, Any] | None = None, scheduler: Any = None, evaluate: Callable[[SupervisedLearner], Mapping[str, float]] | None = None, + training_states: Mapping[str, SupervisedTrainingState] | None = None, ) -> None: self.cfg = cfg or SupervisedTrainerCfg() self.cfg.validate() @@ -101,6 +105,22 @@ def __init__( json.dumps(dict(signature or {}), sort_keys=True, allow_nan=False) ) self.scheduler, self.evaluate = scheduler, evaluate + self.training_states = dict(training_states or {}) + if any(not isinstance(name, str) or not name for name in self.training_states): + raise ValueError("Training state names must be nonempty strings") + for value in self.training_states.values(): + if not all( + callable(getattr(value, name, None)) + for name in ( + "on_update", + "state_dict", + "validate_state_dict", + "load_state_dict", + ) + ): + raise TypeError( + "Training states must implement the supervised update/recovery protocol" + ) self._context = _DistributedContext(self.cfg.device) self.sampler = torch.Generator(device="cpu").manual_seed( self.cfg.seed + self._context.rank @@ -140,6 +160,11 @@ def _prepare_distributed(self) -> None: "max_grad_norm": self.learner.max_grad_norm, "amp_dtype": str(self.learner.amp_dtype), "find_unused_parameters": self.learner.ddp_find_unused_parameters, + "broadcast_buffers": self.learner.ddp_broadcast_buffers, + "training_state_types": [ + (name, f"{type(value).__module__}.{type(value).__qualname__}") + for name, value in self.training_states.items() + ], "scheduler_type": ( None if self.scheduler is None @@ -196,6 +221,11 @@ def admit_output() -> None: ) metrics = self.learner.update(batch) self.sampler.set_state(candidate.get_state()) + for name, component in self.training_states.items(): + self._context.run( + lambda component=component: component.on_update(self.learner), + f"training state update: {name}", + ) if self.scheduler is not None: self._context.run(self.scheduler.step, "scheduler") step = self.learner.num_updates @@ -228,6 +258,7 @@ def append_metrics() -> None: def _local_rank_state(self) -> dict[str, Any]: return { + "training_states": self._training_state_dict(), "sampler": self.sampler.get_state(), "rng": self._rng_state(), "buffers": { @@ -236,6 +267,30 @@ def _local_rank_state(self) -> dict[str, Any]: }, } + def _training_state_dict(self) -> dict[str, Any]: + return { + name: { + "type": f"{type(value).__module__}.{type(value).__qualname__}", + "state": value.state_dict(), + } + for name, value in self.training_states.items() + } + + def _validate_training_states(self, saved: Mapping[str, Any]) -> None: + if tuple(saved) != tuple(self.training_states): + raise ValueError("Supervised checkpoint training state names differ") + for name, value in self.training_states.items(): + identity = f"{type(value).__module__}.{type(value).__qualname__}" + if saved[name].get("type") != identity: + raise ValueError( + f"Supervised checkpoint training state type differs: {name}" + ) + value.validate_state_dict(saved[name]["state"]) + + def _restore_training_states(self, saved: Mapping[str, Any]) -> None: + for name, value in self.training_states.items(): + value.load_state_dict(saved[name]["state"]) + def _gather_rank_states(self) -> list[Any] | None: state = self._context.run(self._local_rank_state, "checkpoint rank state") return self._context.gather(state) @@ -268,6 +323,7 @@ def evaluate() -> dict[str, float]: for name, value in self.learner.model.named_buffers(): value.copy_(buffers[name]) self.sampler.set_state(snapshot[0]["sampler"]) + self._restore_training_states(snapshot[0]["training_states"]) self._restore_rng(rng) self.learner.model.train(was_training) @@ -315,13 +371,19 @@ def _write_checkpoint( self._artifacts.claim() path.parent.mkdir(parents=True, exist_ok=True) state = { - "version": 2 if self._context.enabled else 1, + "version": self._checkpoint_version(), "run_id": self._artifacts.run_id, "learner": self.learner.state_dict(), "sampler": self.sampler.get_state(), "rng": self._rng_state(), "signature": self.signature, "source_id": self.source.source_id, + "ddp_broadcast_buffers": self.learner.ddp_broadcast_buffers, + "training_states": ( + self._training_state_dict() + if rank_states is None + else rank_states[0]["training_states"] + ), "config": self.cfg.to_dict(), "scheduler_type": ( None @@ -390,11 +452,19 @@ def restore() -> None: else: local = state self.sampler.set_state(local["sampler"]) + self._restore_training_states(local.get("training_states", {})) self._restore_rng(local["rng"]) self._context.run(restore, "checkpoint restore") self._artifacts.adopt_resume(run_id) + def _checkpoint_version(self) -> int: + # Older trainers must not silently discard required participant states + # or a nondefault DDP buffer policy. Default recipes keep formats 1/2. + return (2 if self._context.enabled else 1) + ( + 2 if self.training_states or self.learner.ddp_broadcast_buffers else 0 + ) + def _validate_checkpoint(self, state: Mapping[str, Any]) -> None: # This preflight runs collectively before output admission, DDP setup, # or model/optimizer/RNG restoration. Mutable delivery state may not be @@ -402,13 +472,20 @@ def _validate_checkpoint(self, state: Mapping[str, Any]) -> None: validate_source = getattr(self.source, "validate_resume", None) if validate_source is not None: validate_source() + if ( + state.get("ddp_broadcast_buffers", False) + != self.learner.ddp_broadcast_buffers + ): + raise ValueError( + "Supervised checkpoint DDP buffer-broadcast policy differs" + ) distributed = state.get("distributed", {}) if distributed.get("world_size", 1) != self._context.world_size: raise ValueError( "Supervised checkpoint world size differs; exact recovery requires the same world size" ) if ( - state.get("version") != (2 if self._context.enabled else 1) + state.get("version") != self._checkpoint_version() or state["signature"] != self.signature or state["source_id"] != self.source.source_id ): @@ -430,6 +507,12 @@ def _validate_checkpoint(self, state: Mapping[str, Any]) -> None: if state.get("scheduler_type") != scheduler_type: raise ValueError("Supervised checkpoint scheduler type differs") self.learner._validate_state_dict(state["learner"]) + local_state = ( + distributed["rank_states"][self._context.rank] + if self._context.enabled + else state + ) + self._validate_training_states(local_state.get("training_states", {})) if self._context.enabled: rank_states = distributed["rank_states"] if len(rank_states) != self._context.world_size: diff --git a/embodichain/learning/sl/types.py b/embodichain/learning/sl/types.py index 00e58ebf25..1f451cd7bd 100644 --- a/embodichain/learning/sl/types.py +++ b/embodichain/learning/sl/types.py @@ -20,12 +20,20 @@ from collections.abc import Mapping from dataclasses import dataclass, field -from typing import Any, Protocol +from typing import TYPE_CHECKING, Any, Protocol import torch from torch import nn -__all__ = ["SupervisedBatch", "LossResult", "SupervisedObjective"] +if TYPE_CHECKING: + from .learner import SupervisedLearner + +__all__ = [ + "SupervisedBatch", + "LossResult", + "SupervisedObjective", + "SupervisedTrainingState", +] def _map_tensors(value: Any, function: Any) -> Any: @@ -95,6 +103,42 @@ class LossResult: weight: float = 1.0 +class SupervisedTrainingState(Protocol): + """Explicit update participant with owned rank-local checkpoint state. + + Participants commit in insertion order after optimizer success and before + scheduling/evaluation. Validation must not mutate live state. + """ + + def on_update(self, learner: SupervisedLearner) -> None: + """Commit state after the learner's successful optimizer update. + + Args: + learner: Learner at its new committed update boundary. + """ + ... + + def state_dict(self) -> Mapping[str, Any]: + """Return an owned serializable snapshot of the committed state.""" + ... + + def validate_state_dict(self, state: Mapping[str, Any]) -> None: + """Validate compatibility without changing live state. + + Args: + state: Candidate checkpoint state. + """ + ... + + def load_state_dict(self, state: Mapping[str, Any]) -> None: + """Restore a previously validated snapshot. + + Args: + state: State accepted by validate_state_dict. + """ + ... + + class SupervisedObjective(Protocol): """Compute task-specific losses without owning optimizer or data lifecycle.""" diff --git a/extensions/embodichain_models/README.md b/extensions/embodichain_models/README.md index bc3980232c..430238a1f9 100644 --- a/extensions/embodichain_models/README.md +++ b/extensions/embodichain_models/README.md @@ -35,7 +35,7 @@ Use the following canonical packages for the four current model directions: | Ordinary state/visual BC | `embodichain_models.policy.bc` | Target for the BC follow-up; not created yet. | | NeRD | `embodichain_models.dynamics.nerd` | Implemented in this change. | | Robot world model (RWM) | `embodichain_models.dynamics.rwm` | Implemented. | -| RigidFormer | `embodichain_models.dynamics.rigidformer` | Native author architecture and strict official weights. | +| RigidFormer | `embodichain_models.dynamics.rigidformer` | Native author architecture, official weights and shared-SL recipe. | RWM predicts action-conditioned robot state transitions, so it belongs to `dynamics`. World models also describe future-state/observation prediction; diff --git a/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/__init__.py b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/__init__.py index d74413c98d..47a68fea75 100644 --- a/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/__init__.py +++ b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/__init__.py @@ -25,12 +25,22 @@ "RigidFormerCfg", "RigidFormerModel", "load_rigidformer_checkpoint", + "export_rigidformer_checkpoint", + "MOViSequenceSource", + "RigidFormerObjective", + "rigidformer_batch", + "build_rigidformer_trainer", ] _EXPORTS = { "RIGIDFORMER_REVISION": "model", "RigidFormerCfg": "model", "RigidFormerModel": "model", "load_rigidformer_checkpoint": "checkpoint", + "export_rigidformer_checkpoint": "checkpoint", + "MOViSequenceSource": "data", + "RigidFormerObjective": "supervised", + "rigidformer_batch": "supervised", + "build_rigidformer_trainer": "sl_recipe", } diff --git a/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/data/__init__.py b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/data/__init__.py new file mode 100644 index 0000000000..260086a759 --- /dev/null +++ b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/data/__init__.py @@ -0,0 +1,7 @@ +# Copyright (c) 2026 Zhiyang Dou. SPDX-License-Identifier: MIT +# Derived from RigidFormer at 160fe0852c51cffae7b7440452b58d0b49db6df6. +# See the model family's THIRD_PARTY_LICENSES for the complete license. + +"""Private author geometry storage adapters.""" + +from __future__ import annotations diff --git a/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/data/_index.py b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/data/_index.py new file mode 100644 index 0000000000..2f30e4011b --- /dev/null +++ b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/data/_index.py @@ -0,0 +1,67 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Compact author window indexing without materializing repeated paths.""" + +from __future__ import annotations +from bisect import bisect_right +from pathlib import Path +from typing import Any + + +class _SceneWindows: + """Store scene runs while preserving the legacy window-index mapping.""" + + def __init__(self, *, random_sampling: bool) -> None: + self.random_sampling = random_sampling + self.runs: list[tuple[str, tuple[Path, ...] | None, int, int]] = [] + self.ends: list[int] = [] + + def add_scene( + self, name: str, frames: list[Path], sequence_length: int, stride: int + ) -> None: + count = len(range(0, len(frames) - sequence_length + 1, stride)) + if not count: + return + self.runs.append( + ( + name, + None if self.random_sampling else tuple(frames), + sequence_length, + stride, + ) + ) + self.ends.append((self.ends[-1] if self.ends else 0) + count) + + def __len__(self) -> int: + return self.ends[-1] if self.ends else 0 + + def __getitem__(self, index: int) -> dict[str, Any]: + if index < 0: + index += len(self) + if not 0 <= index < len(self): + raise IndexError(index) + run = bisect_right(self.ends, index) + name, frames, length, stride = self.runs[run] + if self.random_sampling: + return {"dataset": name} + start = (index - (self.ends[run - 1] if run else 0)) * stride + selected = frames[start : start + length] + return { + "dataset": name, + "frames": [p.as_posix() for p in selected], + "frame_ids": [int(p.name.split("_")[-1]) for p in selected], + } diff --git a/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/data/dataset.py b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/data/dataset.py new file mode 100644 index 0000000000..209eeefbee --- /dev/null +++ b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/data/dataset.py @@ -0,0 +1,1365 @@ +# Copyright (c) 2026 Zhiyang Dou. SPDX-License-Identifier: MIT +# Derived from RigidFormer at 160fe0852c51cffae7b7440452b58d0b49db6df6. +# See the model family's THIRD_PARTY_LICENSES for the complete license. + +""" +Multi-dataset data loader for multi-body dynamics. +Loads sequences from multiple dataset directories, with train/val/test +membership defined by a split file. +""" + +from __future__ import annotations + +import os +import json +import time +import concurrent.futures as futures +import h5py +import numpy as np +import torch +from torch.utils.data import Dataset, DataLoader, Sampler, RandomSampler +from torch.utils.data.distributed import DistributedSampler +from typing import Dict, List, Tuple, Optional +import glob +from pathlib import Path +import random +from tqdm import tqdm +import sys +from functools import lru_cache +import logging + + +from .object_scale import ObjectScale, check_object_counts + +from .scene_frames import frame_dirs as scene_frame_dirs, is_packed, load_frame +from ._index import _SceneWindows +import zlib + +logger = logging.getLogger(__name__) + +# Padding value for positions: zero contributes nothing to +# pooling/aggregation operations. +PADDING_VALUE = 0.0 + + +# Global LRU cache for frequently loaded numpy arrays +# (~50000 entries at ~5KB each, roughly 250MB). +@lru_cache(maxsize=256) +def _cached_numpy_load(filepath: str) -> np.ndarray: + """Cached numpy array loading. Uses filepath as cache key.""" + return np.load(filepath) + + +def cached_load_npy(path: Path) -> np.ndarray: + """Load numpy array with LRU caching.""" + return _cached_numpy_load(str(path)) + + +def check_scene_set(config: dict) -> str: + """data.scene_set must name the scene set the split file declares ('benchmark' when either is absent). + + 'complete' splits (data_generation/build_complete_splits.py) leave out the MOVi scenes whose metadata omits simulated + bodies; 'benchmark' is the published scene list. Returns the scene set. + """ + want = (config.get("data", {}) or {}).get("scene_set", "benchmark") + if want not in ("benchmark", "complete"): + raise ValueError( + f"data.scene_set must be 'benchmark' or 'complete'; got {want!r}" + ) + with open(config["data"]["split_file"]) as f: + have = json.load(f).get("scene_set", "benchmark") + if have != want: + raise ValueError( + f"data.scene_set is {want!r} but {config['data']['split_file']} declares scene_set {have!r}" + ) + return want + + +class MultiObjectDynamics(Dataset): + """MOVi multi-object dataset with per-object vertices and per-frame padding+mask. + + Expects dataset directories like movi_* with frames frame_00000/... and files: + - meshes/obj_i_faces.npy (per-object faces, static per dataset) + - frame_xxxxx/obj_i_pos.npy (per-object vertices in world space) + """ + + def __init__( + self, + data_path: str, + sequence_length: int = 3, + stride: int = 1, + split: str = "train", + train_ratio: float = 0.9, + val_ratio: float = 0.1, + split_file: Optional[str] = None, + fast_index: bool = True, + skip_first_frames: int = 0, + skip_step_size: int = 1, + max_vertices_per_object: Optional[int] = None, + max_objects_per_frame: Optional[int] = None, + anchor_k: Optional[int] = None, + if_phys_param: bool = False, + **kwargs, + ): + self.data_path = Path(data_path) + self.sequence_length = sequence_length + self.stride = stride + self.split = split + self.skip_first_frames = skip_first_frames + self.skip_step_size = skip_step_size + self.max_vertices_per_object = max_vertices_per_object + self.max_objects_per_frame = max_objects_per_frame + self.if_phys_param = if_phys_param # Strict mode for physics parameters + + # Point Cloud Mode configuration + # When enabled, load point cloud files (obj_i_pc{N}_pos.npy) instead of mesh vertices + if_point_cloud_config = kwargs.get("if_point_cloud", {}) + self.point_cloud_enable = if_point_cloud_config.get("enable", False) + if self.point_cloud_enable: + # the point-cloud path validates neither per-frame point counts nor per-resolution padding (a frame with a + # missing point was accepted; resolutions above the first were truncated to it), and train.py does not + # forward the setting: refused until implemented and validated + raise NotImplementedError( + "point-cloud inputs (if_point_cloud.enable) are not supported by this loader" + ) + self.point_cloud_train_resolution = if_point_cloud_config.get( + "train_resolution", [512] + ) + self.point_cloud_eval_resolution = if_point_cloud_config.get( + "eval_resolution", [512] + ) + + # Ensure resolutions are lists + if isinstance(self.point_cloud_train_resolution, int): + self.point_cloud_train_resolution = [self.point_cloud_train_resolution] + if isinstance(self.point_cloud_eval_resolution, int): + self.point_cloud_eval_resolution = [self.point_cloud_eval_resolution] + + # Select resolution based on split + if self.point_cloud_enable: + if split == "train": + self.point_cloud_resolutions = self.point_cloud_train_resolution + else: + self.point_cloud_resolutions = self.point_cloud_eval_resolution + print(f"[{split}] Point Cloud Mode ENABLED:") + print(f" - Resolutions: {self.point_cloud_resolutions}") + print(f" - Data will be loaded from: frame_XXXXX/obj_i_pc{{N}}_pos.npy") + else: + self.point_cloud_resolutions = None + + # Anchor config (Nq): prefer explicit arg, else env ANCHOR_K, else default 32 + if anchor_k is None: + env_anchor = os.environ.get("ANCHOR_K", "").strip() + if env_anchor: + try: + anchor_k = int(env_anchor) + except Exception: + raise ValueError( + f"Invalid ANCHOR_K environment value: '{env_anchor}' (expected int)" + ) + else: + anchor_k = 32 + if anchor_k <= 0: + raise ValueError(f"anchor_k must be positive; got {anchor_k}") + self.anchor_k = int(anchor_k) + + # Object Permutation Augmentation (train only, at least 2 objects) + augmentation_config = kwargs.get("augmentation", {}) + # The random-walk / independent noise blocks were never applied on this path (the dataset never stored them and + # the helpers expect body1_/body2_ arrays, not MOVi's obj_positions): enabling them is an error, not a no-op. + for _blk in ("random_walk", "independent"): + for _cfg in (augmentation_config, kwargs.get("noise_config") or {}): + if ((_cfg or {}).get(_blk) or {}).get("enable", False): + raise NotImplementedError( + f"data.augmentation.{_blk}.enable: true is not implemented for the MOVi loader " + f"(it would be silently ignored); set it to false" + ) + # Anchor selection rule. 'index0' starts FPS at vertex 0 (order dependent); + # 'centroid_farthest' starts at the vertex farthest from the centroid with lexicographic + # tie-breaks, so the anchors depend on the vertex set only. + self.anchor_fps_start = str(kwargs.get("anchor_fps_start", "index0")) + # Order-independent FPS ordering of every object's vertices (cached as obj_i_fpsorder_cf.npy); + # the encoder's hierarchical levels take the first 50/25/12.5/... percent of this ordering. + self.fps_order = bool(kwargs.get("fps_order", False)) + # data.object_scale_file: apply each object's metadata size to the stored vertices (src/data/object_scale.py) + _scale_file = kwargs.get("object_scale_file") + self._object_scale = ObjectScale(_scale_file) if _scale_file else None + if self._object_scale is not None: + # the folder name is compared without regard to case + _ds_real = os.path.basename(os.path.realpath(str(data_path))) + if (self._object_scale.converted or "").lower() != _ds_real.lower(): + raise ValueError( + f"data.object_scale_file {_scale_file} was built for '{self._object_scale.converted}', " + f"but the dataset path resolves to '{_ds_real}'" + ) + if self._object_scale is not None: + print( + f"[{split}] Object scale ACTIVE: vertices scaled about the object origin by the metadata size ({_scale_file})" + ) + # data.partial_view: partial point clouds of a simulated depth camera instead of the mesh vertices + # (src/data/partial_view.py) + _pv = kwargs.get("partial_view") or {} + self._partial = None + if _pv.get("enable", False): + if self._object_scale is None: + raise ValueError( + "data.partial_view needs data.object_scale_file for the objects' metadata sizes" + ) + self._partial = PartialView(_pv, Path(__file__).resolve().parents[2]) + if self._partial.min_points < self.anchor_k: + raise ValueError( + f"data.partial_view.min_points ({self._partial.min_points}) must be at least " + f"anchor_k ({self.anchor_k})" + ) + self._partial_scenes = {} + print(f"[{split}] {self._partial.describe(train=split == 'train')}") + self._fps_order_by_dataset = {} + if self.anchor_fps_start not in ("index0", "centroid_farthest"): + raise ValueError( + f"data.anchor_fps_start must be 'index0' or 'centroid_farthest'; got '{self.anchor_fps_start}'" + ) + if self.anchor_fps_start != "index0": + print( + f"[{split}] Anchor FPS start rule: {self.anchor_fps_start} (anchor cache suffix '_cf')" + ) + # Anchor resampling augmentation (train split only): with probability `prob` the FPS anchors + # of every object are recomputed from a random start vertex. + resample_config = augmentation_config.get("anchor_resampling", {}) + self.anchor_resample_prob = ( + float(resample_config.get("prob", 0.5)) + if resample_config.get("enable", False) + else 0.0 + ) + if self.anchor_resample_prob > 0 and split == "train": + print( + f"[{split}] Anchor resampling augmentation enabled: prob={self.anchor_resample_prob:.2f}" + ) + permute_config = augmentation_config.get("object_permutation", {}) + self.permute_enabled = permute_config.get("enable", False) and split == "train" + self.permute_prob = permute_config.get("prob_apply", 0.5) + self.permute_debug = permute_config.get("if_debug", False) + if self.permute_enabled: + print( + f"[{split}] Object permutation augmentation enabled: prob={self.permute_prob:.2f}, debug={self.permute_debug}" + ) + + # Load the dataset split file (required; splits are never auto-generated). + if not split_file: + raise ValueError( + "split_file is required. Set 'data.split_file' in the config to a split JSON, " + "e.g. one of the dataset_split_*.json files in the repository root." + ) + + if not os.path.exists(split_file): + raise FileNotFoundError( + f"Split file not found: {split_file}\n" + "Check the 'data.split_file' path in the config. Split files " + "(dataset_split_*.json) are provided in the repository root; new splits " + "can be written for data produced by data_generation/movi_to_sim_data.py." + ) + + with open(split_file, "r") as f: + s = json.load(f) + names = ( + s["train_datasets"] + if split == "train" + else s["val_datasets"] if split == "val" else s["test_datasets"] + ) + self.dataset_dirs = [self.data_path / n for n in names] + if self._object_scale is not None: + check_object_counts(self._object_scale, self.data_path, names) + + # Build sequences + self.all_sequences = ( + [] + ) # each: { 'dataset': name, 'frames': [paths], 'frame_ids': [int] } + if kwargs.get("compact_index", False): + self.all_sequences = _SceneWindows( + random_sampling=bool( + (kwargs.get("random_step_size") or {}).get("enable", False) + ) + ) + + # DDP-aware: only show progress bar on rank 0 + import torch.distributed as dist + + show_progress = True + if dist.is_available() and dist.is_initialized(): + show_progress = dist.get_rank() == 0 + + for ds in tqdm( + self.dataset_dirs, + desc=f"Indexing {split} datasets", + unit="ds", + disable=not show_progress, + ): + frames = scene_frame_dirs(ds) + if self.skip_first_frames > 0: + frames = frames[self.skip_first_frames :] + if self.skip_step_size > 1: + frames = frames[:: self.skip_step_size] + if isinstance(self.all_sequences, _SceneWindows): + self.all_sequences.add_scene( + ds.name, frames, sequence_length, self.stride + ) + continue + for start in range(0, len(frames) - sequence_length + 1, self.stride): + seq = frames[start : start + sequence_length] + if len(seq) == sequence_length: + self.all_sequences.append( + { + "dataset": ds.name, + "frames": [p.as_posix() for p in seq], + "frame_ids": [ + int(Path(p).name.split("_")[-1]) for p in seq + ], + } + ) + + # Precompute/load per-object anchor indices and faces for each dataset from + # the reference frame_00000; cached in memory. + self._anchors_by_dataset = {} + self._faces_by_dataset = {} + + # For point cloud mode, use the first resolution for __init__ validation + init_pc_resolution = ( + self.point_cloud_resolutions[0] if self.point_cloud_enable else None + ) + + for ds in self.dataset_dirs: + ds_name = ds.name + frame0 = ds / "frame_00000" + if not (frame0.exists() or is_packed(ds)): + raise FileNotFoundError( + f"Missing reference frame for dataset '{ds_name}': {frame0}. " + f"Expected a 'frame_00000' directory (or frames.npy) under {ds}." + ) + + # For point cloud mode, load point clouds for validation + # For mesh mode, load mesh vertices + ref_objs = self._load_frame_objects( + frame0, pc_resolution=init_pc_resolution + ) + if len(ref_objs) == 0: + raise RuntimeError( + f"Empty reference frame in {frame0}: no per-object vertex files found." + ) + anchors_dir = ds / "anchors" + require_cache = kwargs.get("require_prepared_cache", False) + if not require_cache: + anchors_dir.mkdir(parents=True, exist_ok=True) + ds_anchor_list = [] + ds_faces_list = [] + ds_order_list = [] + for i, V in enumerate(ref_objs): + nv = V.shape[0] + if ( + self.max_vertices_per_object is not None + and nv > self.max_vertices_per_object + ): + raise ValueError( + f"Dataset '{ds_name}' reference object {i} has {nv} vertices, " + f"exceeding max_vertices_per_object={self.max_vertices_per_object}.\n" + f"Either reduce vertex count in dataset or increase data.max_vertices_per_object in config." + ) + + # Point cloud mode: anchors are unused, so use sequential placeholder + # indices rather than mesh anchors (whose indices may be out of range). + if self.point_cloud_enable: + idx = np.arange(min(self.anchor_k, nv), dtype=np.int32) + if len(idx) < self.anchor_k: + idx = np.pad( + idx, (0, self.anchor_k - len(idx)), mode="constant" + ) + ds_anchor_list.append(idx) + # Create empty faces (not used in point cloud mode) + ds_faces_list.append(np.zeros((0, 3), dtype=np.int64)) + continue + + # Mesh vertex mode: standard anchor and faces loading. + if nv < self.anchor_k: + raise ValueError( + f"Dataset '{ds_name}' object {i} has only {nv} vertices (< anchor_k={self.anchor_k}). " + f"Reduce ANCHOR_K or increase vertex cap (max_vertices_per_object)." + ) + _suffix = "" if self.anchor_fps_start == "index0" else "_cf" + anchor_path = ( + anchors_dir / f"obj_{i}_anchors_fps{self.anchor_k}{_suffix}.npy" + ) + if anchor_path.exists(): + idx = np.load(anchor_path) + if idx.ndim != 1 or idx.shape[0] != self.anchor_k: + raise ValueError( + f"Invalid anchor file shape: {anchor_path} (expected ({self.anchor_k},))" + ) + if idx.dtype.kind not in ("i", "u"): + raise ValueError( + f"Invalid anchor file dtype: {anchor_path} (expected integer)" + ) + if int(idx.max()) >= nv or int(idx.min()) < 0: + raise ValueError( + f"Anchor indices in {anchor_path} out of range [0,{nv-1}] for object {i}" + ) + if np.unique(idx).shape[0] != idx.shape[0]: + raise ValueError( + f"Anchor cache {anchor_path} repeats a vertex: {idx.tolist()}" + ) + ds_anchor_list.append(idx.astype(np.int32, copy=False)) + else: + if require_cache: + raise ValueError( + f"Prepared anchor cache is missing: {anchor_path}" + ) + # Compute deterministic FPS anchors on the STORED (unscaled) frame-0 vertices, as + # pregen_cf_anchors.py does: FPS is invariant to the size scaling only in exact arithmetic (float32 + # rounding of the scaled coordinates breaks the exact ties of symmetric meshes differently) + V_ref = load_frame(ds, 0)[i].astype(np.float32, copy=False) + idx = self._farthest_point_sample_np( + V_ref, self.anchor_k, rule=self.anchor_fps_start + ) + # Atomic write: concurrent DDP ranks must never read a partially written cache. + tmp_path = anchor_path.with_name( + f"{anchor_path.stem}.{os.getpid()}.tmp.npy" + ) + np.save(tmp_path, idx.astype(np.int32, copy=False)) + os.replace(tmp_path, anchor_path) + ds_anchor_list.append(idx.astype(np.int32, copy=False)) + if self.fps_order: + order_path = anchors_dir / f"obj_{i}_fpsorder_cf.npy" + if order_path.exists(): + order = np.load(order_path) + if ( + order.ndim != 1 + or order.shape[0] != nv + or order.dtype.kind not in ("i", "u") + ): + raise ValueError( + f"Invalid FPS-order cache {order_path}: shape {order.shape}, expected ({nv},)" + ) + if not np.array_equal(np.sort(order), np.arange(nv)): + raise ValueError( + f"FPS-order cache {order_path} is not a permutation of the {nv} vertices" + ) + else: + if require_cache: + raise ValueError( + f"Prepared FPS-order cache is missing: {order_path}" + ) + # stored (unscaled) frame-0 vertices, as pregen_fps_order.py does (see the anchors above) + order = self._farthest_point_sample_np( + load_frame(ds, 0)[i].astype(np.float32, copy=False), + nv, + rule="centroid_farthest", + ) + tmp_path = order_path.with_name( + f"{order_path.stem}.{os.getpid()}.tmp.npy" + ) + np.save(tmp_path, order.astype(np.int32, copy=False)) + os.replace(tmp_path, order_path) + ds_order_list.append(order.astype(np.int32, copy=False)) + + # Load faces for this object (static per dataset), and filter to nv cap + faces_path = ds / "meshes" / f"obj_{i}_faces.npy" + if not faces_path.exists(): + raise FileNotFoundError( + f"Missing faces file for dataset '{ds_name}' object {i}: {faces_path}" + ) + faces = np.load(faces_path) + if faces.ndim != 2 or faces.shape[1] != 3: + raise ValueError( + f"Invalid faces array at {faces_path}: shape={faces.shape}, expected (F,3)" + ) + # every face must index this object's vertices (the vertex count is never capped: an object above + # max_vertices_per_object is an error), so an out-of-range face is corrupt data, not a face to drop + if faces.dtype.kind not in ("i", "u"): + raise ValueError( + f"Faces of {faces_path} have dtype {faces.dtype}, expected integers" + ) + faces = faces.astype(np.int64, copy=False) + if faces.size and (int(faces.min()) < 0 or int(faces.max()) >= nv): + raise ValueError( + f"Faces of {faces_path} index vertices outside [0, {nv - 1}]" + ) + ds_faces_list.append(faces) + self._anchors_by_dataset[ds_name] = ds_anchor_list + self._faces_by_dataset[ds_name] = ds_faces_list + if self.fps_order: + self._fps_order_by_dataset[ds_name] = ds_order_list + + # Cache reference frames (frame_00000) per dataset to avoid repeated loading. + self._reference_frames_cache = {} + + # Per-dataset caches: frame lists (avoids glob in __getitem__), frame IDs + # (avoids string parsing per sample), and per-object/max vertex counts. + self._all_frames_by_dataset = {} + self._frame_id_cache = {} + self._max_nv_by_dataset = {} + self._mesh_nv_by_dataset = {} + + for ds in self.dataset_dirs: + ds_name = ds.name + all_frames = scene_frame_dirs(ds) + self._all_frames_by_dataset[ds_name] = all_frames + for fp in all_frames: + frame_id = int(fp.name.split("_")[-1]) # "frame_00123" -> 123 + self._frame_id_cache[(ds_name, fp.name)] = frame_id + # Frames are indexed by position: a missing or extra frame directory would silently change the time step + # and misalign the ground truth (vis_traj.npy, which the evaluator indexes by frame number). + _ids = [self._frame_id_cache[(ds_name, fp.name)] for fp in all_frames] + _T_traj = np.load(ds / "vis_traj.npy", mmap_mode="r").shape[1] + if _ids != list(range(len(_ids))) or len(_ids) != _T_traj: + raise ValueError( + f"{ds}: frame directories must be frame_00000 .. frame_{_T_traj - 1:05d} without gaps " + f"(vis_traj.npy has {_T_traj} frames); found {len(_ids)} directories, " + f"first gap at {next((i for i, f in enumerate(_ids) if f != i), len(_ids))}" + ) + # Load frame_00000 once to record per-object vertex counts + # (at point cloud resolution when enabled). + frame0 = ds / "frame_00000" + ref_objs_data = self._load_frame_objects( + frame0, pc_resolution=init_pc_resolution + ) + nv_list = [obj.shape[0] for obj in ref_objs_data] + self._mesh_nv_by_dataset[ds_name] = nv_list + self._max_nv_by_dataset[ds_name] = max(nv_list) if nv_list else 0 + + # Load physics parameters (mass, friction, restitution) per dataset. + # With if_phys_param=True the file and every object entry must exist. + self._physics_by_dataset = {} + for ds in self.dataset_dirs: + ds_name = ds.name + physics_path = ds / "physics_params.json" + num_objs = len(self._anchors_by_dataset[ds_name]) + + if self.if_phys_param: + if not physics_path.exists(): + raise FileNotFoundError( + f"Physics parameters required but file missing: {physics_path}\n" + f"Dataset: {ds_name} | if_phys_param=True requires physics_params.json\n" + f"Please generate physics parameters or set if_phys_param=False" + ) + + # Parse errors propagate deliberately. + with open(physics_path, "r") as f: + physics_params = json.load(f) + + phys_list = [] + for i in range(num_objs): + obj_key = f"obj_{i}" + if obj_key not in physics_params: + raise KeyError( + f"Missing physics parameters for {obj_key} in {physics_path}\n" + f"Dataset: {ds_name} | Found objects: {list(physics_params.keys())}\n" + f"Required: obj_0 through obj_{num_objs-1}" + ) + + p = physics_params[obj_key] + if "mass" not in p: + raise KeyError( + f"Missing 'mass' field for {obj_key} in {physics_path}" + ) + if "friction" not in p: + raise KeyError( + f"Missing 'friction' field for {obj_key} in {physics_path}" + ) + if "restitution" not in p: + raise KeyError( + f"Missing 'restitution' field for {obj_key} in {physics_path}" + ) + + phys_list.append( + [ + float(p["mass"]), + float(p["friction"]), + float(p["restitution"]), + ] + ) + + self._physics_by_dataset[ds_name] = np.array( + phys_list, dtype=np.float32 + ) # (Nobj, 3) + else: + # Permissive mode: substitute defaults for missing entries. + if physics_path.exists(): + with open(physics_path, "r") as f: + physics_params = json.load(f) + phys_list = [] + for i in range(num_objs): + obj_key = f"obj_{i}" + if obj_key in physics_params: + p = physics_params[obj_key] + phys_list.append( + [ + float(p.get("mass", 1.0)), + float(p.get("friction", 0.5)), + float(p.get("restitution", 0.5)), + ] + ) + else: + phys_list.append([1.0, 0.5, 0.5]) + self._physics_by_dataset[ds_name] = np.array( + phys_list, dtype=np.float32 + ) + else: + # File missing entirely: default values for all objects. + self._physics_by_dataset[ds_name] = np.array( + [[1.0, 0.5, 0.5]] * num_objs, dtype=np.float32 + ) + + # Multi-timescale training configuration; step sizes are expressed in + # normalized time units. + random_step_size_config = kwargs.get("random_step_size", {}) + self.random_step_size_enable = random_step_size_config.get("enable", False) + self.random_step_sizes = random_step_size_config.get( + "step_sizes", [1, 2, 5, 10] + ) + self.random_start_frame = random_step_size_config.get( + "random_start_frame", True + ) + self.sampling_weights = random_step_size_config.get("sampling_weights", None) + + if self.random_step_size_enable: + if ( + not isinstance(self.random_step_sizes, list) + or len(self.random_step_sizes) == 0 + ): + raise ValueError( + f"random_step_size.step_sizes must be a non-empty list, got: {self.random_step_sizes}" + ) + + for ss in self.random_step_sizes: + if not isinstance(ss, int) or ss <= 0: + raise ValueError( + f"All step_sizes must be positive integers, got: {self.random_step_sizes}" + ) + + if self.sampling_weights is None: + # Default to uniform sampling. + self.sampling_weights = [1.0 / len(self.random_step_sizes)] * len( + self.random_step_sizes + ) + else: + if len(self.sampling_weights) != len(self.random_step_sizes): + raise ValueError( + f"sampling_weights length ({len(self.sampling_weights)}) must match " + f"step_sizes length ({len(self.random_step_sizes)})" + ) + + weights_sum = sum(self.sampling_weights) + if abs(weights_sum - 1.0) > 1e-6: + raise ValueError( + f"sampling_weights must sum to 1.0, got sum={weights_sum:.6f}\n" + f"Weights: {self.sampling_weights}" + ) + + print(f"\n{'='*60}") + print(f"[{split}] Multi-timescale training ENABLED:") + print(f" - step_sizes: {self.random_step_sizes}") + print(f" - sampling_weights: {self.sampling_weights}") + print(f" - random_start_frame: {self.random_start_frame}") + print(f" - if_debug: {random_step_size_config.get('if_debug', False)}") + print(f"{'='*60}\n") + + # Every dataset must have enough frames for every configured step size. + print( + f"[{split}] Validating dataset frame counts for multi-timescale training..." + ) + insufficient_datasets = [] + for ds in self.dataset_dirs: + frames = scene_frame_dirs(ds) + if self.skip_first_frames > 0: + frames = frames[self.skip_first_frames :] + available_frames = len(frames) + + for step_size in self.random_step_sizes: + frames_needed = (self.sequence_length - 1) * step_size + 1 + if available_frames < frames_needed: + insufficient_datasets.append( + { + "dataset": ds.name, + "available": available_frames, + "needed": frames_needed, + "step_size": step_size, + "sequence_length": self.sequence_length, + } + ) + + if insufficient_datasets: + error_msg = f"[{split}] Frame validation failed: {len(insufficient_datasets)} datasets have insufficient frames:\n" + for info in insufficient_datasets[:5]: + error_msg += ( + f" - {info['dataset']}: available={info['available']}, " + f"needed={info['needed']} (step_size={info['step_size']}, " + f"sequence_length={info['sequence_length']})\n" + ) + if len(insufficient_datasets) > 5: + error_msg += ( + f" ... and {len(insufficient_datasets) - 5} more datasets\n" + ) + error_msg += ( + f"\nSolutions:\n" + f" 1. Reduce max step_size in config: data.random_step_size.step_sizes\n" + f" 2. Reduce sequence_length in config: data.sequence_length\n" + f" 3. Filter out datasets with insufficient frames\n" + f" 4. Reduce skip_first_frames in config: data.skip_first_frames (currently {self.skip_first_frames})" + ) + raise ValueError(error_msg) + + print( + f"[{split}] All {len(self.dataset_dirs)} datasets have sufficient frames for step_sizes {self.random_step_sizes}" + ) + else: + # Disabled: use the fixed skip_step_size + print( + f"[{split}] Multi-timescale training disabled, using skip_step_size={self.skip_step_size}" + ) + + def __len__(self) -> int: + return len(self.all_sequences) + + def get_cache_stats(self) -> dict: + """Return cache statistics for monitoring.""" + return { + "numpy_cache_info": _cached_numpy_load.cache_info()._asdict(), + "reference_frames_cached": len(self._reference_frames_cache), + "reference_frames_datasets": list(self._reference_frames_cache.keys()), + } + + @staticmethod + def clear_numpy_cache(): + """Clear the global numpy LRU cache to free memory.""" + _cached_numpy_load.cache_clear() + logger.info("Cleared numpy LRU cache") + + def _load_frame_objects( + self, frame_dir: Path, pc_resolution: Optional[int] = None + ) -> List["np.ndarray"]: + """ + Load per-object vertex/point data for a single frame. + + Args: + frame_dir: Path to frame directory + pc_resolution: If provided, load point cloud at this resolution instead of mesh vertices + + Returns: + List of (Ni, 3) arrays, one per object + """ + # Data exceeding the configured limits is rejected rather than truncated. + objs = [] + idx = 0 + + if pc_resolution is not None: + if self._object_scale is not None: + raise NotImplementedError( + "data.object_scale_file is not supported with point-cloud inputs" + ) + # Point cloud mode: load from obj_i_pc{N}_pos.npy + while (frame_dir / f"obj_{idx}_pc{pc_resolution}_pos.npy").exists(): + pc_file = frame_dir / f"obj_{idx}_pc{pc_resolution}_pos.npy" + V = cached_load_npy(pc_file) # (N_points, 3) + + # Point clouds have a fixed size; no vertex limit check needed. + objs.append(V.astype("float32", copy=False)) + idx += 1 + + # Mesh vertices exist but the requested point cloud resolution does not. + if idx == 0 and (frame_dir / "obj_0_pos.npy").exists(): + raise FileNotFoundError( + f"Point cloud files not found at resolution {pc_resolution} in {frame_dir}.\n" + f"Expected: obj_i_pc{pc_resolution}_pos.npy\n" + f"Mesh vertices exist but point cloud mode is enabled.\n" + f"Please use a dataset with pre-computed point clouds." + ) + else: + # Mesh vertex mode: load from obj_i_pos.npy (packed layout: the same arrays from frames.npy) + packed = ( + load_frame(frame_dir.parent, int(frame_dir.name.split("_")[-1])) + if is_packed(frame_dir.parent) + else None + ) + while ( + (idx < len(packed)) + if packed is not None + else (frame_dir / f"obj_{idx}_pos.npy").exists() + ): + V = ( + packed[idx] + if packed is not None + else cached_load_npy(frame_dir / f"obj_{idx}_pos.npy") + ) # (Ni, 3) + if self._object_scale is not None: + V = self._object_scale.apply(frame_dir / f"obj_{idx}_pos.npy", V) + + if ( + self.max_vertices_per_object is not None + and V.shape[0] > self.max_vertices_per_object + ): + raise ValueError( + f"Object {idx} in {frame_dir} has {V.shape[0]} vertices, " + f"exceeding max_vertices_per_object={self.max_vertices_per_object}.\n" + f"Either reduce vertex count in dataset or increase data.max_vertices_per_object in config." + ) + + objs.append(V.astype("float32", copy=False)) + idx += 1 + + if ( + self.max_objects_per_frame is not None + and len(objs) > self.max_objects_per_frame + ): + raise ValueError( + f"Frame {frame_dir} has {len(objs)} objects, " + f"exceeding max_objects_per_frame={self.max_objects_per_frame}.\n" + f"Either reduce object count in dataset or increase data.max_objects_per_frame in config." + ) + + return objs + + def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]: + info = self.all_sequences[idx] + + # Point Cloud Mode: Sample resolution for this sample + pc_resolution = None + if self.point_cloud_enable: + pc_resolution = np.random.choice(self.point_cloud_resolutions) + + # Multi-timescale training: sample a step size for this sequence. + if self.random_step_size_enable: + step_size = np.random.choice( + self.random_step_sizes, p=self.sampling_weights + ) + + ds_name = info["dataset"] + all_frames = self._all_frames_by_dataset[ds_name] + + # Apply skip_first_frames + if self.skip_first_frames > 0: + all_frames = all_frames[self.skip_first_frames :] + + # Total frames spanned by the sequence, e.g. sequence_length=4 with + # step_size=10 uses frames [0, 10, 20, 30] and spans 31 frames. + student_reach = (self.sequence_length - 1) * step_size + total_frames_needed = student_reach + 1 + + # Frame counts were validated in __init__; assert as a sanity check. + assert ( + len(all_frames) >= total_frames_needed + ), f"Unexpected: insufficient frames in {ds_name} (should have been caught in __init__)" + + # Select a start frame within the valid range. + if self.random_start_frame: + max_start_frame = len(all_frames) - total_frames_needed + start_frame = np.random.randint(0, max_start_frame + 1) + else: + # Deterministic start frame derived from idx (for reproducibility). + max_start_frame = len(all_frames) - total_frames_needed + start_frame = idx % (max_start_frame + 1) + + # Frame sequence [start_frame, start_frame + step_size, ...]. + frame_indices = [ + start_frame + i * step_size for i in range(self.sequence_length) + ] + frame_paths = [all_frames[i] for i in frame_indices] + frame_ids = np.array( + [self._frame_id_cache[(ds_name, fp.name)] for fp in frame_paths], + dtype=np.int32, + ) + + selected_step_size = step_size + else: + # Fixed step size: use pre-built sequences with skip_step_size. + ds_name = info["dataset"] + frame_paths = [Path(p) for p in info["frames"]] + frame_ids = np.asarray(info.get("frame_ids", []), dtype=np.int32) + selected_step_size = self.skip_step_size + + if self._partial is not None: + return self._partial_item(ds_name, frame_ids, selected_step_size) + + # Load per-frame objects (ragged); point clouds are loaded instead of mesh + # vertices when pc_resolution is set. + frames_objs: List[List[np.ndarray]] = [ + self._load_frame_objects(fp, pc_resolution=pc_resolution) + for fp in frame_paths + ] + if pc_resolution is None: + # every frame holds the reference frame's objects with the same vertex counts (a missing or truncated file + # would otherwise become padding or a shorter object without notice) + _ref_nv = self._mesh_nv_by_dataset[ds_name] + for fp, fo in zip(frame_paths, frames_objs): + if [V.shape[0] for V in fo] != _ref_nv: + raise ValueError( + f"{fp}: per-object vertex counts {[V.shape[0] for V in fo]} differ from frame_00000 " + f"{_ref_nv}" + ) + num_objs = max(len(f) for f in frames_objs) + + if self.max_vertices_per_object is not None: + max_nv = int(self.max_vertices_per_object) + else: + max_nv = self._max_nv_by_dataset[ds_name] + + T = len(frames_objs) + + # Allocate padded arrays. + obj_positions = np.full( + (T, num_objs, max_nv, 3), PADDING_VALUE, dtype="float32" + ) + valid_obj_mask = np.zeros((num_objs,), dtype="bool") + valid_vtx_mask = np.zeros((T, num_objs, max_nv), dtype="bool") + + for t, fo in enumerate(frames_objs): + for i, V in enumerate(fo): + nv = min(V.shape[0], max_nv) + obj_positions[t, i, :nv, :] = V[:nv] + valid_obj_mask[i] = True + valid_vtx_mask[t, i, :nv] = True + + # Reference pose is always the first frame of the dataset (global frame_00000), + # cached per dataset. Point cloud mode includes the resolution in the cache key. + cache_key = f"{ds_name}_pc{pc_resolution}" if pc_resolution else ds_name + if cache_key not in self._reference_frames_cache: + ds_root = self.data_path / ds_name + frame0 = ds_root / "frame_00000" + if not (frame0.exists() or is_packed(ds_root)): + raise FileNotFoundError( + f"Missing reference frame for dataset '{ds_name}': {frame0}. " + f"Expected a 'frame_00000' directory (or frames.npy) under {ds_root}." + ) + ref_objs = self._load_frame_objects(frame0, pc_resolution=pc_resolution) + if len(ref_objs) == 0: + raise RuntimeError( + f"Empty reference frame in {frame0}: no per-object vertex files found." + ) + self._reference_frames_cache[cache_key] = ref_objs + else: + ref_objs = self._reference_frames_cache[cache_key] + # Every frame in the sequence must contain all reference objects, each stored + # as a non-empty (N, 2|3) array. Point cloud mode is already validated inside + # _load_frame_objects, so it is skipped here, and so is the packed layout (frames.npy holds every object of + # every frame by construction; scene_frames checks its shape against frames_counts.npy). + expected_num_objs = len(ref_objs) + if pc_resolution is None and not is_packed(self.data_path / ds_name): + for t, fdir in enumerate(frame_paths): + for i_ref in range(expected_num_objs): + npy_path = fdir / f"obj_{i_ref}_pos.npy" + if not npy_path.exists(): + raise RuntimeError( + f"Missing expected file for dataset '{ds_name}': {npy_path} " + f"(t={t}, obj={i_ref}). Ensure all frames contain the same set of objects as frame_00000." + ) + try: + arr_chk = np.load(npy_path) + except Exception as e: + raise RuntimeError( + f"Failed to load file: {npy_path} (t={t}, obj={i_ref}). Error: {e}" + ) + if not ( + arr_chk.ndim == 2 + and arr_chk.shape[1] in (2, 3) + and arr_chk.shape[0] > 0 + ): + raise RuntimeError( + f"Invalid/empty vertices at {npy_path}: shape={arr_chk.shape} " + f"(t={t}, obj={i_ref})." + ) + # Build the reference array with fixed Nv (pad/trim to max_nv). + ref_obj_positions = np.full( + (num_objs, max_nv, 3), PADDING_VALUE, dtype="float32" + ) + for i, V in enumerate(ref_objs[:num_objs]): + nv = min(V.shape[0], max_nv) + ref_obj_positions[i, :nv, :] = V[:nv] + + # Load anchors (per object) for this dataset (computed in __init__) + if ds_name not in self._anchors_by_dataset: + raise KeyError( + f"Missing anchors for dataset '{ds_name}'. Initialization may have failed." + ) + ds_anchor_list = self._anchors_by_dataset[ds_name] + if len(ds_anchor_list) < num_objs: + raise ValueError( + f"Anchors count ({len(ds_anchor_list)}) < objects in sequence ({num_objs}) for dataset '{ds_name}'." + ) + anchor_idx = np.zeros((num_objs, self.anchor_k), dtype=np.int32) + for i in range(num_objs): + idx_i = ds_anchor_list[i] + if idx_i.shape[0] != self.anchor_k: + raise ValueError( + f"Anchor length mismatch for dataset '{ds_name}' obj {i}: got {idx_i.shape[0]}, expected {self.anchor_k}" + ) + anchor_idx[i] = idx_i + vtx_rank = None + if self.fps_order: + # rank[i, v] = position of vertex v in object i's FPS ordering (padding: max_nv) + vtx_rank = np.full((num_objs, max_nv), max_nv, dtype=np.int32) + ds_order_list = self._fps_order_by_dataset[ds_name] + for i in range(num_objs): + order_i = ds_order_list[i] + nv_i = min(order_i.shape[0], max_nv) + vtx_rank[i, order_i[:nv_i]] = np.arange(nv_i, dtype=np.int32) + if ( + self.split == "train" + and self.anchor_resample_prob > 0 + and not self.point_cloud_enable + and np.random.rand() < self.anchor_resample_prob + ): + _before = anchor_idx.copy() + for i in range(num_objs): + nv_i = min(ref_objs[i].shape[0], max_nv) + if nv_i >= self.anchor_k: + anchor_idx[i] = self._farthest_point_sample_np( + ref_objs[i][:nv_i].astype(np.float32, copy=False), + self.anchor_k, + rule=self.anchor_fps_start, + start=int(np.random.randint(nv_i)), + ) + if not getattr(self, "_anchor_resample_logged", False): + self._anchor_resample_logged = True + print( + f"[train] Anchor resampling APPLIED (pid {os.getpid()}): obj0 {_before[0].tolist()} -> {anchor_idx[0].tolist()} " + f"(rule={self.anchor_fps_start})", + flush=True, + ) + + # Object permutation must be applied before faces are gathered so that + # faces_list[i] matches the permuted obj_i. Only faces for the current + # sample's objects are taken. + faces_list = self._faces_by_dataset[ds_name][:num_objs] + if self.permute_debug: + print( + f"[object permutation] dataset={ds_name}, num_objs={num_objs}, faces_list length={len(faces_list)}, total_faces={len(self._faces_by_dataset[ds_name])}" + ) + perm = None # Set only if permutation is applied + if self.permute_enabled and np.random.random() < self.permute_prob: + if num_objs >= 2: + # Shuffle all objects; there is no ground mesh to exclude. + perm = np.arange(num_objs) + np.random.shuffle(perm) + + if self.permute_debug: + print( + f"[object permutation] dataset={ds_name}, idx={idx}, num_objs={num_objs}" + ) + print(f" Original order: {list(range(num_objs))}") + print(f" Shuffled order: {perm.tolist()}") + print( + f" Mapping: {' '.join([f'{i}->{perm[i]}' for i in range(num_objs)])}" + ) + + # Permute all object-indexed data. + obj_positions = obj_positions[:, perm] # (T, Nobj, Nv, 3) + ref_obj_positions = ref_obj_positions[perm] # (Nobj, Nv, 3) + ref_objs = [ref_objs[i] for i in perm] # Permute ref_objs list + # Permute frames_objs (list of lists) + frames_objs = [[fo[i] for i in perm] for fo in frames_objs] + valid_obj_mask = valid_obj_mask[perm] # (Nobj,) + valid_vtx_mask = valid_vtx_mask[:, perm] # (T, Nobj, Nv) + anchor_idx = anchor_idx[perm] # (Nobj, Nq) + if vtx_rank is not None: + vtx_rank = vtx_rank[perm] # (Nobj, Nv) + # Permute faces_list + faces_list = [faces_list[i] for i in perm] + + # time_step_id: first global frame id of this sequence (scalar) + ts_id_val = int(frame_ids[0]) if frame_ids.size > 0 else -1 + + # Get physics parameters for this dataset (permuted if augmentation applied) + phys_params = self._physics_by_dataset[ds_name] # (Nobj_full, 3) + phys_params = phys_params[:num_objs] # Trim to current sample's object count + if perm is not None: + phys_params = phys_params[ + perm + ] # Apply permutation to match permuted objects + + sample = { + "obj_positions": torch.from_numpy(obj_positions), + "valid_obj_mask": torch.from_numpy(valid_obj_mask), + "valid_vtx_mask": torch.from_numpy(valid_vtx_mask), + "ref_obj_positions": torch.from_numpy(ref_obj_positions), + "dt": torch.tensor(1.0, dtype=torch.float32), + "frame_ids": ( + torch.from_numpy(frame_ids) + if frame_ids.size > 0 + else torch.tensor([], dtype=torch.int32) + ), + "time_step_id": torch.tensor(ts_id_val, dtype=torch.int64), + "anchor_idx": torch.from_numpy(anchor_idx), + **( + {"vtx_rank": torch.from_numpy(vtx_rank)} if vtx_rank is not None else {} + ), + "dataset_name": ds_name, + "faces_list": faces_list, # Store permuted faces (list of numpy arrays) + "permutation": ( + perm if perm is not None else np.arange(num_objs) + ), # Store permutation order (identity if not permuted) + "phys_params": torch.from_numpy( + phys_params + ), # (Nobj, 3): [mass, friction, restitution] + "skip_step_size": torch.tensor( + selected_step_size, dtype=torch.float32 + ), # Time step size for temporal encoding (randomly sampled when multi-timescale enabled) + "is_point_cloud": pc_resolution + is not None, # Whether this sample uses point cloud data + } + + if pc_resolution is not None: + sample["point_cloud_resolution"] = pc_resolution + + # Apply noise augmentation (train split only) + if getattr(self, "split", "train") == "train" and getattr( + self, "noise_config", None + ): + rw_cfg = self.noise_config.get("random_walk", {}) + if rw_cfg.get("enable", False): + start_epoch = rw_cfg.get("start_epoch", 0) + cur_epoch = getattr(self, "current_epoch", 0) + if cur_epoch >= start_epoch: + sample, noise_info = apply_random_walk_noise_to_batch( + sample, self.noise_config, is_training=True + ) + sample["noise_info"] = noise_info + elif self.noise_config.get("independent", {}).get("enable", False): + sample = apply_independent_noise( + sample, self.noise_config, is_training=True + ) + + return sample + + def set_epoch(self, epoch: int): + """Set current epoch for noise scheduling.""" + self.current_epoch = int(epoch) + + @staticmethod + def _farthest_point_sample_np( + vertices: "np.ndarray", + k: int, + rule: str = "index0", + start: Optional[int] = None, + ) -> "np.ndarray": + """Deterministic farthest point sampling on an (N,3) array. Returns indices of shape (k,). + + rule='index0': start at vertex 0, ties resolved by the lowest index. + rule='centroid_farthest': start at the vertex farthest from the centroid and resolve every + tie by the lexicographically smallest coordinates, so the result depends on the vertex set + only, not on its ordering. An explicit `start` overrides the rule's start vertex. + """ + if vertices.ndim != 2 or vertices.shape[1] != 3: + raise ValueError(f"Expected vertices shape (N,3); got {vertices.shape}") + n = vertices.shape[0] + if k > n: + raise ValueError(f"FPS requested k={k} but only N={n} vertices available") + V = vertices.astype(np.float64, copy=False) + lex = rule == "centroid_farthest" + + def argmax(d): + if not lex: + return int(np.argmax(d)) + cand = np.flatnonzero(d >= d.max() - 1e-9 * max(1.0, abs(float(d.max())))) + if cand.shape[0] == 1: + return int(cand[0]) + order = np.lexsort((V[cand, 2], V[cand, 1], V[cand, 0])) + return int(cand[order[0]]) + + selected = np.empty((k,), dtype=np.int64) + dists = np.full((n,), np.inf, dtype=np.float64) + if start is not None: + selected[0] = int(start) + elif lex: + selected[0] = argmax(((V - V.mean(axis=0)) ** 2).sum(axis=1)) + else: + selected[0] = 0 + last = V[selected[0]] + for i in range(1, k): + delta = V - last + dists = np.minimum(dists, (delta * delta).sum(axis=1)) + # Never re-select a vertex: meshes may contain coincident duplicates, whose distance to an + # already selected vertex is exactly zero. Excluding selected indices keeps the output a + # permutation; duplicates are appended last (ties resolved by the rule above). + dists[selected[:i]] = -np.inf + idx = argmax(dists) + selected[i] = idx + last = V[idx] + return selected.astype(np.int32, copy=False) + + +def make_collate_fn( + max_objects_per_frame: Optional[int] = None, + max_vertices_per_object: Optional[int] = None, +): + """Factory function to create a collate_fn with fixed padding dimensions. + + Args: + max_objects_per_frame: Fixed maximum number of objects per frame (None = dynamic) + max_vertices_per_object: Fixed maximum vertices per object (None = dynamic) + + Returns: + Configured collate function + """ + + def collate_fn(batch: List[Dict]) -> Dict[str, torch.Tensor]: + """Custom collate function for batching sequences.""" + # Detect multi-object sample: obj_positions present with per-sample padding + if "obj_positions" in batch[0]: + # Determine dimensions + T = batch[0]["obj_positions"].shape[0] + + # Use fixed padding if configured, otherwise compute from batch + if max_objects_per_frame is not None: + max_nobj = int(max_objects_per_frame) + # Validate that no sample exceeds the fixed limit + for s in batch: + _, nobj, _, _ = s["obj_positions"].shape + if nobj > max_nobj: + raise ValueError( + f"Sample has {nobj} objects, exceeding max_objects_per_frame={max_nobj}. " + f"Increase data.max_objects_per_frame in config or reduce object count in dataset." + ) + else: + # Dynamic padding + max_nobj = 0 + for s in batch: + _, nobj, _, _ = s["obj_positions"].shape + if nobj > max_nobj: + max_nobj = nobj + + if max_vertices_per_object is not None: + max_nv = int(max_vertices_per_object) + # Validate that no sample exceeds the fixed limit + for s in batch: + _, _, nv, _ = s["obj_positions"].shape + if nv > max_nv: + raise ValueError( + f"Sample has {nv} vertices, exceeding max_vertices_per_object={max_nv}. " + f"Increase data.max_vertices_per_object in config or reduce vertex count in dataset." + ) + else: + # Dynamic padding + max_nv = 0 + for s in batch: + _, _, nv, _ = s["obj_positions"].shape + if nv > max_nv: + max_nv = nv + + B = len(batch) + device = batch[0]["obj_positions"].device + + # Allocate padded tensors. + obj_positions = torch.full( + (B, T, max_nobj, max_nv, 3), + PADDING_VALUE, + dtype=torch.float32, + device=device, + ) + valid_obj_mask = torch.zeros((B, max_nobj), dtype=torch.bool, device=device) + valid_vtx_mask = torch.zeros( + (B, T, max_nobj, max_nv), dtype=torch.bool, device=device + ) + ref_obj_positions = torch.full( + (B, max_nobj, max_nv, 3), + PADDING_VALUE, + dtype=torch.float32, + device=device, + ) + # Optional anchors, padded per object to max_nobj with fixed Nq. + # anchor_idx pads with index 0; anchor_mask pads with False. + has_anchors = "anchor_idx" in batch[0] + has_rank = "vtx_rank" in batch[0] + if has_rank: + vtx_rank = torch.full( + (B, max_nobj, max_nv), int(max_nv), dtype=torch.int64, device=device + ) + if has_anchors: + anchor_k = int(batch[0]["anchor_idx"].shape[1]) + anchor_idx = torch.zeros( + (B, max_nobj, anchor_k), dtype=torch.int64, device=device + ) + anchor_mask = torch.zeros( + (B, max_nobj, anchor_k), dtype=torch.bool, device=device + ) + + # Fill per-sample + for b, s in enumerate(batch): + t_, nobj, nv, _ = s["obj_positions"].shape + obj_positions[b, :, :nobj, :nv, :] = s["obj_positions"] + valid_obj_mask[b, :nobj] = s["valid_obj_mask"] + valid_vtx_mask[b, :, :nobj, :nv] = s["valid_vtx_mask"] + ref_obj_positions[b, :nobj, :nv, :] = s["ref_obj_positions"] + if has_rank: + if "vtx_rank" not in s: + raise KeyError( + "Batch sample missing 'vtx_rank' while others provide it" + ) + r = s["vtx_rank"].to(device=device, dtype=torch.int64) + vtx_rank[b, :nobj, :nv] = torch.where( + r[:, :nv] >= nv, + torch.full_like(r[:, :nv], int(max_nv)), + r[:, :nv], + ) + if has_anchors: + if "anchor_idx" not in s: + raise KeyError( + "Batch sample missing 'anchor_idx' while others provide it" + ) + if s["anchor_idx"].dim() != 2: + raise ValueError( + f"anchor_idx must be (Nobj, Nq); got {tuple(s['anchor_idx'].shape)}" + ) + Nq = int(s["anchor_idx"].shape[1]) + if Nq != anchor_k: + raise ValueError( + f"Inconsistent anchor_k across batch: {Nq} vs {anchor_k}" + ) + if s["anchor_idx"].shape[0] != nobj: + raise ValueError( + f"anchor_idx Nobj ({s['anchor_idx'].shape[0]}) != sample Nobj ({nobj})" + ) + anchor_idx[b, :nobj, :] = s["anchor_idx"].to( + device=device, dtype=torch.int64 + ) + # With strict loader checks, all anchors are valid; set mask True + anchor_mask[b, :nobj, :] = True + + collated = { + "obj_positions": obj_positions, + "valid_obj_mask": valid_obj_mask, + "valid_vtx_mask": valid_vtx_mask, + "ref_obj_positions": ref_obj_positions, + "dt": torch.stack([s["dt"] for s in batch], dim=0), + } + if has_anchors: + collated["anchor_idx"] = anchor_idx + collated["anchor_mask"] = anchor_mask + if has_rank: + collated["vtx_rank"] = vtx_rank + + # Handle phys_params (Nobj, 3) - pad to max_nobj with PADDING_VALUE + if "phys_params" in batch[0]: + phys_params_padded = torch.full( + (B, max_nobj, 3), PADDING_VALUE, dtype=torch.float32 + ) + for b, s in enumerate(batch): + nobj = s["phys_params"].shape[0] + phys_params_padded[b, :nobj, :] = s["phys_params"] + collated["phys_params"] = phys_params_padded + + # Handle skip_step_size (scalar) - can be stacked directly + if "skip_step_size" in batch[0]: + collated["skip_step_size"] = torch.stack( + [s["skip_step_size"] for s in batch], dim=0 + ) + + # Pass through any remaining keys; tensors are stacked and must agree in shape. + for key in batch[0].keys(): + if key in collated: + continue + if isinstance(batch[0][key], torch.Tensor): + collated[key] = torch.stack( + [sample[key] for sample in batch], dim=0 + ) + else: + collated[key] = [sample[key] for sample in batch] + + return collated + + return collate_fn diff --git a/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/data/gpu_augmentation.py b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/data/gpu_augmentation.py new file mode 100644 index 0000000000..4a00012cb1 --- /dev/null +++ b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/data/gpu_augmentation.py @@ -0,0 +1,363 @@ +# Copyright (c) 2026 Zhiyang Dou. SPDX-License-Identifier: MIT +# Derived from the pinned author release; see THIRD_PARTY_LICENSES. + +""" +GPU-accelerated data augmentation for training. + +All augmentations operate on GPU tensors and are applied in the training +loop after the batch has been moved to the GPU. + +data.augmentation keys: +- rotation_augmentation: rotation about the vertical (z) axis. One angle per + batch, a multiple of angle_step_size, with probability prob_use_rotation; + with per_sample: true every sample gets its own uniform angle in [0, 360) + (probability prob_use_rotation per sample). +- reflection_augmentation: per sample with probability prob, the scene is + mirrored in the vertical x = 0 plane (x -> -x). Mirrored rigid bodies under + gravity, contact and isotropic friction follow the mirrored trajectories, + so the mirrored sequence is again a valid one. +- rigid_noise: the input frames (all but the last frame of a sample) get a + rigid random-walk perturbation per object (rotation about the object's + centroid plus a translation), written to batch['obj_positions_input']; + batch['obj_positions'] stays the ground truth the losses compare against. + pos_std (m) and rot_std_deg are the standard deviations at the last input + frame for step size 1, scaled by step_size ** step_size_exponent. +""" + +from __future__ import annotations + +import torch +import math +from typing import Dict +import logging + +logger = logging.getLogger(__name__) + + +def rotvec_to_matrix(rotvec: torch.Tensor) -> torch.Tensor: + """Rotation matrices (..., 3, 3) of rotation vectors (..., 3) (Rodrigues' formula).""" + theta = rotvec.norm(dim=-1, keepdim=True) # (..., 1) + small = theta < 1e-8 + axis = rotvec / torch.where(small, torch.ones_like(theta), theta) + x, y, z = axis.unbind(-1) + zero = torch.zeros_like(x) + K = torch.stack([zero, -z, y, z, zero, -x, -y, x, zero], dim=-1).view( + *rotvec.shape[:-1], 3, 3 + ) + s = torch.sin(theta).unsqueeze(-1) + c = torch.cos(theta).unsqueeze(-1) + eye = torch.eye(3, dtype=rotvec.dtype, device=rotvec.device).expand(K.shape) + R = eye + s * K + (1.0 - c) * (K @ K) + return torch.where(small.unsqueeze(-1), eye, R) + + +class GPUAugmentation: + """ + GPU-accelerated augmentation pipeline. + + Features: + - Rotation augmentation (Z-axis only, no ground mesh), per batch or per sample + - Reflection augmentation (x -> -x, per sample) + - Rigid random-walk noise on the input frames + - All operations on GPU tensors + """ + + def __init__(self, config: Dict, device: torch.device): + """ + Initialize GPU augmentation. + + Args: + config: Augmentation configuration + device: Target GPU device + """ + self.device = device + self.config = config + aug_cfg = config.get("data", {}).get("augmentation", {}) + + # Rotation augmentation settings + rotation_cfg = aug_cfg.get("rotation_augmentation", {}) + self.rotation_enabled = rotation_cfg.get("enable", False) + self.rotation_angle_step = rotation_cfg.get("angle_step_size", 30) + self.rotation_prob = rotation_cfg.get("prob_use_rotation", 0.3) + self.rotation_per_sample = bool(rotation_cfg.get("per_sample", False)) + + # Pre-compute rotation angles (only if rotation enabled and valid step size) + if self.rotation_enabled and self.rotation_angle_step > 0: + self.rotation_angles = torch.arange( + self.rotation_angle_step, + 360, + self.rotation_angle_step, + dtype=torch.float32, + device=device, + ) + else: + # Rotation disabled or invalid step size - create empty tensor + self.rotation_angles = torch.tensor([], dtype=torch.float32, device=device) + + reflection_cfg = aug_cfg.get("reflection_augmentation", {}) or {} + self.reflection_enabled = bool(reflection_cfg.get("enable", False)) + self.reflection_prob = float(reflection_cfg.get("prob", 0.5)) + + noise_cfg = aug_cfg.get("rigid_noise", {}) or {} + self.noise_enabled = bool(noise_cfg.get("enable", False)) + self.noise_pos_std = float(noise_cfg.get("pos_std", 0.0)) + self.noise_rot_std = math.radians(float(noise_cfg.get("rot_std_deg", 0.0))) + self.noise_step_exponent = float(noise_cfg.get("step_size_exponent", 1.0)) + if self.noise_enabled and config.get("data", {}).get( + "center_scene_by_t0", False + ): + # centring shifts the network inputs in place; with separate input frames the targets would not follow + raise ValueError( + "data.augmentation.rigid_noise cannot be combined with data.center_scene_by_t0" + ) + + # Counters for periodic logging. + self.rotation_total_calls = 0 + self.rotation_applied_count = 0 + + logger.info(f"\n{'='*80}") + logger.info(f"[GPU Augmentation] Initialization") + logger.info(f"{'='*80}") + logger.info(f"Rotation Augmentation:") + logger.info(f" - Enabled: {self.rotation_enabled}") + if self.rotation_enabled: + if self.rotation_per_sample: + logger.info(f" - Per sample, uniform angle in [0, 360)") + else: + logger.info(f" - Angle step size: {self.rotation_angle_step}°") + logger.info( + f" - Rotation angles: {self.rotation_angles.cpu().tolist()}" + ) + logger.info(f" - Probability: {self.rotation_prob:.2%}") + logger.info( + f"Reflection Augmentation (x -> -x): enabled={self.reflection_enabled}" + + ( + f", probability {self.reflection_prob:.2%} per sample" + if self.reflection_enabled + else "" + ) + ) + logger.info( + f"Rigid input noise: enabled={self.noise_enabled}" + + ( + f", pos_std={self.noise_pos_std} m, rot_std={math.degrees(self.noise_rot_std):.3f} deg at the last " + f"input frame (step size 1), x step_size^{self.noise_step_exponent:g}" + if self.noise_enabled + else "" + ) + ) + logger.info(f"{'='*80}\n") + + @property + def enabled(self) -> bool: + return self.rotation_enabled or self.reflection_enabled or self.noise_enabled + + def apply_rotation_augmentation( + self, batch: Dict[str, torch.Tensor], split: str = "train" + ) -> Dict[str, torch.Tensor]: + """ + Apply rotation augmentation around Z-axis (GPU version, no ground mesh). + + Args: + batch: Batch dict with GPU tensors + split: 'train', 'val', or 'test' (only train gets augmented) + + Returns: + Augmented batch dict + """ + # Only apply to training split + if split != "train": + batch["rotation_angle_deg"] = 0.0 + return batch + + self.rotation_total_calls += 1 + + if not self.rotation_enabled: + batch["rotation_angle_deg"] = 0.0 + return batch + + if self.rotation_per_sample: + return self._rotate_per_sample(batch) + + if len(self.rotation_angles) == 0: + raise ValueError( + f"Rotation enabled but no valid angles available. " + f"angle_step_size={self.rotation_angle_step} must be > 0" + ) + + # Decide whether to rotate (probabilistic) + rand_val = torch.rand(1, device=self.device).item() + if rand_val > self.rotation_prob: + batch["rotation_angle_deg"] = 0.0 + return batch + + # Randomly select rotation angle + angle_idx = torch.randint( + 0, len(self.rotation_angles), (1,), device=self.device + ) + angle_deg = self.rotation_angles[angle_idx].item() + + if angle_deg == 0: + batch["rotation_angle_deg"] = 0.0 + return batch + + self.rotation_applied_count += 1 + if self.rotation_applied_count <= 5 or self.rotation_total_calls % 100 == 0: + apply_rate = self.rotation_applied_count / self.rotation_total_calls + logger.info( + f"[GPU Rotation Aug] Applying {angle_deg:.0f}° | " + f"Batch {self.rotation_total_calls} | " + f"Applied: {self.rotation_applied_count}/{self.rotation_total_calls} ({apply_rate:.2%})" + ) + + # Create rotation matrix (Z-axis) + angle_rad = math.radians(angle_deg) + cos_a = math.cos(angle_rad) + sin_a = math.sin(angle_rad) + + R = torch.tensor( + [[cos_a, -sin_a, 0], [sin_a, cos_a, 0], [0, 0, 1]], + dtype=torch.float32, + device=self.device, + ) + + # Rotate obj_positions (B, T, Nobj, Nv, 3); all objects rotate since + # there is no ground mesh. + obj_pos = batch["obj_positions"] + obj_pos_flat = obj_pos.reshape(-1, 3) + obj_pos_rotated = torch.matmul(obj_pos_flat, R.T) + batch["obj_positions"] = obj_pos_rotated.view(obj_pos.shape) + + # Rotate ref_obj_positions: (B, Nobj, Nv, 3) + ref_pos = batch["ref_obj_positions"] + ref_pos_flat = ref_pos.reshape(-1, 3) + ref_pos_rotated = torch.matmul(ref_pos_flat, R.T) + batch["ref_obj_positions"] = ref_pos_rotated.view(ref_pos.shape) + + # Record the applied angle so ground-truth overlays can be rotated to match. + batch["rotation_angle_deg"] = angle_deg + + return batch + + def _per_sample_linear( + self, batch: Dict[str, torch.Tensor], M: torch.Tensor + ) -> None: + """Apply the per-sample linear maps M (B, 3, 3) to obj_positions (B,T,Nobj,Nv,3) and ref_obj_positions.""" + # einsum may return non-contiguous strides; the model reshapes these tensors with .view() + batch["obj_positions"] = torch.einsum( + "bij,btonj->btoni", M, batch["obj_positions"] + ).contiguous() + ref = batch["ref_obj_positions"] + if ref.dim() == 4: + batch["ref_obj_positions"] = torch.einsum( + "bij,bonj->boni", M, ref + ).contiguous() + else: + raise ValueError( + f"per-sample augmentation needs per-sample ref_obj_positions (B,Nobj,Nv,3); got {tuple(ref.shape)}" + ) + + def _rotate_per_sample( + self, batch: Dict[str, torch.Tensor] + ) -> Dict[str, torch.Tensor]: + B = batch["obj_positions"].shape[0] + angle = torch.rand(B, device=self.device) * (2 * math.pi) + angle = torch.where( + torch.rand(B, device=self.device) < self.rotation_prob, + angle, + torch.zeros_like(angle), + ) + c, s = torch.cos(angle), torch.sin(angle) + zero, one = torch.zeros_like(c), torch.ones_like(c) + R = torch.stack([c, -s, zero, s, c, zero, zero, zero, one], dim=-1).view( + B, 3, 3 + ) + self._per_sample_linear(batch, R) + batch["rotation_angle_deg"] = ( + 0.0 # no single angle: ground-truth overlays cannot be rotated to match + ) + batch["rotation_angles_deg_per_sample"] = torch.rad2deg(angle) + return batch + + def apply_reflection( + self, batch: Dict[str, torch.Tensor], split: str = "train" + ) -> Dict[str, torch.Tensor]: + if split != "train" or not self.reflection_enabled: + return batch + B = batch["obj_positions"].shape[0] + flip = torch.rand(B, device=self.device) < self.reflection_prob + sx = torch.where( + flip, -torch.ones(B, device=self.device), torch.ones(B, device=self.device) + ) + M = torch.diag_embed( + torch.stack([sx, torch.ones_like(sx), torch.ones_like(sx)], dim=-1) + ) + self._per_sample_linear(batch, M) + batch["reflected"] = flip + return batch + + def apply_rigid_noise( + self, batch: Dict[str, torch.Tensor], split: str = "train" + ) -> Dict[str, torch.Tensor]: + if split != "train" or not self.noise_enabled: + return batch + pos = batch["obj_positions"] # (B,T,Nobj,Nv,3) + B, T, Nobj, Nv, _ = pos.shape + Tin = T - 1 # the network inputs are frames 0..T-2; frame T-1 is only a target + if Tin < 1: + return batch + scale = torch.ones(B, device=pos.device) + if "skip_step_size" in batch: + scale = ( + batch["skip_step_size"].to(pos.device).float() + ** self.noise_step_exponent + ) + scale = scale.view(B, 1, 1, 1) + # random walk over the input frames: the last input frame has the configured standard deviation + rotvec = ( + torch.randn(B, Tin, Nobj, 3, device=pos.device) + * (self.noise_rot_std / math.sqrt(Tin)) + ).cumsum(1) * scale + shift = ( + torch.randn(B, Tin, Nobj, 3, device=pos.device) + * (self.noise_pos_std / math.sqrt(Tin)) + ).cumsum(1) * scale + R = rotvec_to_matrix(rotvec) # (B,Tin,Nobj,3,3) + x = pos[:, :Tin] + vmask = batch["valid_vtx_mask"][:, :Tin].to(torch.bool) # (B,Tin,Nobj,Nv) + w = vmask.unsqueeze(-1).to(x.dtype) + centroid = (x * w).sum(3, keepdim=True) / w.sum(3, keepdim=True).clamp_min( + 1.0 + ) # (B,Tin,Nobj,1,3) + xn = ( + torch.einsum("btoij,btonj->btoni", R, x - centroid) + + centroid + + shift.unsqueeze(3) + ) + xn = torch.where(vmask.unsqueeze(-1), xn, x) # padding stays as it is + batch["obj_positions_input"] = torch.cat([xn, pos[:, Tin:]], dim=1).contiguous() + return batch + + def __call__( + self, batch: Dict[str, torch.Tensor], split: str = "train" + ) -> Dict[str, torch.Tensor]: + """ + Apply all enabled augmentations. + + Args: + batch: Batch dict with GPU tensors + split: 'train', 'val', or 'test' + + Returns: + Augmented batch dict + """ + if self.rotation_enabled: + batch = self.apply_rotation_augmentation(batch, split) + else: + # Rotation disabled: record a zero angle for consistency. + batch["rotation_angle_deg"] = 0.0 + batch = self.apply_reflection(batch, split) + # last: the noise perturbs the (rotated, mirrored) input frames + batch = self.apply_rigid_noise(batch, split) + + return batch diff --git a/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/data/object_scale.py b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/data/object_scale.py new file mode 100644 index 0000000000..7de594aa6a --- /dev/null +++ b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/data/object_scale.py @@ -0,0 +1,160 @@ +# Copyright (c) 2026 Zhiyang Dou. SPDX-License-Identifier: MIT +# Derived from RigidFormer at 160fe0852c51cffae7b7440452b58d0b49db6df6. +# See the model family's THIRD_PARTY_LICENSES for the complete license. + +"""Per-object metadata scale for converted MOVi vertex data (data.object_scale_file). + +data_generation/movi_to_sim_data.py wrote every object's vertices as R v_asset + p and never applied the metadata +`size`, so stored meshes do not have the simulated dimensions (MOVi-Sphere: all spheres ~1.0 m across, simulated radii +0.35 / 0.7 m; resting mesh bottoms at z ~ -0.15 / +0.2 m instead of the floor). With a scale file from +data_generation/build_object_scales.py, every array read from /frame_/obj__pos.npy is scaled about the +object origin p_t (vis_traj.npy[0, t, i, :3], equal to the metadata positions): + + v' = p_t + size_i (v - p_t) + +after which resting objects of every size have their lowest vertex at z ~ 0.002 m. Anchor and FPS-order caches are +unchanged (farthest-point sampling is invariant to a uniform scaling about any point). Without a scale file nothing +changes. +""" + +from __future__ import annotations +import json +from pathlib import Path +from typing import Dict, Optional + +import numpy as np + +try: + from .scene_frames import load_frame +except ImportError: # imported as src.data.object_scale + from .scene_frames import load_frame + + +class ObjectScale: + def __init__(self, scale_file: str): + data = json.load(open(scale_file)) + self.scale_file = str(scale_file) + # converted dataset the sizes belong to (scene names movi_ repeat across MOVi-A/B/Sphere) + self.converted = data.get("converted") + self.scales: Dict[str, np.ndarray] = { + k: np.asarray(v, dtype=np.float64) for k, v in data["scales"].items() + } + self.fingerprint = { + k: np.asarray(v, dtype=np.float64) + for k, v in (data.get("fingerprint") or {}).items() + } + if "radius0" not in data: + raise ValueError( + f"{scale_file} has no radius0 entries (the guard against vertex data that already carries the " + f"size); rebuild it with data_generation/build_object_scales.py" + ) + self.radius0 = { + k: np.asarray(v, dtype=np.float64) for k, v in data["radius0"].items() + } + for name, table in ( + ("scales", self.scales), + ("fingerprint", self.fingerprint), + ("radius0", self.radius0), + ): + bad = [ + k + for k, v in table.items() + if not np.isfinite(v).all() + or (name != "fingerprint" and not (v > 0).all()) + ] + if bad: + raise ValueError( + f"{scale_file}: non-finite or non-positive {name} for {len(bad)} scene(s), e.g. {bad[:3]}" + ) + self._origins: Dict[str, np.ndarray] = {} + + def origins(self, scene_dir: Path) -> np.ndarray: + key = scene_dir.name + if key not in self._origins: + o = np.load(scene_dir / "vis_traj.npy")[0, :, :, :3].astype( + np.float64 + ) # (T, Nobj, 3) + if not np.isfinite(o).all(): + raise ValueError( + f'ObjectScale: non-finite object origins in {scene_dir / "vis_traj.npy"}' + ) + self._origins[key] = o + return self._origins[key] + + def apply(self, vertex_path: Path, V: np.ndarray) -> np.ndarray: + """Scaled copy of V, the array stored at /frame_/obj__pos.npy (V itself is never modified).""" + vertex_path = Path(vertex_path) + frame_dir = vertex_path.parent + scene_dir = frame_dir.parent + if not frame_dir.name.startswith("frame_") or not vertex_path.name.startswith( + "obj_" + ): + raise ValueError(f"ObjectScale: not a per-frame vertex file: {vertex_path}") + if scene_dir.name not in self.scales: + raise KeyError( + f"ObjectScale: scene {scene_dir.name} missing from {self.scale_file}" + ) + t = int(frame_dir.name.split("_")[1]) + i = int(vertex_path.name.split("_")[1]) + s = self.scales[scene_dir.name][i] + p = self.origins(scene_dir)[t, i] + return (p + s * (V.astype(np.float64) - p)).astype(V.dtype) + + +def from_config(config: dict) -> Optional[ObjectScale]: + f = (config.get("data", {}) or {}).get("object_scale_file") + return ObjectScale(f) if f else None + + +def check_object_counts(object_scale: "ObjectScale", data_path, scenes) -> None: + """For every scene: as many sizes as objects in frame_00000; object 0's frame-0 origin equal to the sidecar's + fingerprint (MOVi-A and MOVi-Sphere share seeds and hence object counts, so the counts alone cannot tell them + apart); and every object's largest frame-0 vertex distance from its origin equal to the sidecar's radius0, i.e. the + stored vertices are still the unscaled ones the sizes apply to (data re-converted with the size would otherwise be + scaled twice).""" + import os + + bad = [] + for sc in scenes: + if sc not in object_scale.scales: + bad.append(f"{sc}: missing") + continue + frame0 = load_frame( + Path(data_path) / sc, 0 + ) # either layout (src/data/scene_frames.py) + n = len(frame0) + if n != len(object_scale.scales[sc]): + bad.append(f"{sc}: {n} objects vs {len(object_scale.scales[sc])} sizes") + continue + if sc not in object_scale.fingerprint or sc not in object_scale.radius0: + bad.append(f"{sc}: no fingerprint / radius0 in the sidecar") + continue + p0 = np.load(os.path.join(str(data_path), sc, "vis_traj.npy"))[ + 0, 0, :, :3 + ].astype( + np.float64 + ) # (Nobj, 3) + if ( + not np.abs(p0[0] - object_scale.fingerprint[sc]).max() <= 1e-4 + ): # (a NaN fails this test too) + bad.append( + f"{sc}: frame-0 origin {p0[0].round(4).tolist()} vs fingerprint {object_scale.fingerprint[sc].tolist()}" + ) + continue + r = np.array( + [ + np.linalg.norm(frame0[o].astype(np.float64) - p0[o], axis=1).max() + for o in range(n) + ] + ) + if ( + not np.abs(r - object_scale.radius0[sc]).max() <= 1e-4 + ): # (a NaN fails this test too) + bad.append( + f"{sc}: frame-0 vertex radii {r.round(4).tolist()} vs radius0 {object_scale.radius0[sc].tolist()} " + f"(sizes {object_scale.scales[sc].round(4).tolist()})" + ) + if bad: + raise ValueError( + f"{object_scale.scale_file}: {len(bad)} scene(s) do not match the data, e.g. {bad[:3]}" + ) diff --git a/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/data/scene_frames.py b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/data/scene_frames.py new file mode 100644 index 0000000000..c47ee39e04 --- /dev/null +++ b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/data/scene_frames.py @@ -0,0 +1,96 @@ +# Copyright (c) 2026 Zhiyang Dou. SPDX-License-Identifier: MIT +# Derived from RigidFormer at 160fe0852c51cffae7b7440452b58d0b49db6df6. +# See the model family's THIRD_PARTY_LICENSES for the complete license. + +"""Per-frame vertices of a converted scene, in either layout written by data_generation/movi_to_sim_data.py: + + per-file /frame_/obj__pos.npy one array per object and frame (T x N files) + packed /frames.npy (T, sum_i Nv_i, 3): the same arrays, concatenated per frame + /frames_counts.npy (N,): the vertex count of every object + +Both hold the same float32 values (the stored, unscaled vertices R(q) v_asset + p). The packed layout needs two files +per scene instead of ~5,000 (MOVi-B: 1,200 x 2 files instead of 3.8 million), so it fits on any local disk or a +network file system. The frame paths /frame_ keep naming the frames in both layouts; with the packed layout +they are names only (no folders exist). +""" + +from __future__ import annotations +from collections import OrderedDict +from pathlib import Path +from typing import List + +import numpy as np + +_OPEN_LIMIT = 128 # packed scenes kept memory-mapped per process +_open = OrderedDict() # scene dir -> (frames memmap, offsets) + + +def clear_cache(root: Path) -> None: + """Forget packed scenes under a newly verified source root.""" + root = Path(root).resolve() + for key in list(_open): + if Path(key).resolve().is_relative_to(root): + del _open[key] + + +def is_packed(scene_dir: Path) -> bool: + return (Path(scene_dir) / "frames.npy").is_file() + + +def _packed(scene_dir: Path): + key = str(scene_dir) + if key in _open: + _open.move_to_end(key) + return _open[key] + frames = np.load(Path(scene_dir) / "frames.npy", mmap_mode="r") + counts = np.load(Path(scene_dir) / "frames_counts.npy").astype(np.int64) + if ( + frames.ndim != 3 + or frames.shape[2] != 3 + or frames.shape[1] != int(counts.sum()) + or (counts <= 0).any() + ): + raise ValueError( + f"{scene_dir}: frames.npy {frames.shape} does not hold the vertex counts {counts.tolist()}" + ) + entry = (frames, np.concatenate([[0], np.cumsum(counts)])) + _open[key] = entry + while len(_open) > _OPEN_LIMIT: + _open.popitem(last=False) + return entry + + +def num_frames(scene_dir: Path) -> int: + if is_packed(scene_dir): + return int(_packed(scene_dir)[0].shape[0]) + return len(frame_dirs(scene_dir)) + + +def frame_dirs(scene_dir: Path) -> List[Path]: + """frame_ paths of a scene, t = 0 .. T-1 (per-file layout: the existing folders, sorted).""" + scene_dir = Path(scene_dir) + if is_packed(scene_dir): + return [ + scene_dir / f"frame_{t:05d}" for t in range(_packed(scene_dir)[0].shape[0]) + ] + return sorted(p for p in scene_dir.glob("frame_*") if p.is_dir()) + + +def load_frame(scene_dir: Path, t: int) -> List[np.ndarray]: + """The stored vertices (Nv_i, 3) of every object at frame t, float32 (copies for the packed layout).""" + scene_dir = Path(scene_dir) + if is_packed(scene_dir): + frames, off = _packed(scene_dir) + if not 0 <= t < frames.shape[0]: + raise IndexError( + f"{scene_dir}: frame {t} outside 0 .. {frames.shape[0] - 1}" + ) + row = np.asarray(frames[t]) + return [row[off[i] : off[i + 1]].copy() for i in range(len(off) - 1)] + frame_dir = scene_dir / f"frame_{t:05d}" + out = [] + while (frame_dir / f"obj_{len(out)}_pos.npy").exists(): + out.append(np.load(frame_dir / f"obj_{len(out)}_pos.npy")) + if not out: + raise FileNotFoundError(f"{frame_dir}: no obj__pos.npy") + return out diff --git a/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/checkpoint.py b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/checkpoint.py index 8f5fe6ea22..8b205a2541 100644 --- a/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/checkpoint.py +++ b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/checkpoint.py @@ -22,7 +22,7 @@ import torch from .model import RigidFormerCfg, RigidFormerModel -__all__ = ["load_rigidformer_checkpoint"] +__all__ = ["load_rigidformer_checkpoint", "export_rigidformer_checkpoint"] def load_rigidformer_checkpoint( @@ -76,3 +76,68 @@ def load_rigidformer_checkpoint( raise ValueError(f"Official checkpoint tensor differs: {name}") model.network.load_state_dict(state, strict=True) return model.to(device).eval() + + +def export_rigidformer_checkpoint( + source: str | Path, destination: str | Path, *, weights: str = "ema" +) -> Path: + """Export a native SL training checkpoint to tensor/config inference weights. + + Args: + source: Checkpoint produced by the native author-based SL recipe. + destination: Fresh output path; existing files are rejected. + weights: Explicit ema or raw model selection. + + Returns: + Written inference checkpoint, loadable by load_rigidformer_checkpoint. + """ + import numpy as np + + if weights not in ("ema", "raw"): + raise ValueError("weights must be ema or raw") + output = Path(destination) + if output.exists(): + raise FileExistsError(output) + try: + reconstruct = np._core.multiarray._reconstruct + except AttributeError: # NumPy 1.x + reconstruct = np.core.multiarray._reconstruct + safe = [ + (reconstruct, "numpy.core.multiarray._reconstruct"), + (reconstruct, "numpy._core.multiarray._reconstruct"), + np.ndarray, + np.dtype, + type(np.dtype(np.uint32)), + ] + with torch.serialization.safe_globals(safe): + state = torch.load(source, map_location="cpu", weights_only=True, mmap=True) + signature = state.get("signature", {}) + if signature.get("family") != "rigidformer-author-v1": + raise ValueError("Expected a native author-based RigidFormer SL checkpoint") + raw = state["learner"]["model"] + if weights == "ema": + extra = state.get("training_states", {}).get("ema", {}).get("state") + if extra is None: + raise ValueError("Checkpoint has no EMA; select raw explicitly") + selected = {k: v.to(raw[k].dtype) for k, v in extra["shadow"].items()} + else: + selected = raw + if not all(k.startswith("network.") for k in selected): + raise ValueError("Native model weight namespace differs") + payload = { + "model_state_dict": { + k.removeprefix("network."): v.detach().clone() for k, v in selected.items() + }, + "config": signature["reference_config"], + "source": { + "format": "embodichain.sl", + "author_revision": signature["author_revision"], + "num_updates": state["learner"]["num_updates"], + "weights": weights, + }, + } + output.parent.mkdir(parents=True, exist_ok=True) + temporary = output.with_suffix(output.suffix + ".tmp") + torch.save(payload, temporary) + temporary.replace(output) + return output diff --git a/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/data.py b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/data.py new file mode 100644 index 0000000000..4e512ad991 --- /dev/null +++ b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/data.py @@ -0,0 +1,314 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Replayable author MOVi packed sequences for unified supervised learning.""" + +from __future__ import annotations +from collections.abc import Mapping +from copy import deepcopy +from pathlib import Path +from typing import Any +import hashlib +import json +import random +import numpy as np +import torch +from embodichain.learning.sl import SupervisedBatch +from .supervised import rigidformer_batch + +__all__ = ["MOViSequenceSource"] + + +def _check_manifest_coverage(root: Path, manifest: Mapping[str, Any]) -> None: + """Require every readable geometry payload to belong to the committed set.""" + payloads = { + str(path.relative_to(root)) + for path in root.rglob("*") + if path.is_file() + and ( + path.suffix == ".npy" + or path.name in {"splits.json", "object_scales.json", "physics_params.json"} + ) + } + missing = payloads - manifest["files"].keys() + if missing: + raise ValueError( + f"Prepared geometry is not committed by manifest: {sorted(missing)[0]}" + ) + + +def _verified_manifest(root: Path) -> dict[str, Any]: + """Read a committed manifest and verify its actual prepared payloads.""" + path = root / "manifest.json" + if not path.is_file(): + raise ValueError("Prepared data requires manifest.json with file hashes") + manifest = json.loads(path.read_text()) + _check_manifest_coverage(root, manifest) + for rel, expected in manifest["files"].items(): + h = hashlib.sha256() + with (root / rel).open("rb") as f: + for block in iter(lambda: f.read(1 << 20), b""): + h.update(block) + if h.hexdigest() != expected: + raise ValueError(f"Prepared geometry hash differs: {rel}") + return manifest + + +class MOViSequenceSource: + """Read author packed scenes and checkpoint a rank-local shuffled traversal. + + Args: + root: Converted MOVi-Sphere, MOVi-A or MOVi-B directory. + split_file: Frozen author scene split JSON. + object_scale_file: Frozen author size sidecar for unscaled stored vertices. + split: train, val or test. + data_config: Author sampling and augmentation settings. + seed: Epoch-shuffle seed; per-batch augmentation RNG is caller-owned. + max_batch_bytes: Bound before collating a padded geometry batch. + + Register this source in SupervisedTrainer.training_states so sampling is + committed after successful updates. Core SL never imports this adapter. + """ + + def __init__( + self, + root: str | Path, + split_file: str | Path, + object_scale_file: str | Path, + *, + split: str = "train", + data_config: Mapping[str, Any] | None = None, + seed: int = 42, + max_batch_bytes: int = 512 * 1024 * 1024, + ) -> None: + from ._native.data.dataset import ( + MultiObjectDynamics, + _cached_numpy_load, + make_collate_fn, + ) + from ._native.data.scene_frames import clear_cache + + self.root = Path(root) + self.split_file, self.scale_file = Path(split_file), Path(object_scale_file) + self.config = deepcopy(dict(data_config or {})) + self.seed, self.max_batch_bytes = seed, max_batch_bytes + if split not in ("train", "val", "test") or max_batch_bytes < 1: + raise ValueError("Invalid split or batch memory limit") + for name in ("skip_step_size", "stride"): + if name in self.config and ( + type(self.config[name]) is not int or self.config[name] < 1 + ): + raise ValueError(f"{name} must be a positive integer") + if self.config.get("partial_view", {}).get("enable", False) or self.config.get( + "if_point_cloud", {} + ).get("enable", False): + raise ValueError("This recipe supports full mesh-vertex point inputs") + self.rank = ( + torch.distributed.get_rank() if torch.distributed.is_initialized() else 0 + ) + self.world_size = ( + torch.distributed.get_world_size() + if torch.distributed.is_initialized() + else 1 + ) + self.manifest = _verified_manifest(self.root) + # A previous source may have cached the old inode or per-frame arrays + # at these paths before the newly verified manifest was committed. + clear_cache(self.root) + _cached_numpy_load.cache_clear() + identity = { + "manifest": self.manifest, + "split": split, + "split_sha": hashlib.sha256(self.split_file.read_bytes()).hexdigest(), + "scale_sha": hashlib.sha256(self.scale_file.read_bytes()).hexdigest(), + "sampling": self.config, + "seed": seed, + } + self.source_id = ( + "rigidformer-movi-v1:" + + hashlib.sha256( + json.dumps(identity, sort_keys=True, allow_nan=False).encode() + ).hexdigest() + ) + options = { + k: self.config[k] + for k in ( + "sequence_length", + "stride", + "skip_first_frames", + "skip_step_size", + "random_step_size", + "augmentation", + "anchor_fps_start", + "fps_order", + "max_vertices_per_object", + "max_objects_per_frame", + ) + if k in self.config + } + self.dataset = MultiObjectDynamics( + str(self.root), + split=split, + split_file=str(self.split_file), + object_scale_file=str(self.scale_file), + anchor_k=self.config.get("anchor_k", 4), + if_phys_param=True, + compact_index=True, + require_prepared_cache=True, + **options, + ) + # The author loader can generate FPS caches. Prepared sources must + # already commit those caches, rather than acquire untracked inputs. + _check_manifest_coverage(self.root, self.manifest) + maximum = self.config.get("max_objects_per_frame") + if maximum is not None and any( + len(counts) > maximum + for counts in self.dataset._mesh_nv_by_dataset.values() + ): + raise ValueError( + "max_objects_per_frame cannot truncate physical scene objects" + ) + self.collate = make_collate_fn( + max_objects_per_frame=self.config.get("max_objects_per_frame", 16), + max_vertices_per_object=self.config.get("max_vertices_per_object"), + ) + self.epoch = 0 + self.cursor = 0 + self._pending: tuple[int, int] | None = None + self._order_epoch = -1 + self._order = torch.empty(0, dtype=torch.long) + self._per_rank = (len(self.dataset) + self.world_size - 1) // self.world_size + + def _verify_files(self) -> None: + if _verified_manifest(self.root) != self.manifest: + raise ValueError("Prepared geometry manifest differs") + + def _indices(self, epoch: int) -> torch.Tensor: + if self._order_epoch != epoch: + generator = torch.Generator().manual_seed(self.seed + epoch) + order = torch.randperm(len(self.dataset), generator=generator) + total = self._per_rank * self.world_size + if total > len(order): + padding = total - len(order) + repeats = (padding + len(order) - 1) // len(order) + order = torch.cat((order, order.repeat(repeats)[:padding])) + self._order = order[self.rank : total : self.world_size] + self._order_epoch = epoch + return self._order + + def updates_per_epoch(self, batch_size: int) -> int: + """Return full-batch updates per rank, matching drop-last training.""" + if type(batch_size) is not int or batch_size < 1 or self._per_rank < batch_size: + raise ValueError("Every rank needs at least one complete batch") + return self._per_rank // batch_size + + def sample(self, batch_size: int, *, generator: torch.Generator) -> SupervisedBatch: + """Stage one bounded batch; on_update commits its traversal position. + + Args: + batch_size: Number of windows on this rank. + generator: Caller-owned sampling/augmentation RNG. + + Returns: + Owned CPU shifted geometry supervision. + """ + if self._pending is not None: + raise RuntimeError("Previous sampled batch has not been committed") + self.updates_per_epoch(batch_size) + epoch, cursor = self.epoch, self.cursor + if cursor + batch_size > self._per_rank: + epoch, cursor = epoch + 1, 0 + indices = self._indices(epoch)[cursor : cursor + batch_size] + names = [self.dataset.all_sequences[int(i)]["dataset"] for i in indices] + objects = self.config.get("max_objects_per_frame", 16) + if objects is None: + objects = max(len(self.dataset._mesh_nv_by_dataset[name]) for name in names) + points = self.config.get("max_vertices_per_object") or max( + self.dataset._max_nv_by_dataset[name] for name in names + ) + frames = self.dataset.sequence_length + estimate = batch_size * ( + objects * (points * (frames * 13 + 24) + 9 * self.dataset.anchor_k + 128) + + 64 + ) + if estimate > self.max_batch_bytes: + raise ValueError( + "Padded geometry exceeds max_batch_bytes before reading payloads" + ) + py, npstate, ts = ( + random.getstate(), + np.random.get_state(), + torch.get_rng_state(), + ) + seed = int(torch.randint(0, 2**31, (1,), generator=generator)) + try: + random.seed(seed) + np.random.seed(seed) + torch.set_rng_state(torch.Generator().manual_seed(seed).get_state()) + values = [self.dataset[int(i)] for i in indices] + finally: + random.setstate(py) + np.random.set_state(npstate) + torch.set_rng_state(ts) + batch = self.collate(values) + self._pending = (epoch, cursor + batch_size) + return rigidformer_batch( + batch, + metadata={ + "source_id": self.source_id, + "scenes": batch["dataset_name"], + "frame_ids": batch["frame_ids"], + "physical_frame_rate": 240, + "model_dt": 1, + }, + ) + + def on_update(self, learner: Any) -> None: + """Commit the staged window traversal after a successful optimizer step.""" + if self._pending is None: + raise RuntimeError("No geometry batch is pending") + self.epoch, self.cursor = self._pending + self._pending = None + + def state_dict(self) -> dict[str, Any]: + """Return the committed rank-local sampling boundary.""" + return { + "source_id": self.source_id, + "epoch": self.epoch, + "cursor": self.cursor, + "rank": self.rank, + "world_size": self.world_size, + } + + def validate_state_dict(self, state: Mapping[str, Any]) -> None: + """Reject incompatible geometry or traversal before training mutation.""" + if ( + state.get("source_id") != self.source_id + or state.get("rank") != self.rank + or state.get("world_size") != self.world_size + or type(state.get("epoch")) is not int + or state["epoch"] < 0 + or type(state.get("cursor")) is not int + or not 0 <= state["cursor"] <= self._per_rank + ): + raise ValueError("MOVi sampling state differs") + self._verify_files() + + def load_state_dict(self, state: Mapping[str, Any]) -> None: + """Restore a prevalidated committed traversal and discard staged reads.""" + self.epoch, self.cursor = state["epoch"], state["cursor"] + self._pending = None + self._order_epoch = -1 diff --git a/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/ema.py b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/ema.py new file mode 100644 index 0000000000..5420bca366 --- /dev/null +++ b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/ema.py @@ -0,0 +1,92 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Author-equivalent EMA as a recoverable shared-trainer state.""" + +from __future__ import annotations +from collections.abc import Mapping +from typing import Any +import torch +from torch import nn + +__all__ = ["RigidFormerEMA"] + + +class RigidFormerEMA: + """Maintain the author's ramped FP32 EMA, including model buffers. + + Args: + model: Model whose parameter and buffer layout remains fixed. + decay: Maximum decay; initial updates use the author's warmup. + """ + + def __init__(self, model: nn.Module, decay: float = 0.9999) -> None: + if not 0 < decay < 1: + raise ValueError("EMA decay must be in (0,1)") + self.model, self.decay, self.num_updates = model, decay, 0 + self.shadow = { + k: ( + v.detach().float().clone() + if v.is_floating_point() + else v.detach().clone() + ) + for k, v in model.state_dict().items() + } + + @torch.no_grad() + def on_update(self, learner: Any) -> None: + """Update the average once after a successful parameter update.""" + self.num_updates += 1 + decay = min(self.decay, (1 + self.num_updates) / (10 + self.num_updates)) + floating, values = [], [] + for name, value in self.model.state_dict().items(): + if value.is_floating_point(): + floating.append(self.shadow[name]) + values.append(value.detach().float()) + else: + self.shadow[name].copy_(value) + torch._foreach_mul_(floating, decay) + torch._foreach_add_(floating, values, alpha=1 - decay) + + def state_dict(self) -> dict[str, Any]: + """Return owned FP32 shadows and the exact EMA update cursor.""" + return { + "decay": self.decay, + "num_updates": self.num_updates, + "shadow": {k: v.detach().cpu().clone() for k, v in self.shadow.items()}, + } + + def validate_state_dict(self, state: Mapping[str, Any]) -> None: + """Validate EMA policy and every shadow tensor before restoration.""" + if ( + state.get("decay") != self.decay + or type(state.get("num_updates")) is not int + or state["num_updates"] < 0 + or state["shadow"].keys() != self.shadow.keys() + or any( + v.shape != state["shadow"][k].shape + or v.dtype != state["shadow"][k].dtype + for k, v in self.shadow.items() + ) + ): + raise ValueError("RigidFormer EMA state differs") + + @torch.no_grad() + def load_state_dict(self, state: Mapping[str, Any]) -> None: + """Restore a prevalidated EMA without changing live model weights.""" + self.num_updates = state["num_updates"] + for key, value in self.shadow.items(): + value.copy_(state["shadow"][key]) diff --git a/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/sl_recipe.py b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/sl_recipe.py new file mode 100644 index 0000000000..a4da7857f4 --- /dev/null +++ b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/sl_recipe.py @@ -0,0 +1,263 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Installable author MOVi recipe using the shared supervised trainer.""" + +from __future__ import annotations +from collections.abc import Mapping +from copy import deepcopy +from typing import Any, TYPE_CHECKING +import torch +from .model import RigidFormerCfg, RigidFormerModel, RIGIDFORMER_REVISION +from .data import MOViSequenceSource +from .ema import RigidFormerEMA +from .supervised import RigidFormerObjective + +if TYPE_CHECKING: + from embodichain.learning.sl import SupervisedTrainer +__all__ = ["build_rigidformer_trainer"] + + +class _EpochSchedule: + def __init__( + self, + optimizer: torch.optim.Optimizer, + steps: int, + config: Mapping[str, Any], + max_epochs: int, + ) -> None: + self.steps, self.num_updates = steps, 0 + self.policy = { + "scheduler": dict(config), + "steps_per_epoch": steps, + "max_epochs": max_epochs, + } + warmup = int(config.get("warmup_epochs", 0)) + if config.get("type", "cosine") == "constant": + self.scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lambda _: 1.0) + elif config.get("type", "cosine") == "cosine": + # Match the author's construction order as well as the LR curve: + # Cosine records its initial rates after LinearLR's initialization. + linear = None + if warmup: + linear = torch.optim.lr_scheduler.LinearLR( + optimizer, + start_factor=float(config.get("start_factor", 0.1)), + total_iters=warmup, + ) + cosine = torch.optim.lr_scheduler.CosineAnnealingLR( + optimizer, + T_max=max(1, max_epochs - warmup), + eta_min=float(config.get("eta_min", config.get("min_lr", 1e-6))), + ) + if linear is not None: + self.scheduler = torch.optim.lr_scheduler.SequentialLR( + optimizer, [linear, cosine], milestones=[warmup] + ) + else: + self.scheduler = cosine + else: + raise ValueError( + "Only the released constant/cosine epoch schedules are supported" + ) + + def on_update(self, learner: Any) -> None: + self.num_updates = learner.num_updates + if self.num_updates % self.steps == 0: + self.scheduler.step() + + def state_dict(self) -> dict[str, Any]: + return { + "policy": deepcopy(self.policy), + "num_updates": self.num_updates, + "scheduler": deepcopy(self.scheduler.state_dict()), + } + + def validate_state_dict(self, state: Mapping[str, Any]) -> None: + if ( + state.get("policy") != self.policy + or type(state.get("num_updates")) is not int + or state["num_updates"] < 0 + ): + raise ValueError("Author epoch schedule differs") + + def load_state_dict(self, state: Mapping[str, Any]) -> None: + self.num_updates = state["num_updates"] + self.scheduler.load_state_dict(state["scheduler"]) + + +def _optimizer(model: RigidFormerModel, cfg: Mapping[str, Any]) -> torch.optim.AdamW: + if cfg.get("type", "AdamW").lower() != "adamw": + raise ValueError("The author recipe requires AdamW") + lr = float(cfg.get("lr", 1e-4)) + buckets: dict[float, list[torch.nn.Parameter]] = {1.0: []} + for name, param in model.network.named_parameters(): + if not param.requires_grad: + continue + scale = 1.0 + if any( + k in name + for k in ( + "predictor.fusion_norm1", + "predictor.fusion_w1", + "predictor.fusion_g1", + ) + ): + scale = float(cfg.get("fusion_lr_mult", 10.0)) + elif "shared_obj_encoder" in name: + scale = float(cfg.get("encoder_lr_mult", 1.0)) + elif name.startswith(("decoder.", "predictor.")): + scale = float(cfg.get("decoder_lr_mult", 1.0)) + buckets.setdefault(scale, []).append(param) + embeddings = { + id(p) for n, p in model.network.named_parameters() if "register_tokens" in n + } + groups = [] + for scale, params in buckets.items(): + if not params: + continue + if cfg.get("decay_matrices_only", True): + matrices = [p for p in params if p.ndim >= 2 and id(p) not in embeddings] + vectors = [p for p in params if p.ndim < 2 or id(p) in embeddings] + if matrices: + groups.append({"params": matrices, "lr": lr * scale}) + if vectors: + groups.append( + {"params": vectors, "lr": lr * scale, "weight_decay": 0.0} + ) + else: + groups.append({"params": params, "lr": lr * scale}) + return torch.optim.AdamW( + groups, + lr=lr, + betas=tuple(cfg.get("betas", (0.9, 0.999))), + eps=float(cfg.get("eps", 1e-8)), + weight_decay=float(cfg.get("weight_decay", 0.01)), + ) + + +def build_rigidformer_trainer(config: Mapping[str, Any]) -> SupervisedTrainer: + """Build an author-based MOVi recipe without importing example helpers. + + Args: + config: Model settings; prepared data root, split_file and scale_file; + optional author training overrides; common trainer configuration. + Relative paths resolve from the invoking working directory. + + Returns: + Shared trainer with committed sampling, epoch schedule and FP32 EMA. + """ + from embodichain.learning.sl import ( + SupervisedLearner, + SupervisedTrainer, + SupervisedTrainerCfg, + ) + from ._native.data.gpu_augmentation import GPUAugmentation + + model_cfg = RigidFormerCfg(**config.get("model", {})) + reference = model_cfg.reference_config() + training = {**reference["training"], **config.get("training", {})} + for key in ("optimizer", "scheduler", "loss_weights", "ema"): + training[key] = { + **reference["training"].get(key, {}), + **config.get("training", {}).get(key, {}), + } + if training.get("scheduled_sampling", {}).get("enabled", False): + raise ValueError( + "The released MOVi recipe uses teacher forcing; scheduled sampling is not part of this adapter" + ) + raw_data = dict(config["data"]) + sampling = {**reference["data"], **raw_data.pop("sampling", {})} + if sampling.get("anchor_k", 4) != reference["data"].get("anchor_k", 4): + raise ValueError("data.sampling.anchor_k must match the model architecture") + seed = int(config.get("trainer", {}).get("seed", 42)) + source = MOViSequenceSource( + raw_data.pop("root"), + raw_data.pop("split_file"), + raw_data.pop("scale_file"), + split=raw_data.pop("split", "train"), + data_config=sampling, + seed=seed, + **raw_data, + ) + kwargs = { + "device": "cuda", + "batch_size": int(training["train_batch_size"]), + **config.get("trainer", {}), + } + steps = source.updates_per_epoch(kwargs["batch_size"]) + kwargs.setdefault("total_updates", int(training["max_epochs"]) * steps) + trainer_cfg = SupervisedTrainerCfg(**kwargs) + torch.manual_seed(seed) + model = RigidFormerModel(model_cfg).to(trainer_cfg.device) + if source.world_size > 1: + model.network = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model.network) + optimizer = _optimizer(model, training["optimizer"]) + augmentation_cfg = deepcopy(reference) + augmentation_cfg["data"] = sampling + objective = RigidFormerObjective( + { + name: training["loss_weights"][name] + for name in ( + "position_before_icp", + "position_after_icp", + "acceleration_before_icp", + "acceleration_after_icp", + ) + }, + time_step_scale=training["loss_weights"].get("if_time_step_scale", True), + augmentation=GPUAugmentation( + augmentation_cfg, torch.device(trainer_cfg.device) + ), + ) + precision = config.get("precision", "bf16") + if precision not in ("bf16", "fp32"): + raise ValueError("precision must be bf16 or fp32") + learner = SupervisedLearner( + model, + objective, + optimizer, + max_grad_norm=float(training.get("gradient_clip", 1.0)), + amp_dtype=torch.bfloat16 if precision == "bf16" else None, + ddp_find_unused_parameters=True, + ddp_broadcast_buffers=True, + ) + schedule = _EpochSchedule( + optimizer, steps, training["scheduler"], int(training["max_epochs"]) + ) + states: dict[str, Any] = {"geometry_source": source, "epoch_schedule": schedule} + if training["ema"].get("enable", True): + states["ema"] = RigidFormerEMA( + model, float(training["ema"].get("decay", 0.9999)) + ) + return SupervisedTrainer( + learner, + source, + trainer_cfg, + signature={ + "family": "rigidformer-author-v1", + "author_revision": RIGIDFORMER_REVISION, + "architecture": reference["model"], + "reference_config": reference, + "sampling": sampling, + "training": training, + "precision": precision, + "attention_backend": model_cfg.attention_backend, + "steps_per_epoch": steps, + "ddp_broadcast_buffers": True, + }, + training_states=states, + ) diff --git a/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/supervised.py b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/supervised.py new file mode 100644 index 0000000000..d3a85b3b74 --- /dev/null +++ b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/supervised.py @@ -0,0 +1,188 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Shifted author supervision and valid-anchor loss reduction for unified SL.""" + +from __future__ import annotations +from collections.abc import Mapping +from typing import Any +import torch +from torch import nn +from torch.nn import functional as F +from embodichain.learning.sl import LossResult, SupervisedBatch + +__all__ = ["rigidformer_batch", "RigidFormerObjective"] + + +def rigidformer_batch( + values: Mapping[str, Any], *, metadata: Mapping[str, Any] | None = None +) -> SupervisedBatch: + """Separate shifted conditioning history and next-frame supervision. + + Args: + values: Author batch including an aligned sequence of at least 3 frames. + metadata: Optional immutable data and sampling provenance. + + Returns: + Supervision for frames 2..T-1 with causal conditioning through T-2. + Earlier targets also occur as subsequent teacher-forcing inputs. + """ + positions = values["obj_positions"] + if positions.ndim != 5 or positions.shape[1] < 3: + raise ValueError("Supervision requires obj_positions [B,T>=3,O,P,3]") + inputs = {k: v for k, v in values.items() if isinstance(v, torch.Tensor)} + inputs["obj_positions"] = positions[:, :-1] + inputs["valid_vtx_mask"] = values["valid_vtx_mask"][:, :-1] + return SupervisedBatch( + inputs=inputs, targets=positions[:, 2:], metadata=dict(metadata or {}) + ) + + +class RigidFormerObjective: + """Reproduce the author's four anchor-only smooth-L1 losses. + + Args: + weights: Position/acceleration before/after Kabsch weights. + time_step_scale: Divide errors by step_size squared before smooth-L1. + """ + + def __init__( + self, + weights: Mapping[str, float] | None = None, + *, + time_step_scale: bool = True, + augmentation: Any = None, + ) -> None: + self.weights = { + "position_before_icp": 10.0, + "position_after_icp": 10.0, + "acceleration_before_icp": 1.0, + "acceleration_after_icp": 1.0, + **dict(weights or {}), + } + if set(self.weights) != { + "position_before_icp", + "position_after_icp", + "acceleration_before_icp", + "acceleration_after_icp", + } or any( + not torch.isfinite(torch.tensor(v)) or v < 0 for v in self.weights.values() + ): + raise ValueError( + "Only the four finite nonnegative author anchor loss weights are supported" + ) + self.time_step_scale = time_step_scale + self.augmentation = augmentation + + def compute_loss(self, model: nn.Module, batch: SupervisedBatch) -> LossResult: + """Compute one local anchor mean; the shared learner owns DDP reduction. + + Args: + model: Native RigidFormer model. + batch: Shifted conditioning and owned target geometry. + + Returns: + Weighted scalar loss, diagnostic terms and valid-anchor count. + """ + values = dict(batch.inputs) + history = values["obj_positions"] + target = batch.targets + if self.augmentation is not None: + raw = dict(values) + raw["obj_positions"] = torch.cat((history[:, :2], target), dim=1) + raw["valid_vtx_mask"] = torch.cat( + (values["valid_vtx_mask"], values["valid_vtx_mask"][:, -1:]), dim=1 + ) + # The reference augments FP32 geometry before entering model autocast. + with torch.autocast(history.device.type, enabled=False): + augmented = rigidformer_batch(self.augmentation(raw, split="train")) + values = dict(augmented.inputs) + history, target = values["obj_positions"], augmented.targets + if target.shape != history[:, 1:].shape or batch.masks: + raise ValueError( + "Target layout must align with shifted author history; masks use geometry fields" + ) + # The last frame is never a conditioning input. Fill it without labels. + values["obj_positions"] = torch.cat( + (history, torch.zeros_like(history[:, -1:])), dim=1 + ) + values["valid_vtx_mask"] = torch.cat( + (values["valid_vtx_mask"], values["valid_vtx_mask"][:, -1:]), dim=1 + ) + # The native network centers its own input in-place when configured; + # its position outputs use that centered frame, while target is owned + # separately and remains in world coordinates. + reference = getattr(model, "reference_config", {}) + centered = reference.get("data", {}).get("center_scene_by_t0", False) + centers = None + if centered: + positions = values.get("obj_positions_input", values["obj_positions"]) + centers = torch.stack( + [ + positions[i, 0][values["valid_vtx_mask"][i, 0]].mean(dim=0) + for i in range(positions.shape[0]) + ] + ) + predictions = model(values) + b, t, o, p, _ = target.shape + anchors = values["anchor_idx"][:, None, :, :, None].expand(b, t, o, -1, 3) + + def at_anchor(value: torch.Tensor) -> torch.Tensor: + return torch.gather(value, 3, anchors) + + gt_pos = at_anchor(target) + if centers is not None: + gt_pos = gt_pos - centers[:, None, None, None, :] + gt_acc = at_anchor(target - 2 * history[:, 1:] + history[:, :-1]) + errors = { + "position_before_icp": at_anchor(predictions["obj_positions_norm"]) + - gt_pos, + "position_after_icp": at_anchor(predictions["obj_positions_icp_norm"]) + - gt_pos, + "acceleration_before_icp": predictions[ + "obj_accelerations_no_icp_norm_anchor" + ] + - gt_acc, + "acceleration_after_icp": at_anchor(predictions["obj_accelerations_norm"]) + - gt_acc, + } + mask = ( + values["valid_obj_mask"][:, None, :, None] + .expand(b, t, o, anchors.shape[3]) + .float() + ) + count = mask.sum() + if count <= 0: + raise ValueError("At least one valid supervised anchor is required") + scale = ( + values["skip_step_size"][:, None, None, None, None].float().square() + if self.time_step_scale + else 1.0 + ) + terms = { + name: ( + F.smooth_l1_loss( + error / scale, torch.zeros_like(error), reduction="none" + ) + .sum(-1) + .float() + * mask + ).sum() + / (count + 1e-8) + for name, error in errors.items() + } + loss = sum(terms[name] * weight for name, weight in self.weights.items()) + return LossResult(loss=loss, metrics=terms, weight=float(count.detach())) diff --git a/pyproject.toml b/pyproject.toml index e0d6695510..1bc0109d39 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -60,6 +60,8 @@ dependencies = [ rigidformer = [ "einops>=0.8.1", "roma==1.5.4", + "pyarrow>=15", + "trimesh>=4.9", ] policy-deploy = [ # ORT 1.27+ defaults to CUDA 13; EmbodiChain's supported stack uses CUDA 12. diff --git a/scripts/benchmark/learning/rigidformer_index.py b/scripts/benchmark/learning/rigidformer_index.py new file mode 100644 index 0000000000..6030fdecb4 --- /dev/null +++ b/scripts/benchmark/learning/rigidformer_index.py @@ -0,0 +1,205 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Measure author MOVi index construction memory and latency. + +Run: python -m scripts.benchmark.learning.rigidformer_index --output /tmp/index.json +Scope: CPU metadata construction, excluding data IO, model work and imports. +""" + +from __future__ import annotations +import argparse +import gc +import json +import os +from pathlib import Path +import platform +import statistics +import time +import tracemalloc +from typing import Any +from embodichain_models.dynamics.rigidformer._native.data._index import _SceneWindows + + +def _build(scenes: int, compact: bool) -> Any: + """Construct the released random-sampling scene mapping or its list baseline.""" + result = _SceneWindows(random_sampling=True) if compact else [] + for i in range(scenes): + name = f"scene_{i}" + frames = [Path(name) / f"frame_{j:05d}" for j in range(2, 480)] + if compact: + result.add_scene(name, frames, 8, 1) + else: + for start in range(len(frames) - 8 + 1): + selected = frames[start : start + 8] + result.append( + { + "dataset": name, + "frames": [p.as_posix() for p in selected], + "frame_ids": [int(p.name.split("_")[-1]) for p in selected], + } + ) + return result + + +def _report(path: Path, record: dict[str, Any]) -> None: + """Write one Markdown report alongside the machine-readable measurements.""" + rows = record["results"] + lines = [ + "# RigidFormer random-window index", + "", + record["scope"], + "", + f"Python {record['python']}; CPU {record['cpu']}; one Python thread; " + f"{record['repeats']} warmed repeats. No tensor dtype/device is involved.", + "", + "Memory is a separate construction with tracemalloc: Python allocations only, " + "not process RSS or total peak memory. GPU metrics are inapplicable.", + "", + ] + tables = [ + ( + "Time & Memory", + [ + "scenes", + "implementation", + "cost_time_ms", + "cpu_delta_mb", + "gpu_delta_mb", + "peak_gpu_mb", + "retained_python_mb", + "peak_python_mb", + ], + [ + [ + r["scenes"], + r["implementation"], + r["median_seconds"] * 1000, + "unmeasured", + "N/A", + "N/A", + r["retained_python_bytes"] / 2**20, + r["peak_python_bytes"] / 2**20, + ] + for r in rows + ], + ), + ( + "Success & Other Metrics", + ["scenes", "implementation", "windows", "success_rate"], + [[r["scenes"], r["implementation"], r["windows"], 1.0] for r in rows], + ), + ( + "Leaderboard", + ["scenes", "rank", "implementation", "success_rate", "cost_time_ms"], + [ + [scenes, i + 1, r["implementation"], 1.0, r["median_seconds"] * 1000] + for scenes in sorted({r["scenes"] for r in rows}) + for i, r in enumerate( + sorted( + (r for r in rows if r["scenes"] == scenes), + key=lambda r: r["median_seconds"], + ) + ) + ], + ), + ] + for title, columns, values in tables: + lines += [ + f"## {title}", + "", + "| " + " | ".join(columns) + " |", + "| " + " | ".join("---" for _ in columns) + " |", + ] + lines += [ + "| " + + " | ".join(f"{v:.6g}" if isinstance(v, float) else str(v) for v in row) + + " |" + for row in values + ] + lines += [""] + path.write_text("\n".join(lines)) + + +def main() -> None: + """Measure equivalent window counts and scene selection at three scales.""" + args = argparse.ArgumentParser(description=__doc__) + args.add_argument("--scenes", type=int, nargs="+", default=[1, 32, 960]) + args.add_argument("--repeats", type=int, default=3) + args.add_argument("--output", type=Path, required=True) + cfg = args.parse_args() + if cfg.repeats < 1 or any(s < 1 for s in cfg.scenes): + args.error("scenes and repeats must be positive") + results = [] + for scenes in cfg.scenes: + for compact in (False, True): + warm = _build(min(scenes, 2), compact) + del warm + samples = [] + for _ in range(cfg.repeats): + gc.collect() + start = time.perf_counter() + result = _build(scenes, compact) + elapsed = time.perf_counter() - start + assert len(result) == scenes * 471 + for idx in (0, len(result) // 2, len(result) - 1): + assert result[idx]["dataset"] == f"scene_{idx//471}" + samples.append(elapsed) + del result + gc.collect() + tracemalloc.start() + result = _build(scenes, compact) + retained, peak = tracemalloc.get_traced_memory() + tracemalloc.stop() + del result + results.append( + { + "scenes": scenes, + "windows": scenes * 471, + "implementation": "compact" if compact else "author-list", + "median_seconds": statistics.median(samples), + "samples_seconds": samples, + "retained_python_bytes": retained, + "peak_python_bytes": peak, + } + ) + record = { + "python": platform.python_version(), + "cpu": platform.processor(), + "cpu_count": os.cpu_count(), + "python_threads": 1, + "omp_num_threads": os.environ.get("OMP_NUM_THREADS"), + "seed": None, + "inputs": "deterministic 480-frame scene paths; no random payloads", + "source_sha256": __import__("hashlib") + .sha256(Path(__file__).read_bytes()) + .hexdigest(), + "scope": "CPU index construction only; no IO/model/GPU; single process", + "frames": 480, + "skip_first": 2, + "sequence_length": 8, + "repeats": cfg.repeats, + "results": results, + } + cfg.output.parent.mkdir(parents=True, exist_ok=True) + cfg.output.write_text(json.dumps(record, indent=2)) + _report(cfg.output.with_suffix(".md"), record) + print(json.dumps(record, indent=2)) + print(f"Markdown report saved: {cfg.output.with_suffix('.md')}") + + +if __name__ == "__main__": + main() diff --git a/tests/learning/sl/test_distributed.py b/tests/learning/sl/test_distributed.py index a3b0d62e17..550d95ccb9 100644 --- a/tests/learning/sl/test_distributed.py +++ b/tests/learning/sl/test_distributed.py @@ -166,6 +166,30 @@ def sample(self, batch_size: int, *, generator: torch.Generator) -> SupervisedBa return SupervisedBatch(batch.inputs + offset, batch.targets) +class _RankUpdateState: + def __init__(self, rank: int) -> None: + self.rank, self.updates, self.total = rank, 0, torch.zeros(()) + + def on_update(self, learner: SupervisedLearner) -> None: + self.updates += 1 + self.total += next(learner.model.parameters()).detach().sum() + self.rank + + def state_dict(self) -> dict[str, Any]: + return {"rank": self.rank, "updates": self.updates, "total": self.total.clone()} + + def validate_state_dict(self, state: dict[str, Any]) -> None: + if ( + state["rank"] != self.rank + or type(state["updates"]) is not int + or state["updates"] < 0 + ): + raise ValueError("Rank-local update state differs") + + def load_state_dict(self, state: dict[str, Any]) -> None: + self.updates = state["updates"] + self.total.copy_(state["total"]) + + def _build_recovery(path: Path, rank: int) -> SupervisedTrainer: random.seed(101 + rank) np.random.seed(201 + rank) @@ -187,6 +211,7 @@ def _build_recovery(path: Path, rank: int) -> SupervisedTrainer: _StochasticSource(), cfg, signature={"regression": "rank-local recovery"}, + training_states={"rank_update": _RankUpdateState(rank)}, ) def evaluate(current: SupervisedLearner) -> dict[str, float]: @@ -197,6 +222,7 @@ def evaluate(current: SupervisedLearner) -> dict[str, float]: np.random.random(3) torch.rand(5) torch.randint(10, (3,), generator=trainer.sampler) + trainer.training_states["rank_update"].updates += 100 # BC-style checkpoint publication inside a rank-zero callback must not hang. trainer.save(path / "selected.pt") with (path.parent / f"{path.name}-validation-rank{rank}.jsonl").open( @@ -213,6 +239,10 @@ def _snapshot(trainer: SupervisedTrainer) -> dict[str, Any]: return copy.deepcopy( { "learner": trainer.learner.state_dict(), + "training_states": { + name: value.state_dict() + for name, value in trainer.training_states.items() + }, "sampler": trainer.sampler.get_state(), "rng": { "python": random.getstate(), @@ -419,15 +449,34 @@ def _source_resume_worker(rank: int, directory: Path) -> None: dist.barrier() +def _chunked_checkpoint_worker(rank: int) -> None: + from embodichain.learning.sl.distributed import _DistributedContext + + context = _DistributedContext("cpu") + # Slightly over 16 MiB, so this exercises multiple transfer chunks. + value = {"rank": rank, "tensor": torch.arange(4_194_307, dtype=torch.float32)} + states = context.gather(value) + if rank == 0: + assert states is not None + assert [state["rank"] for state in states] == list(range(dist.get_world_size())) + for state in states[1:]: + torch.testing.assert_close(state["tensor"], value["tensor"], rtol=0, atol=0) + else: + assert states is None + + with pytest.raises(RuntimeError, match="checkpoint state serialization"): + context.gather((lambda: None) if rank == 1 else value) + + def _distributed_worker( - rank: int, group_file: str, directory: str, scenario: str + rank: int, group_file: str, directory: str, scenario: str, world_size: int ) -> None: torch.set_num_threads(1) dist.init_process_group( "gloo", init_method=f"file://{group_file}", rank=rank, - world_size=2, + world_size=world_size, timeout=timedelta(seconds=45), ) try: @@ -439,17 +488,24 @@ def _distributed_worker( _artifact_worker(rank, Path(directory)) elif scenario == "source-resume": _source_resume_worker(rank, Path(directory)) + elif scenario == "chunked-checkpoint": + _chunked_checkpoint_worker(rank) else: _recovery_worker(rank, Path(directory), scenario) finally: dist.destroy_process_group() -def _spawn(scenario: str, directory: Path) -> None: +def _spawn(scenario: str, directory: Path, world_size: int = 2) -> None: context = mp.spawn( _distributed_worker, - args=(str(directory / f"group-{scenario}"), str(directory), scenario), - nprocs=2, + args=( + str(directory / f"group-{scenario}"), + str(directory), + scenario, + world_size, + ), + nprocs=world_size, join=False, ) deadline = time.monotonic() + 100 @@ -469,6 +525,10 @@ def test_unequal_weight_single_and_multistep_match_global_batch(tmp_path: Path) _spawn("weighted", tmp_path) +def test_chunked_checkpoint_gathers_and_propagates_failures(tmp_path: Path) -> None: + _spawn("chunked-checkpoint", tmp_path, world_size=3) + + def test_restart_preserves_rank_rng_sampling_buffers_and_single_writer( tmp_path: Path, ) -> None: diff --git a/tests/learning/sl/test_training_states.py b/tests/learning/sl/test_training_states.py new file mode 100644 index 0000000000..cc79636606 --- /dev/null +++ b/tests/learning/sl/test_training_states.py @@ -0,0 +1,111 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Commit and restore explicit update participants independently of objectives.""" + +from __future__ import annotations +import copy +import torch +import pytest +from .test_trainer import build + + +class UpdateAverage: + """Independent running-weight average exercising trainer state recovery.""" + + def __init__(self): + self.count = 0 + self.total = torch.zeros(()) + + def on_update(self, learner): + self.count += 1 + self.total += next(learner.model.parameters()).detach().cpu().sum() + + def state_dict(self): + return {"count": self.count, "total": self.total.clone()} + + def validate_state_dict(self, state): + if ( + type(state["count"]) is not int + or state["count"] < 0 + or state["total"].shape != () + ): + raise ValueError("Invalid average state") + + def load_state_dict(self, state): + self.count = state["count"] + self.total.copy_(state["total"]) + + +def test_update_participant_restores_exact_committed_average(tmp_path): + full = build(tmp_path / "state-full") + full.training_states["average"] = UpdateAverage() + full.fit() + split = build(tmp_path / "state-split") + split.training_states["average"] = UpdateAverage() + split.fit(stop_after=7) + saved = torch.load(tmp_path / "state-split/last.pt", weights_only=False) + assert saved["version"] == 3 + resumed = build(tmp_path / "state-resume") + resumed.training_states["average"] = UpdateAverage() + resumed.resume(tmp_path / "state-split/last.pt") + assert resumed.training_states["average"].count == 7 + resumed.fit() + assert resumed.training_states["average"].count == 20 + torch.testing.assert_close( + resumed.training_states["average"].total, + full.training_states["average"].total, + atol=0, + rtol=0, + ) + for key, value in full.learner.model.state_dict().items(): + torch.testing.assert_close( + resumed.learner.model.state_dict()[key], value, atol=0, rtol=0 + ) + + +def test_participant_layout_rejected_before_model_mutation(tmp_path): + source = build(tmp_path / "participant-source") + source.training_states["average"] = UpdateAverage() + source.fit(stop_after=2) + target = build(tmp_path / "participant-target") + target.training_states["another_average"] = UpdateAverage() + before = copy.deepcopy(target.learner.model.state_dict()) + with pytest.raises(ValueError, match="training state names"): + target.resume(tmp_path / "participant-source/last.pt") + assert target.learner.num_updates == 0 + for key, value in before.items(): + torch.testing.assert_close( + target.learner.model.state_dict()[key], value, atol=0, rtol=0 + ) + + +def test_default_recipe_resumes_checkpoint_without_new_optional_fields(tmp_path): + source = build(tmp_path / "legacy-source") + source.fit(stop_after=2) + state = torch.load(tmp_path / "legacy-source/last.pt", weights_only=False) + assert state["version"] == 1 + state.pop("training_states") + state.pop("ddp_broadcast_buffers") + path = tmp_path / "legacy.pt" + torch.save(state, path) + target = build(tmp_path / "legacy-target") + target.resume(path) + assert target.learner.num_updates == 2 + for key, value in source.learner.model.state_dict().items(): + torch.testing.assert_close( + target.learner.model.state_dict()[key], value, atol=0, rtol=0 + ) diff --git a/tests/models/dynamics/_rigidformer_training.py b/tests/models/dynamics/_rigidformer_training.py new file mode 100644 index 0000000000..af2899ef1b --- /dev/null +++ b/tests/models/dynamics/_rigidformer_training.py @@ -0,0 +1,170 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +import hashlib +import json +from pathlib import Path +import numpy as np +from embodichain_models.dynamics.rigidformer.data import MOViSequenceSource +from ._rigidformer_fixtures import tiny_cfg + + +def make_packed(root: Path, frames_count: int = 120) -> Path: + folder = root / "scene_0" + folder.mkdir(parents=True) + (folder / "anchors").mkdir() + (folder / "meshes").mkdir() + shape = np.array( + [[0, 0, 0], [0.1, 0, 0], [0, 0.1, 0], [0, 0, 0.1]], dtype=np.float32 + ) + frames = np.stack( + [ + np.concatenate( + [shape + [t * 0.001, 0, 1.0], shape + [t * 0.001, 1.0, 1.0]], axis=0 + ) + for t in range(frames_count) + ] + ).astype(np.float32) + np.save(folder / "frames.npy", frames) + np.save(folder / "frames_counts.npy", np.array([4, 4], dtype=np.int64)) + traj = np.zeros((1, frames_count, 2, 7), dtype=np.float32) + traj[0, :, :, 3] = 1 + for t in range(frames_count): + traj[0, t, :, :3] = [[t * 0.001, 0, 1.0], [t * 0.001, 1.0, 1.0]] + np.save(folder / "vis_traj.npy", traj) + physics = { + f"obj_{i}": {"mass": 1.0, "friction": 0.8, "restitution": 0.3, "size": size} + for i, size in enumerate((0.7, 1.4)) + } + (folder / "physics_params.json").write_text(json.dumps(physics)) + for i in range(2): + np.save( + folder / "meshes" / f"obj_{i}_faces.npy", + np.array([[0, 1, 2], [0, 1, 3]], dtype=np.int32), + ) + # All four tetrahedron vertices are anchors; the stored order is an + # explicit fixture input, included in the manifest like converted data. + order = np.array([1, 2, 3, 0], dtype=np.int32) + np.save(folder / "anchors" / f"obj_{i}_fpsorder_cf.npy", order) + np.save(folder / "anchors" / f"obj_{i}_anchors_fps4_cf.npy", order) + splits = { + "scene_set": "complete", + "train_datasets": ["scene_0"], + "val_datasets": [], + "test_datasets": [], + } + (root / "splits.json").write_text(json.dumps(splits)) + p0 = traj[0, 0, :, :3].astype(np.float64) + radius = [ + float( + np.linalg.norm( + frames[0, i * 4 : (i + 1) * 4].astype(np.float64) - p0[i], axis=1 + ).max() + ) + for i in range(2) + ] + (root / "scales.json").write_text( + json.dumps( + { + "converted": root.name, + "scales": {"scene_0": [0.7, 1.4]}, + "fingerprint": {"scene_0": p0[0].tolist()}, + "radius0": {"scene_0": radius}, + } + ) + ) + manifest = { + "files": { + str(p.relative_to(root)): hashlib.sha256(p.read_bytes()).hexdigest() + for p in root.rglob("*") + if p.is_file() + } + } + (root / "manifest.json").write_text(json.dumps(manifest)) + return root + + +def source(root, *, random_step=True, **kwargs): + sampling = { + "sequence_length": 8, + "skip_first_frames": 2, + "random_step_size": { + "enable": random_step, + "step_sizes": [1, 5, 10], + "random_start_frame": True, + "sampling_weights": [0.34, 0.33, 0.33], + }, + "anchor_k": 4, + "anchor_fps_start": "centroid_farthest", + "fps_order": True, + "max_objects_per_frame": 2, + "augmentation": {}, + } + return MOViSequenceSource( + root, root / "splits.json", root / "scales.json", data_config=sampling, **kwargs + ) + + +def recovery_config(root: Path, output: Path, device: str) -> dict: + """Use four windows per scene to cross epochs and warmup in six updates.""" + architecture = tiny_cfg().architecture + return { + "model": {"architecture": architecture}, + "data": { + "root": str(root), + "split_file": str(root / "splits.json"), + "scale_file": str(root / "scales.json"), + "sampling": { + "sequence_length": 8, + "skip_first_frames": 2, + "max_objects_per_frame": 2, + "random_step_size": {"enable": False}, + "augmentation": {}, + }, + }, + "training": {"max_epochs": 3, "scheduler": {"warmup_epochs": 2}}, + "precision": "fp32", + "trainer": { + "device": device, + "batch_size": 1, + "total_updates": 6, + "seed": 42, + "checkpoint_interval": 0, + "output_dir": str(output), + }, + } + + +def assert_exact(actual, expected) -> None: + """Check exact resumed tensor, optimizer, RNG and participant state.""" + import torch + + if isinstance(actual, torch.Tensor): + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + elif isinstance(actual, np.ndarray): + np.testing.assert_array_equal(actual, expected) + elif isinstance(actual, dict): + assert actual.keys() == expected.keys() + for key in actual: + assert_exact(actual[key], expected[key]) + elif isinstance(actual, (tuple, list)): + assert len(actual) == len(expected) + for a, b in zip(actual, expected): + assert_exact(a, b) + else: + assert actual == expected diff --git a/tests/models/dynamics/test_rigidformer_distributed.py b/tests/models/dynamics/test_rigidformer_distributed.py new file mode 100644 index 0000000000..725101f999 --- /dev/null +++ b/tests/models/dynamics/test_rigidformer_distributed.py @@ -0,0 +1,280 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Real NCCL/SyncBatchNorm weighted updates and model-recipe recovery.""" + +from __future__ import annotations + +from datetime import timedelta +import ctypes +import faulthandler +import json +from pathlib import Path +import time +import traceback + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp + +from embodichain.learning.sl import SupervisedLearner +from embodichain_models.dynamics.rigidformer import RigidFormerModel +from embodichain_models.dynamics.rigidformer.sl_recipe import build_rigidformer_trainer +from embodichain_models.dynamics.rigidformer.supervised import ( + RigidFormerObjective, + rigidformer_batch, +) +from ._rigidformer_fixtures import tiny_cfg, geometry_batch +from ._rigidformer_training import assert_exact, make_packed, recovery_config + +pytestmark = [pytest.mark.gpu, pytest.mark.slow] + + +def _chunked_checkpoint_worker(rank: int, path: str) -> None: + from embodichain.learning.sl.distributed import _DistributedContext + + torch.cuda.set_device(rank) + dist.init_process_group( + "nccl", + init_method=(Path(path) / "chunked-rendezvous").as_uri(), + rank=rank, + world_size=2, + timeout=timedelta(seconds=90), + ) + try: + context = _DistributedContext(f"cuda:{rank}") + context.gather({"warmup": rank}) + values = {"rank": rank, "tensor": torch.arange(4_194_307, dtype=torch.float32)} + torch.cuda.synchronize() + baseline = torch.cuda.memory_allocated() + torch.cuda.reset_peak_memory_stats() + gathered = context.gather(values) + torch.cuda.synchronize() + peak_delta = torch.cuda.max_memory_allocated() - baseline + # The rank state exceeds one 16 MiB transfer chunk. + assert peak_delta < 20 * 1024 * 1024 + if rank == 0: + assert [state["rank"] for state in gathered] == [0, 1] + torch.testing.assert_close(gathered[1]["tensor"], values["tensor"]) + else: + assert gathered is None + finally: + dist.destroy_process_group() + + +@pytest.mark.skipif( + torch.cuda.device_count() < 2 or not dist.is_nccl_available(), + reason="Requires two CUDA devices and NCCL", +) +def test_nccl_chunked_checkpoint_limits_device_memory(tmp_path): + context = mp.spawn( + _chunked_checkpoint_worker, args=(str(tmp_path),), nprocs=2, join=False + ) + deadline = time.monotonic() + 120 + try: + while not context.join(timeout=1): + if time.monotonic() > deadline: + raise TimeoutError("NCCL checkpoint transfer exceeded 120 seconds") + finally: + for process in context.processes: + if process.is_alive(): + process.kill() + process.join() + + +def _disable_dropout(model: torch.nn.Module) -> None: + # Encoder regularization is hard-coded in author code. Disable stochastic + # masks for the concatenated-batch oracle; the recovery case keeps them. + for module in model.modules(): + if isinstance(module, torch.nn.Dropout): + module.p = 0.0 + + +def _weighted_batch(device: str) -> dict: + values = geometry_batch(device) + for key in ("obj_positions", "valid_vtx_mask"): + values[key] = torch.cat((values[key], values[key][:, :, :1].clone()), dim=2) + for key in ("ref_obj_positions", "anchor_idx", "phys_params", "vtx_rank"): + values[key] = torch.cat((values[key], values[key][:, :1].clone()), dim=1) + values["obj_positions"][:, :, 2, :, 0] += 1 + values["ref_obj_positions"][:, 2, :, 0] += 1 + values["valid_obj_mask"] = torch.tensor( + [[True, True, True], [True, True, False]], device=device + ) + return values + + +def _oracle(path: Path) -> None: + # Compute outside any process group. The author's grouped encoder otherwise + # all-reduces its moments even with ordinary BatchNorm modules. + torch.manual_seed(73) + reference = RigidFormerModel(tiny_cfg()).cuda().train() + _disable_dropout(reference) + initial = {n: v.detach().cpu().clone() for n, v in reference.state_dict().items()} + optimizer = torch.optim.SGD(reference.parameters(), lr=0.001) + expected = RigidFormerObjective().compute_loss( + reference, rigidformer_batch(_weighted_batch("cuda:0")) + ) + expected.loss.backward() + torch.nn.utils.clip_grad_norm_(reference.parameters(), 1.0) + optimizer.step() + torch.save( + { + "initial": initial, + "loss": float(expected.loss), + "model": {n: v.detach().cpu() for n, v in reference.state_dict().items()}, + "gradients": { + n: None if v.grad is None else v.grad.detach().cpu() + for n, v in reference.named_parameters() + }, + }, + path / "oracle.pt", + ) + + +def _weighted_update(rank: int, path: Path) -> dict: + torch.manual_seed(73) + model = RigidFormerModel(tiny_cfg()).cuda(rank) + _disable_dropout(model) + model.network = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model.network) + expected = torch.load( + path / "oracle.pt", weights_only=True, map_location=f"cuda:{rank}" + ) + for name, value in model.state_dict().items(): + torch.testing.assert_close(value, expected["initial"][name], rtol=0, atol=0) + values = _weighted_batch(f"cuda:{rank}") + learner = SupervisedLearner( + model, + RigidFormerObjective(), + torch.optim.SGD(model.parameters(), lr=0.001), + max_grad_norm=1.0, + ddp_find_unused_parameters=True, + ddp_broadcast_buffers=True, + ) + local = {key: value[rank : rank + 1] for key, value in values.items()} + metrics = learner.update(rigidformer_batch(local)) + assert metrics["loss"] == pytest.approx(expected["loss"], rel=1e-4, abs=1e-5), ( + metrics["loss"], + expected["loss"], + ) + error = 0.0 + for name, value in model.state_dict().items(): + target = expected["model"][name] + torch.testing.assert_close(value, target, rtol=1e-4, atol=1e-5) + if value.numel(): + error = max(error, float((value - target).abs().max())) + for name, actual in model.named_parameters(): + target = expected["gradients"][name] + if actual.grad is None: + assert target is None + else: + torch.testing.assert_close(actual.grad, target, rtol=1e-4, atol=1e-5) + return { + "global_loss": metrics["loss"], + "state_max_abs": error, + "valid_objects_by_rank": [3, 2], + } + + +def _worker(rank: int, path: str) -> None: + faulthandler.dump_traceback_later(30, repeat=True) + torch.set_num_threads(1) + torch.cuda.set_device(rank) + root = Path(path) + dist.init_process_group( + "nccl", + init_method=(root / "rendezvous").as_uri(), + rank=rank, + world_size=2, + timeout=timedelta(seconds=90), + ) + try: + weighted = _weighted_update(rank, root) + config = recovery_config(root / "data", root / "full", f"cuda:{rank}") + full = build_rigidformer_trainer(config) + assert full.signature["steps_per_epoch"] == 2 + full.fit() + config["trainer"]["output_dir"] = str(root / "split") + split = build_rigidformer_trainer(config) + split.fit(stop_after=3) + config["trainer"]["output_dir"] = str(root / "resumed") + resumed = build_rigidformer_trainer(config) + resumed.resume(root / "split/last.pt") + resumed.fit() + a = torch.load(root / "full/last.pt", weights_only=False) + b = torch.load(root / "resumed/last.pt", weights_only=False) + for key in ("learner", "training_states", "sampler", "rng", "distributed"): + assert_exact(a[key], b[key]) + assert resumed.training_states["epoch_schedule"].scheduler.last_epoch == 3 + if rank == 0: + runtime = ctypes.c_int() + ctypes.CDLL("libnccl.so.2").ncclGetVersion(ctypes.byref(runtime)) + (root / "summary.json").write_text( + json.dumps( + { + "backend": "nccl", + "world_size": 2, + "sync_batch_norm": True, + "loaded_nccl_version": runtime.value, + "weighted_update": weighted, + "recovery": "exact", + "updates": 6, + "split_after": 3, + "epochs": 3, + "warmup_epochs": 2, + "compared": [ + "model", + "optimizer", + "EMA", + "schedule", + "rank RNG", + "rank source cursor", + "rank buffers", + ], + }, + indent=2, + ) + ) + except BaseException: + traceback.print_exc() + raise + finally: + faulthandler.cancel_dump_traceback_later() + dist.destroy_process_group() + + +@pytest.mark.skipif( + torch.cuda.device_count() < 2 or not dist.is_nccl_available(), + reason="Requires two CUDA devices and NCCL", +) +def test_nccl_weighted_syncbn_and_exact_rank_recipe_recovery(tmp_path): + make_packed(tmp_path / "data", frames_count=13) + _oracle(tmp_path) + context = mp.spawn(_worker, args=(str(tmp_path),), nprocs=2, join=False) + deadline = time.monotonic() + 120 + try: + while not context.join(timeout=1): + if time.monotonic() > deadline: + raise TimeoutError( + "Two-rank RigidFormer qualification exceeded 120 seconds" + ) + finally: + for process in context.processes: + if process.is_alive(): + process.kill() + process.join() + assert json.loads((tmp_path / "summary.json").read_text())["recovery"] == "exact" diff --git a/tests/models/dynamics/test_rigidformer_training.py b/tests/models/dynamics/test_rigidformer_training.py new file mode 100644 index 0000000000..32a388c982 --- /dev/null +++ b/tests/models/dynamics/test_rigidformer_training.py @@ -0,0 +1,496 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +import hashlib +import json +import os +import numpy as np +import pytest +import torch +from embodichain_models.dynamics.rigidformer.supervised import ( + RigidFormerObjective, + rigidformer_batch, +) +from ._rigidformer_fixtures import tiny_cfg, geometry_batch +from ._rigidformer_training import make_packed, source, recovery_config, assert_exact + + +@pytest.fixture +def packed_data(tmp_path): + return make_packed(tmp_path / "ToyData") + + +def test_shifted_objective_matches_independent_anchor_huber_oracle(): + v = geometry_batch() + + class Predictions(torch.nn.Module): + def __init__(self): + super().__init__() + self.offset = torch.nn.Parameter(torch.tensor(0.02)) + + def forward(self, values): + raw = values["obj_positions"][:, 1:-1] + self.offset + return { + "obj_positions_norm": raw, + "obj_positions_icp_norm": raw + self.offset, + "obj_accelerations_no_icp_norm_anchor": self.offset.expand( + 2, 2, 2, 4, 3 + ), + "obj_accelerations_norm": self.offset.expand_as(raw) * 2, + } + + model = Predictions() + batch = rigidformer_batch(v) + got = RigidFormerObjective().compute_loss(model, batch) + + def huber(e): + a = e.abs() + return torch.where(a < 1, 0.5 * e.square(), a - 0.5).sum(-1) + + expected = 0.0 + terms = [] + for b, objects in enumerate((2, 1)): + for t in range(2): + for o in range(objects): + scale = float(v["skip_step_size"][b]) ** 2 + gt = v["obj_positions"][b, t + 2, o, :4] + prev = v["obj_positions"][b, t, o, :4] + curr = v["obj_positions"][b, t + 1, o, :4] + raw = curr + model.offset + accel = gt - 2 * curr + prev + terms.append( + 10 * huber((raw - gt) / scale) + + 10 * huber((raw + model.offset - gt) / scale) + + huber((model.offset - accel) / scale) + + huber((2 * model.offset - accel) / scale) + ) + expected = torch.cat(terms).mean() + assert got.weight == 24.0 + torch.testing.assert_close(got.loss, expected) + got.loss.backward() + actual = model.offset.grad.clone() + model.offset.grad = None + expected.backward() + torch.testing.assert_close(model.offset.grad, actual) + # The final target is absent from conditioning. + assert batch.inputs["obj_positions"].shape[1] == 3 + + +@pytest.mark.parametrize("random_step", [True, False]) +def test_packed_geometry_scales_once_and_replays_owned_windows( + packed_data, random_step +): + a = source(packed_data, random_step=random_step) + generator = torch.Generator().manual_seed(999) + first = a.sample(2, generator=generator) + a.on_update(None) + saved = a.state_dict() + rng = generator.get_state().clone() + expected = a.sample(2, generator=generator) + b = source(packed_data, random_step=random_step) + b.validate_state_dict(saved) + b.load_state_dict(saved) + generator.set_state(rng) + actual = b.sample(2, generator=generator) + for key, value in expected.inputs.items(): + torch.testing.assert_close(actual.inputs[key], value, rtol=0, atol=0) + torch.testing.assert_close(expected.targets, actual.targets, rtol=0, atol=0) + assert ( + first.inputs["obj_positions"].untyped_storage().data_ptr() + != actual.inputs["obj_positions"].untyped_storage().data_ptr() + ) + assert float( + first.inputs["obj_positions"][0, 0, 0, :, 2].max() - 1.0 + ) == pytest.approx(0.07, abs=1e-6) + assert float( + first.inputs["obj_positions"][0, 0, 1, :, 2].max() - 1.0 + ) == pytest.approx(0.14, abs=1e-6) + + +def test_changed_geometry_rejected_before_sampling_recovery(packed_data): + a = source(packed_data) + state = a.state_dict() + p = packed_data / "scene_0" / "frames.npy" + values = np.load(p) + values[0, 0, 0] += 1 + np.save(p, values) + with pytest.raises(ValueError, match="hash differs"): + a.validate_state_dict(state) + + +def test_geometry_missing_from_manifest_is_not_a_verified_source(packed_data): + path = packed_data / "manifest.json" + manifest = json.loads(path.read_text()) + manifest["files"].pop("scene_0/frames.npy") + path.write_text(json.dumps(manifest)) + with pytest.raises(ValueError, match="not committed by manifest"): + source(packed_data) + + +@pytest.mark.parametrize("missing_cache", ["anchor", "fps_order"]) +def test_missing_prepared_cache_does_not_modify_source(packed_data, missing_cache): + from embodichain_models.dynamics.rigidformer.data import MOViSequenceSource + + original = source(packed_data, random_step=False) + valid = dict(original.config) + if missing_cache == "anchor": + config = {**valid, "anchor_k": 2} + message = "Prepared anchor cache is missing" + else: + path = packed_data / "scene_0" / "anchors" / "obj_0_fpsorder_cf.npy" + path.unlink() + manifest_path = packed_data / "manifest.json" + manifest = json.loads(manifest_path.read_text()) + manifest["files"].pop(str(path.relative_to(packed_data))) + manifest_path.write_text(json.dumps(manifest)) + config = dict(valid) + valid["fps_order"] = False + message = "Prepared FPS-order cache is missing" + before = {p.relative_to(packed_data) for p in packed_data.rglob("*")} + with pytest.raises(ValueError, match=message): + MOViSequenceSource( + packed_data, + packed_data / "splits.json", + packed_data / "scales.json", + data_config=config, + ) + assert {p.relative_to(packed_data) for p in packed_data.rglob("*")} == before + MOViSequenceSource( + packed_data, + packed_data / "splits.json", + packed_data / "scales.json", + data_config=valid, + ) + + +def test_new_manifest_reloads_replaced_packed_geometry(packed_data): + first = source(packed_data, random_step=False) + before = first.sample(1, generator=torch.Generator().manual_seed(17)) + scene = packed_data / "scene_0" + manifest_path = packed_data / "manifest.json" + manifest = json.loads(manifest_path.read_text()) + for name in ("frames.npy", "vis_traj.npy"): + path = scene / name + values = np.load(path).copy() + if name == "frames.npy": + values[1:, :, 0] += 0.5 + else: + values[0, 1:, :, 0] += 0.5 + replacement = scene / "replacement.tmp.npy" + np.save(replacement, values) + os.replace(replacement, path) + manifest["files"][str(path.relative_to(packed_data))] = hashlib.sha256( + path.read_bytes() + ).hexdigest() + manifest_path.write_text(json.dumps(manifest)) + second = source(packed_data, random_step=False) + after = second.sample(1, generator=torch.Generator().manual_seed(17)) + assert second.source_id != first.source_id + torch.testing.assert_close( + after.inputs["obj_positions"][:, :, :, :, 0], + before.inputs["obj_positions"][:, :, :, :, 0] + 0.5, + ) + + +def test_new_manifest_reloads_replaced_per_frame_geometry(tmp_path): + root = make_packed(tmp_path / "ToyData", frames_count=10) + scene = root / "scene_0" + frames = np.load(scene / "frames.npy") + manifest_path = root / "manifest.json" + manifest = json.loads(manifest_path.read_text()) + for frame, values in enumerate(frames): + directory = scene / f"frame_{frame:05d}" + directory.mkdir() + for obj in range(2): + path = directory / f"obj_{obj}_pos.npy" + np.save(path, values[obj * 4 : (obj + 1) * 4]) + manifest["files"][str(path.relative_to(root))] = hashlib.sha256( + path.read_bytes() + ).hexdigest() + for name in ("frames.npy", "frames_counts.npy"): + (scene / name).unlink() + manifest["files"].pop(f"scene_0/{name}") + manifest_path.write_text(json.dumps(manifest)) + first = source(root, random_step=False) + before = first.sample(1, generator=torch.Generator().manual_seed(17)) + path = scene / "frame_00002" / "obj_0_pos.npy" + changed = np.load(path).copy() + changed[:, 0] += 0.5 + replacement = path.with_suffix(".tmp.npy") + np.save(replacement, changed) + os.replace(replacement, path) + manifest["files"][str(path.relative_to(root))] = hashlib.sha256( + path.read_bytes() + ).hexdigest() + manifest_path.write_text(json.dumps(manifest)) + second = source(root, random_step=False) + after = second.sample(1, generator=torch.Generator().manual_seed(17)) + assert second.source_id != first.source_id + assert after.inputs["obj_positions"][0, 0, 0, 0, 0].item() == pytest.approx( + before.inputs["obj_positions"][0, 0, 0, 0, 0].item() + 0.35 + ) + + +def test_fixed_step_and_stride_are_applied(packed_data): + from embodichain_models.dynamics.rigidformer.data import MOViSequenceSource + + original = source(packed_data, random_step=False) + updated = MOViSequenceSource( + packed_data, + packed_data / "splits.json", + packed_data / "scales.json", + data_config={**original.config, "skip_step_size": 5, "stride": 3}, + ) + batch = updated.sample(1, generator=torch.Generator().manual_seed(17)) + assert batch.inputs["skip_step_size"].tolist() == [5.0] + assert np.diff(batch.metadata["frame_ids"][0]).tolist() == [5] * 7 + assert len(updated.dataset) == 6 + + +def test_small_dataset_pads_every_distributed_rank(tmp_path, monkeypatch): + root = make_packed(tmp_path / "ToyData", frames_count=10) + monkeypatch.setattr(torch.distributed, "is_initialized", lambda: True) + monkeypatch.setattr(torch.distributed, "get_world_size", lambda: 4) + for rank in range(4): + monkeypatch.setattr(torch.distributed, "get_rank", lambda: rank) + dataset = source(root, random_step=False) + assert dataset.updates_per_epoch(1) == 1 + batch = dataset.sample(1, generator=torch.Generator().manual_seed(17)) + assert batch.metadata["scenes"] == ["scene_0"] + + +def test_default_object_padding_counts_toward_memory_limit(packed_data): + from embodichain_models.dynamics.rigidformer.data import MOViSequenceSource + + original = source(packed_data, random_step=False) + config = dict(original.config) + config.pop("max_objects_per_frame") + config["sequence_length"] = 3 + bounded = MOViSequenceSource( + packed_data, + packed_data / "splits.json", + packed_data / "scales.json", + data_config=config, + max_batch_bytes=2000, + ) + with pytest.raises(ValueError, match="max_batch_bytes"): + bounded.sample(1, generator=torch.Generator().manual_seed(17)) + + +def test_recipe_rejects_anchor_count_mismatched_with_model(packed_data, tmp_path): + from embodichain_models.dynamics.rigidformer.sl_recipe import ( + build_rigidformer_trainer, + ) + + config = recovery_config(packed_data, tmp_path / "output", "cpu") + config["data"]["sampling"]["anchor_k"] = 2 + with pytest.raises(ValueError, match="anchor_k must match"): + build_rigidformer_trainer(config) + + +def test_centered_predictions_use_centered_position_targets(): + class CenteredPrediction(torch.nn.Module): + reference_config = {"data": {"center_scene_by_t0": True}} + + def forward(self, values): + positions = values["obj_positions"] + center = positions[:, 0].mean(dim=(1, 2))[:, None, None, None] + predicted = positions[:, 1:-1] - center + zero = torch.zeros_like(predicted) + return { + "obj_positions_norm": predicted, + "obj_positions_icp_norm": predicted, + "obj_accelerations_no_icp_norm_anchor": zero, + "obj_accelerations_norm": zero, + } + + positions = torch.tensor([1.0, 2.0, 3.0]).view(1, 1, 1, 1, 3) + positions = positions.expand(1, 3, 2, 1, 3).clone() + values = { + "obj_positions": positions, + "valid_vtx_mask": torch.ones(1, 3, 2, 1, dtype=torch.bool), + "valid_obj_mask": torch.ones(1, 2, dtype=torch.bool), + "anchor_idx": torch.zeros(1, 2, 1, dtype=torch.long), + "skip_step_size": torch.ones(1), + } + result = RigidFormerObjective().compute_loss( + CenteredPrediction(), rigidformer_batch(values) + ) + assert result.loss.item() == pytest.approx(0.0) + + +@pytest.mark.gpu +def test_centered_model_loss_is_translation_invariant(): + pytest.importorskip("roma") + from embodichain_models.dynamics.rigidformer import RigidFormerModel + + cfg = tiny_cfg() + cfg.architecture["data"]["center_scene_by_t0"] = True + model = RigidFormerModel(cfg).cuda().eval() + values = geometry_batch("cuda") + values["valid_obj_mask"][:] = True + moved = {key: value.clone() for key, value in values.items()} + translation = torch.tensor([0.4, -0.6, 1.2], device="cuda") + moved["obj_positions"] += translation + moved["ref_obj_positions"] += translation + objective = RigidFormerObjective() + with torch.no_grad(): + original = objective.compute_loss(model, rigidformer_batch(values)).loss + translated = objective.compute_loss(model, rigidformer_batch(moved)).loss + torch.testing.assert_close(translated, original, rtol=1e-4, atol=1e-5) + + +def test_export_accepts_legacy_numpy_rng_pickle(tmp_path, monkeypatch): + from embodichain_models.dynamics.rigidformer.checkpoint import ( + export_rigidformer_checkpoint, + ) + + try: + reconstruct = np._core.multiarray._reconstruct + except AttributeError: + reconstruct = np.core.multiarray._reconstruct + checkpoint = tmp_path / "training.pt" + original_module = reconstruct.__module__ + monkeypatch.setattr(reconstruct, "__module__", "numpy.core.multiarray") + raw, average = torch.tensor([1.0]), torch.tensor([2.0]) + torch.save( + { + "rng": {"numpy": np.random.get_state()}, + "signature": { + "family": "rigidformer-author-v1", + "reference_config": {}, + "author_revision": "test", + }, + "learner": {"model": {"network.weight": raw}, "num_updates": 1}, + "training_states": { + "ema": {"state": {"shadow": {"network.weight": average}}} + }, + }, + checkpoint, + ) + monkeypatch.setattr(reconstruct, "__module__", original_module) + for selection, expected in (("raw", raw), ("ema", average)): + output = export_rigidformer_checkpoint( + checkpoint, tmp_path / f"{selection}.pt", weights=selection + ) + payload = torch.load(output, weights_only=True) + torch.testing.assert_close(payload["model_state_dict"]["weight"], expected) + + +@pytest.mark.parametrize("random_sampling", [False, True]) +def test_compact_windows_preserve_scene_stride_and_boundaries(random_sampling): + from pathlib import Path + from embodichain_models.dynamics.rigidformer._native.data._index import ( + _SceneWindows, + ) + + windows = _SceneWindows(random_sampling=random_sampling) + windows.add_scene("a", [Path(f"a/frame_{i:05d}") for i in range(7)], 4, 2) + windows.add_scene("b", [Path(f"b/frame_{i:05d}") for i in range(8)], 4, 2) + expected = [ + ("a", [0, 1, 2, 3]), + ("a", [2, 3, 4, 5]), + ("b", [0, 1, 2, 3]), + ("b", [2, 3, 4, 5]), + ("b", [4, 5, 6, 7]), + ] + assert len(windows) == len(expected) + for i, (scene, frames) in enumerate(expected): + item = windows[i] + assert item["dataset"] == scene + if not random_sampling: + assert item["frame_ids"] == frames + assert item["frames"] == [f"{scene}/frame_{j:05d}" for j in frames] + assert windows[-1] == windows[len(expected) - 1] + + +@pytest.mark.gpu +def test_recipe_resume_crosses_epoch_and_warmup_with_exact_ema(tmp_path): + from embodichain_models.dynamics.rigidformer.sl_recipe import ( + build_rigidformer_trainer, + ) + from embodichain_models.dynamics.rigidformer import ( + export_rigidformer_checkpoint, + load_rigidformer_checkpoint, + ) + + root = make_packed(tmp_path / "short-scenes", frames_count=13) + config = recovery_config(root, tmp_path / "full", "cuda") + config["trainer"]["batch_size"] = 2 + full = build_rigidformer_trainer(config) + assert full.signature["steps_per_epoch"] == 2 + full.fit() + config["trainer"]["output_dir"] = str(tmp_path / "split") + split = build_rigidformer_trainer(config) + split.fit(stop_after=3) + config["trainer"]["output_dir"] = str(tmp_path / "resumed") + resumed = build_rigidformer_trainer(config) + resumed.resume(tmp_path / "split/last.pt") + resumed.fit() + for key in ("learner", "training_states", "sampler", "rng"): + a = torch.load(tmp_path / "full/last.pt", weights_only=False)[key] + b = torch.load(tmp_path / "resumed/last.pt", weights_only=False)[key] + assert_exact(a, b) + schedule = resumed.training_states["epoch_schedule"] + assert schedule.scheduler.last_epoch == 3 + for selection in ("raw", "ema"): + path = export_rigidformer_checkpoint( + tmp_path / "resumed/last.pt", + tmp_path / f"{selection}.pt", + weights=selection, + ) + model = load_rigidformer_checkpoint(path) + expected = ( + resumed.learner.model.state_dict() + if selection == "raw" + else resumed.training_states["ema"].shadow + ) + for key, value in model.state_dict().items(): + torch.testing.assert_close(value, expected[key].cpu(), rtol=0, atol=0) + + +@pytest.mark.gpu +def test_fp32_geometry_augmentation_precedes_bf16_model_autocast(): + pytest.importorskip("roma") + from embodichain_models.dynamics.rigidformer import RigidFormerModel + from embodichain_models.dynamics.rigidformer._native.data.gpu_augmentation import ( + GPUAugmentation, + ) + + cfg = tiny_cfg() + architecture = cfg.architecture + architecture["data"]["augmentation"] = { + "rotation_augmentation": { + "enable": True, + "angle_step_size": 90, + "prob_use_rotation": 1.0, + } + } + model = RigidFormerModel(cfg).cuda().train() + values = geometry_batch("cuda") + values["valid_obj_mask"][:] = True + objective = RigidFormerObjective( + augmentation=GPUAugmentation(architecture, torch.device("cuda")) + ) + with torch.autocast("cuda", dtype=torch.bfloat16): + result = objective.compute_loss(model, rigidformer_batch(values)) + result.loss.backward() + assert torch.isfinite(result.loss) + assert all( + p.grad is None or torch.isfinite(p.grad).all() for p in model.parameters() + )