RigidFormer geometry dynamics#
The model family implements the author’s RigidFormer architecture from revision 160fe0852c51cffae7b7440452b58d0b49db6df6. Licensed implementation code is private to the model package; public adapters do not import the research showcase or a simulator. The optional extra installs numerical/data dependencies, without the community rigidformer package.
Geometry uses aligned material points, explicit reference geometry, valid object/vertex masks and physical mass/friction/restitution. The author KNN path requires CUDA and at least two valid objects per scene. Model construction and 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 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#
Configure a native model using an author architecture. |
|
Predict rigid geometry with the author architecture and weight layout. |
|
|
Load an official model/EMA weights file and its architecture. |
|
Export a native SL training checkpoint to tensor/config inference weights. |
Read author packed scenes and checkpoint a rank-local shuffled traversal. |
|
Reproduce the author's four anchor-only smooth-L1 losses. |
|
|
Separate shifted conditioning history and next-frame supervision. |
|
Build an author-based MOVi recipe without importing example helpers. |
- embodichain_models.dynamics.rigidformer.RIGIDFORMER_REVISION = '160fe0852c51cffae7b7440452b58d0b49db6df6'#
str(object=’’) -> str str(bytes_or_buffer[, encoding[, errors]]) -> str
Create a new string object from the given object. If encoding or errors is specified, then the object must expose a data buffer that will be decoded using the given encoding and error handler. Otherwise, returns the result of object.__str__() (if defined) or repr(object). encoding defaults to sys.getdefaultencoding(). errors defaults to ‘strict’.
- class embodichain_models.dynamics.rigidformer.RigidFormerCfg[source]#
Configure a native model using an author architecture.
- Parameters:
preset (
str) – Author MOVi architecture/configuration to load.architecture (
dict[str,Any] |None) – Optional full author configuration. Checkpoint imports supply their saved configuration rather than guessing a preset.attention_backend (
str) – Explicit sdpa or optional flash backend.max_batch_bytes (
int) – Bound on geometry and context tensors before forwarding.
Physical intervals are step_size / frame_rate seconds. The author model integrates in its learned unit system with dt=1; these are distinct units.
Methods:
__init__([preset, architecture, ...])copy(**kwargs)Return a new object replacing specified fields with new values.
Return an owned author configuration without opening data assets.
replace(**kwargs)Return a new object replacing specified fields with new values.
to_dict()Convert an object into dictionary recursively.
validate([prefix])Check the validity of configclass object.
- __init__(preset=<factory>, architecture=<factory>, attention_backend=<factory>, max_batch_bytes=<factory>)#
- copy(**kwargs)#
Return a new object replacing specified fields with new values.
This is especially useful for frozen classes. Example usage:
@configclass(frozen=True) class C: x: int y: int c = C(1, 2) c1 = c.replace(x=3) assert c1.x == 3 and c1.y == 2
- Parameters:
obj (
object) – The object to replace.**kwargs – The fields to replace and their new values.
- Return type:
object- Returns:
The new object.
- reference_config()[source]#
Return an owned author configuration without opening data assets.
- Return type:
dict[str,Any]- Returns:
Complete author configuration, including explicit architecture overrides.
- replace(**kwargs)#
Return a new object replacing specified fields with new values.
This is especially useful for frozen classes. Example usage:
@configclass(frozen=True) class C: x: int y: int c = C(1, 2) c1 = c.replace(x=3) assert c1.x == 3 and c1.y == 2
- Parameters:
obj (
object) – The object to replace.**kwargs – The fields to replace and their new values.
- Return type:
object- Returns:
The new object.
- to_dict()#
Convert an object into dictionary recursively.
Note
Ignores all names starting with “__” (i.e. built-in methods).
- Parameters:
obj (
object) – An instance of a class to convert.- Raises:
ValueError – When input argument is not an object.
- Return type:
dict[str,Any]- Returns:
Converted dictionary mapping.
- validate(prefix='')#
Check the validity of configclass object.
This function checks if the object is a valid configclass object. A valid configclass object contains no MISSING entries.
- Parameters:
obj (
object) – The object to check.prefix (
str) – The prefix to add to the missing fields. Defaults to ‘’.
- Return type:
list[str]- Returns:
A list of missing fields.
- Raises:
TypeError – When the object is not a valid configuration object.
- class embodichain_models.dynamics.rigidformer.RigidFormerModel[source]#
Predict rigid geometry with the author architecture and weight layout.
- Parameters:
cfg (
RigidFormerCfg|None) – Native architecture, optional attention backend and geometry budget.
Construction loads optional numerical dependencies, but no data, weights, examples or simulation. Inputs follow the author’s batch layout. Training uses shifted conditioning frames; free rollout takes only warmup frames.
Methods:
__init__([cfg])Initialize internal Module state, shared by both nn.Module and ScriptModule.
forward(batch, *[, teacher_forcing_ratio])Predict shifted training frames using explicit causal conditioning.
rollout(batch, *, num_frames[, warmup_frames])Predict a free trajectory from warmup geometry.
- __init__(cfg=None)[source]#
Initialize internal Module state, shared by both nn.Module and ScriptModule.
- forward(batch, *, teacher_forcing_ratio=1.0)[source]#
Predict shifted training frames using explicit causal conditioning.
- Parameters:
batch (
Mapping[str,Any]) – Author-layout geometry, masks, anchors and physical context.teacher_forcing_ratio (
float) – Ratio in [0,1]; one matches released recipes.
- Returns:
,2:].
- Return type:
Author prediction fields with frames aligned to positions[
- rollout(batch, *, num_frames, warmup_frames=4)[source]#
Predict a free trajectory from warmup geometry.
- Parameters:
batch (
Mapping[str,Any]) – Geometry containing exactly the declared warmup frames.num_frames (
int) – Total returned frames including warmup.warmup_frames (
int) – Initial ground-truth frames; four matches evaluation.
- Return type:
dict[str,Tensor]- Returns:
Author rollout fields. This inference method does not retain autograd.
- embodichain_models.dynamics.rigidformer.load_rigidformer_checkpoint(path, *, device='cpu', attention_backend='sdpa')[source]#
Load an official model/EMA weights file and its architecture.
- Parameters:
path (
str|Path) – Official tensor/config checkpoint; decoded with weights_only=True.device (
str|device) – Destination device, independent of the file’s saved device.attention_backend (
str) – Explicit native attention implementation.
- Return type:
- Returns:
Evaluation model with every author weight loaded strictly.
This is a weights import, not recovery of an official training process.
- embodichain_models.dynamics.rigidformer.export_rigidformer_checkpoint(source, destination, *, weights='ema')[source]#
Export a native SL training checkpoint to tensor/config inference weights.
- Parameters:
source (
str|Path) – Checkpoint produced by the native author-based SL recipe.destination (
str|Path) – Fresh output path; existing files are rejected.weights (
str) – Explicit ema or raw model selection.
- Return type:
Path- Returns:
Written inference checkpoint, loadable by load_rigidformer_checkpoint.
- class embodichain_models.dynamics.rigidformer.MOViSequenceSource[source]#
Read author packed scenes and checkpoint a rank-local shuffled traversal.
- Parameters:
root (
str|Path) – Converted MOVi-Sphere, MOVi-A or MOVi-B directory.split_file (
str|Path) – Frozen author scene split JSON.object_scale_file (
str|Path) – Frozen author size sidecar for unscaled stored vertices.split (
str) – train, val or test.data_config (
Mapping[str,Any] |None) – Author sampling and augmentation settings.seed (
int) – Epoch-shuffle seed; per-batch augmentation RNG is caller-owned.max_batch_bytes (
int) – 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.
Methods:
__init__(root, split_file, object_scale_file, *)load_state_dict(state)Restore a prevalidated committed traversal and discard staged reads.
on_update(learner)Commit the staged window traversal after a successful optimizer step.
sample(batch_size, *, generator)Stage one bounded batch; on_update commits its traversal position.
Return the committed rank-local sampling boundary.
updates_per_epoch(batch_size)Return full-batch updates per rank, matching drop-last training.
validate_state_dict(state)Reject incompatible geometry or traversal before training mutation.
- __init__(root, split_file, object_scale_file, *, split='train', data_config=None, seed=42, max_batch_bytes=536870912)[source]#
- load_state_dict(state)[source]#
Restore a prevalidated committed traversal and discard staged reads.
- Parameters:
state (
Mapping[str,Any]) – Snapshot accepted by validate_state_dict.- Return type:
None
- on_update(learner)[source]#
Commit the staged window traversal after a successful optimizer step.
- Parameters:
learner (
Any) – Host update context; progress comes from the staged batch.- Return type:
None
- sample(batch_size, *, generator)[source]#
Stage one bounded batch; on_update commits its traversal position.
- Parameters:
batch_size (
int) – Number of windows on this rank.generator (
Generator) – Caller-owned sampling/augmentation RNG.
- Return type:
- Returns:
Owned CPU shifted geometry supervision.
- state_dict()[source]#
Return the committed rank-local sampling boundary.
- Return type:
dict[str,Any]- Returns:
Owned source identity, epoch, cursor and rank/world-size values.
- class embodichain_models.dynamics.rigidformer.RigidFormerObjective[source]#
Reproduce the author’s four anchor-only smooth-L1 losses.
- Parameters:
weights (
Mapping[str,float] |None) – Position/acceleration before/after Kabsch weights.time_step_scale (
bool) – Divide errors by step_size squared before smooth-L1.augmentation (
Any) – Optional author geometry augmentation, applied in FP32 before the network’s numerical context.
Methods:
__init__([weights, time_step_scale, ...])compute_loss(model, batch)Compute one local anchor mean; the shared learner owns DDP reduction.
- compute_loss(model, batch)[source]#
Compute one local anchor mean; the shared learner owns DDP reduction.
- Parameters:
model (
Module) – Native RigidFormer model.batch (
SupervisedBatch) – Shifted conditioning and owned target geometry.
- Return type:
- Returns:
Weighted scalar loss, diagnostic terms and valid-anchor count.
- embodichain_models.dynamics.rigidformer.rigidformer_batch(values, *, metadata=None)[source]#
Separate shifted conditioning history and next-frame supervision.
- Parameters:
values (
Mapping[str,Any]) – Author batch including an aligned sequence of at least 3 frames.metadata (
Mapping[str,Any] |None) – Optional immutable data and sampling provenance.
- Return type:
- Returns:
Supervision for frames 2..T-1 with causal conditioning through T-2. Earlier targets also occur as subsequent teacher-forcing inputs.
- embodichain_models.dynamics.rigidformer.build_rigidformer_trainer(config)[source]#
Build an author-based MOVi recipe without importing example helpers.
- Parameters:
config (
Mapping[str,Any]) – 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.- Return type:
- Returns:
Shared trainer with committed sampling, epoch schedule and FP32 EMA.
Architecture and geometry contract#
Configure a native model using an author architecture. |
|
Predict rigid geometry with the author architecture and weight layout. |
- embodichain_models.dynamics.rigidformer.model.RIGIDFORMER_REVISION = '160fe0852c51cffae7b7440452b58d0b49db6df6'#
str(object=’’) -> str str(bytes_or_buffer[, encoding[, errors]]) -> str
Create a new string object from the given object. If encoding or errors is specified, then the object must expose a data buffer that will be decoded using the given encoding and error handler. Otherwise, returns the result of object.__str__() (if defined) or repr(object). encoding defaults to sys.getdefaultencoding(). errors defaults to ‘strict’.
- class embodichain_models.dynamics.rigidformer.model.RigidFormerCfg[source]#
Configure a native model using an author architecture.
- Parameters:
preset (
str) – Author MOVi architecture/configuration to load.architecture (
dict[str,Any] |None) – Optional full author configuration. Checkpoint imports supply their saved configuration rather than guessing a preset.attention_backend (
str) – Explicit sdpa or optional flash backend.max_batch_bytes (
int) – Bound on geometry and context tensors before forwarding.
Physical intervals are step_size / frame_rate seconds. The author model integrates in its learned unit system with dt=1; these are distinct units.
Methods:
__init__([preset, architecture, ...])copy(**kwargs)Return a new object replacing specified fields with new values.
Return an owned author configuration without opening data assets.
replace(**kwargs)Return a new object replacing specified fields with new values.
to_dict()Convert an object into dictionary recursively.
validate([prefix])Check the validity of configclass object.
- __init__(preset=<factory>, architecture=<factory>, attention_backend=<factory>, max_batch_bytes=<factory>)#
- copy(**kwargs)#
Return a new object replacing specified fields with new values.
This is especially useful for frozen classes. Example usage:
@configclass(frozen=True) class C: x: int y: int c = C(1, 2) c1 = c.replace(x=3) assert c1.x == 3 and c1.y == 2
- Parameters:
obj (
object) – The object to replace.**kwargs – The fields to replace and their new values.
- Return type:
object- Returns:
The new object.
- reference_config()[source]#
Return an owned author configuration without opening data assets.
- Return type:
dict[str,Any]- Returns:
Complete author configuration, including explicit architecture overrides.
- replace(**kwargs)#
Return a new object replacing specified fields with new values.
This is especially useful for frozen classes. Example usage:
@configclass(frozen=True) class C: x: int y: int c = C(1, 2) c1 = c.replace(x=3) assert c1.x == 3 and c1.y == 2
- Parameters:
obj (
object) – The object to replace.**kwargs – The fields to replace and their new values.
- Return type:
object- Returns:
The new object.
- to_dict()#
Convert an object into dictionary recursively.
Note
Ignores all names starting with “__” (i.e. built-in methods).
- Parameters:
obj (
object) – An instance of a class to convert.- Raises:
ValueError – When input argument is not an object.
- Return type:
dict[str,Any]- Returns:
Converted dictionary mapping.
- validate(prefix='')#
Check the validity of configclass object.
This function checks if the object is a valid configclass object. A valid configclass object contains no MISSING entries.
- Parameters:
obj (
object) – The object to check.prefix (
str) – The prefix to add to the missing fields. Defaults to ‘’.
- Return type:
list[str]- Returns:
A list of missing fields.
- Raises:
TypeError – When the object is not a valid configuration object.
- class embodichain_models.dynamics.rigidformer.model.RigidFormerModel[source]#
Predict rigid geometry with the author architecture and weight layout.
- Parameters:
cfg (
RigidFormerCfg|None) – Native architecture, optional attention backend and geometry budget.
Construction loads optional numerical dependencies, but no data, weights, examples or simulation. Inputs follow the author’s batch layout. Training uses shifted conditioning frames; free rollout takes only warmup frames.
Methods:
__init__([cfg])Initialize internal Module state, shared by both nn.Module and ScriptModule.
forward(batch, *[, teacher_forcing_ratio])Predict shifted training frames using explicit causal conditioning.
rollout(batch, *, num_frames[, warmup_frames])Predict a free trajectory from warmup geometry.
- __init__(cfg=None)[source]#
Initialize internal Module state, shared by both nn.Module and ScriptModule.
- forward(batch, *, teacher_forcing_ratio=1.0)[source]#
Predict shifted training frames using explicit causal conditioning.
- Parameters:
batch (
Mapping[str,Any]) – Author-layout geometry, masks, anchors and physical context.teacher_forcing_ratio (
float) – Ratio in [0,1]; one matches released recipes.
- Returns:
,2:].
- Return type:
Author prediction fields with frames aligned to positions[
- rollout(batch, *, num_frames, warmup_frames=4)[source]#
Predict a free trajectory from warmup geometry.
- Parameters:
batch (
Mapping[str,Any]) – Geometry containing exactly the declared warmup frames.num_frames (
int) – Total returned frames including warmup.warmup_frames (
int) – Initial ground-truth frames; four matches evaluation.
- Return type:
dict[str,Tensor]- Returns:
Author rollout fields. This inference method does not retain autograd.
Official weights and native SL export#
|
Load an official model/EMA weights file and its architecture. |
|
Export a native SL training checkpoint to tensor/config inference weights. |
- embodichain_models.dynamics.rigidformer.checkpoint.load_rigidformer_checkpoint(path, *, device='cpu', attention_backend='sdpa')[source]#
Load an official model/EMA weights file and its architecture.
- Parameters:
path (
str|Path) – Official tensor/config checkpoint; decoded with weights_only=True.device (
str|device) – Destination device, independent of the file’s saved device.attention_backend (
str) – Explicit native attention implementation.
- Return type:
- Returns:
Evaluation model with every author weight loaded strictly.
This is a weights import, not recovery of an official training process.
- embodichain_models.dynamics.rigidformer.checkpoint.export_rigidformer_checkpoint(source, destination, *, weights='ema')[source]#
Export a native SL training checkpoint to tensor/config inference weights.
- Parameters:
source (
str|Path) – Checkpoint produced by the native author-based SL recipe.destination (
str|Path) – Fresh output path; existing files are rejected.weights (
str) – Explicit ema or raw model selection.
- Return type:
Path- Returns:
Written inference checkpoint, loadable by load_rigidformer_checkpoint.
Prepared packed scenes#
Read author packed scenes and checkpoint a rank-local shuffled traversal. |
- class embodichain_models.dynamics.rigidformer.data.MOViSequenceSource[source]#
Read author packed scenes and checkpoint a rank-local shuffled traversal.
- Parameters:
root (
str|Path) – Converted MOVi-Sphere, MOVi-A or MOVi-B directory.split_file (
str|Path) – Frozen author scene split JSON.object_scale_file (
str|Path) – Frozen author size sidecar for unscaled stored vertices.split (
str) – train, val or test.data_config (
Mapping[str,Any] |None) – Author sampling and augmentation settings.seed (
int) – Epoch-shuffle seed; per-batch augmentation RNG is caller-owned.max_batch_bytes (
int) – 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.
Methods:
__init__(root, split_file, object_scale_file, *)load_state_dict(state)Restore a prevalidated committed traversal and discard staged reads.
on_update(learner)Commit the staged window traversal after a successful optimizer step.
sample(batch_size, *, generator)Stage one bounded batch; on_update commits its traversal position.
Return the committed rank-local sampling boundary.
updates_per_epoch(batch_size)Return full-batch updates per rank, matching drop-last training.
validate_state_dict(state)Reject incompatible geometry or traversal before training mutation.
- __init__(root, split_file, object_scale_file, *, split='train', data_config=None, seed=42, max_batch_bytes=536870912)[source]#
- load_state_dict(state)[source]#
Restore a prevalidated committed traversal and discard staged reads.
- Parameters:
state (
Mapping[str,Any]) – Snapshot accepted by validate_state_dict.- Return type:
None
- on_update(learner)[source]#
Commit the staged window traversal after a successful optimizer step.
- Parameters:
learner (
Any) – Host update context; progress comes from the staged batch.- Return type:
None
- sample(batch_size, *, generator)[source]#
Stage one bounded batch; on_update commits its traversal position.
- Parameters:
batch_size (
int) – Number of windows on this rank.generator (
Generator) – Caller-owned sampling/augmentation RNG.
- Return type:
- Returns:
Owned CPU shifted geometry supervision.
- state_dict()[source]#
Return the committed rank-local sampling boundary.
- Return type:
dict[str,Any]- Returns:
Owned source identity, epoch, cursor and rank/world-size values.
Shifted supervision and anchor losses#
Reproduce the author's four anchor-only smooth-L1 losses. |
|
|
Separate shifted conditioning history and next-frame supervision. |
- class embodichain_models.dynamics.rigidformer.supervised.RigidFormerObjective[source]#
Reproduce the author’s four anchor-only smooth-L1 losses.
- Parameters:
weights (
Mapping[str,float] |None) – Position/acceleration before/after Kabsch weights.time_step_scale (
bool) – Divide errors by step_size squared before smooth-L1.augmentation (
Any) – Optional author geometry augmentation, applied in FP32 before the network’s numerical context.
Methods:
__init__([weights, time_step_scale, ...])compute_loss(model, batch)Compute one local anchor mean; the shared learner owns DDP reduction.
- compute_loss(model, batch)[source]#
Compute one local anchor mean; the shared learner owns DDP reduction.
- Parameters:
model (
Module) – Native RigidFormer model.batch (
SupervisedBatch) – Shifted conditioning and owned target geometry.
- Return type:
- Returns:
Weighted scalar loss, diagnostic terms and valid-anchor count.
- embodichain_models.dynamics.rigidformer.supervised.rigidformer_batch(values, *, metadata=None)[source]#
Separate shifted conditioning history and next-frame supervision.
- Parameters:
values (
Mapping[str,Any]) – Author batch including an aligned sequence of at least 3 frames.metadata (
Mapping[str,Any] |None) – Optional immutable data and sampling provenance.
- Return type:
- Returns:
Supervision for frames 2..T-1 with causal conditioning through T-2. Earlier targets also occur as subsequent teacher-forcing inputs.
Recoverable EMA state#
Maintain the author's ramped FP32 EMA, including model buffers. |
- class embodichain_models.dynamics.rigidformer.ema.RigidFormerEMA[source]#
Maintain the author’s ramped FP32 EMA, including model buffers.
- Parameters:
model (
Module) – Model whose parameter and buffer layout remains fixed.decay (
float) – Maximum decay; initial updates use the author’s warmup.
Methods:
__init__(model[, decay])load_state_dict(state)Restore a prevalidated EMA without changing live model weights.
on_update(learner)Update the average once after a successful parameter update.
Return owned FP32 shadows and the exact EMA update cursor.
validate_state_dict(state)Validate EMA policy and every shadow tensor before restoration.
- load_state_dict(state)[source]#
Restore a prevalidated EMA without changing live model weights.
- Parameters:
state (
Mapping[str,Any]) – Snapshot accepted by validate_state_dict.- Return type:
None
- on_update(learner)[source]#
Update the average once after a successful parameter update.
- Parameters:
learner (
Any) – Successful host update context. Values are read from the model supplied at construction.- Return type:
None
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.
The showcase owns asset preparation, the author evaluator and independent reference comparison. Fixed-checkpoint test evaluation does not select a checkpoint. New learning recipes should select using validation data. Reports record source/data/checkpoint identities, attention backend and subset/full scope. Short numerical comparisons do not establish published accuracy.