Official model library

The main embodichain distribution bundles embodichain_models with the same release version and dependencies. This package owns concrete model families, their physical representations, model-specific data adapters and supervised recipes. General MLPs and feature normalization remain in embodichain.learning.models; shared optimizer updates, recovery and training orchestration remain in Supervised learning (SL).

Import a selected family explicitly. Importing the model-library root does not load every model or initialize simulation. NeRD’s raw teacher assets, data generation and original-repository comparison runners belong to the learning/nerd showcase, outside the runtime model package.

NeRD public entry points

embodichain_models.dynamics.nerd lazily exposes the causal network, its physical codec, portable checkpoints and shared-SL integration. Inputs are named history tensors; predictions are physical state-increment targets. Model evaluation retains input gradients. Frozen training statistics travel with the network as buffers.

NeRDModelCfg	Network architecture and ordered input schema for NeRD.
NeRDModel	Causal NeRD network with immutable named inputs and portable statistics.
FeatureNormalizer	Frozen mean/variance buffers with the exact official sqrt(var + 1e-5).
export_model(path, model, manifest)	Write tensor weights and JSON schema without executable Python objects.
load_model(path[, device])	Verify and load an exported model; returned model is in evaluation mode.
convert_legacy_checkpoint(checkpoint, ...[, ...])	Convert a trusted official Transformer checkpoint into a portable export.
NeRDStateCodec	Encode Newton child-COM velocities as the official NeRD spatial twist.
NeRDObjective	Preserve the NeRD normalized physical-target loss.
make_nerd_batch(batch, codec, device)	Encode a raw NeRD sequence without changing its sampling contract.
NeRDWindowSource	Adapt an immutable NeRD window reader without owning sampling RNG.
build_supervised_trainer(config)	Build a simulation-independent NeRD recipe for the shared trainer.

class embodichain_models.dynamics.nerd.NeRDModelCfg

Network architecture and ordered input schema for NeRD.

The default architecture matches the official Cartpole checkpoint; input_dims and output_dim must be supplied. Ant uses n_embd=384 and its contact schema. attention_backend="math" persists an explicit SDPA choice in portable exports; default preserves Torch’s backend selection.

Methods:

__init__([input_dims, output_dim, n_embd, ...])
__new__(**kwargs)

__init__(input_dims=<factory>, output_dim=<factory>, n_embd=<factory>, n_layer=<factory>, n_head=<factory>, block_size=<factory>, history_length=<factory>, head_hidden_dims=<factory>, bias=<factory>, dropout=<factory>, normalize_input=<factory>, normalize_output=<factory>, output_tanh=<factory>, attention_backend=<factory>)

__new__(**kwargs)

class embodichain_models.dynamics.nerd.NeRDModel

Causal NeRD network with immutable named inputs and portable statistics.

Parameters:

cfg (NeRDModelCfg) – Explicit input layout and Transformer configuration.

All inputs have shape (batch, time, width). Right-padding is permitted: causal attention preserves prefix outputs, and the caller masks padded loss. Features follow the network’s floating dtype on their existing device; dtype conversion preserves input gradients independently of normalization. evaluate gathers the final valid prefix token independently per row. Neither training nor evaluation implicitly detaches gradients.

Methods:

__init__(cfg)	Initialize internal Module state, shared by both nn.Module and ScriptModule.
__new__(**kwargs)

__init__(cfg)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

__new__(**kwargs)

class embodichain_models.dynamics.nerd.FeatureNormalizer

Frozen mean/variance buffers with the exact official sqrt(var + 1e-5).

Methods:

__init__(width)	Initialize internal Module state, shared by both nn.Module and ScriptModule.
__new__(**kwargs)

__init__(width)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

__new__(**kwargs)

embodichain_models.dynamics.nerd.export_model(path, model, manifest)

Write tensor weights and JSON schema without executable Python objects.

Parameters:

• path (str | Path) – Output directory, which must not already contain an export.

• model (NeRDModel) – Model whose parameters and statistics are exported on CPU.

• manifest (dict[str, Any]) – Deployment metadata, including codec and timestep contracts.

Return type:

None

embodichain_models.dynamics.nerd.load_model(path, device='cpu')

Verify and load an exported model; returned model is in evaluation mode.

Return type:

tuple[NeRDModel, dict[str, Any]]

embodichain_models.dynamics.nerd.convert_legacy_checkpoint(checkpoint, destination, manifest, *, trust_pickle=False)

Convert a trusted official Transformer checkpoint into a portable export.

Parameters:

• checkpoint (str | Path) – Official model.pt, containing a serialized Python model.

• destination (str | Path) – New export directory.

• manifest (dict[str, Any]) – Deployment metadata validated by the eventual provider.

• trust_pickle (bool) – Must be explicitly true: legacy pickle can execute code.

Return type:

None

Only the official identity encoder/ReLU head/Transformer architecture is accepted. Arbitrary architectures must be migrated explicitly.

class embodichain_models.dynamics.nerd.NeRDStateCodec

Encode Newton child-COM velocities as the official NeRD spatial twist.

Parameters:

• robot (str) – cartpole or ant; fixes q/qd layout to 2/2 or 15/14.

• base_com (Sequence[float]) – Root body’s COM offset in its local body frame.

• continuous_q_indices (Sequence[int] | None) – Coordinates wrapped to [-pi, pi). Defaults to the Cartpole pole coordinate and no Ant coordinates.

State tensors may have any leading batch/time dimensions. Ant q uses xyz, xyzw, eight joint coordinates. Native qd uses world COM linear velocity, world angular velocity, eight joint speeds. The legacy spatial twist uses angular velocity then linear velocity referenced to the world origin.

Methods:

__init__([robot, base_com, continuous_q_indices])
__new__(**kwargs)

__init__(robot='cartpole', base_com=(0.0, 0.0, 0.0), continuous_q_indices=None)

__new__(**kwargs)

class embodichain_models.dynamics.nerd.NeRDObjective

Preserve the NeRD normalized physical-target loss.

Parameters:

include_prefix (bool) – Include valid startup positions using the batch mask. False retains the reference full-window MSE reduction order.

Input and target normalization remain owned by the NeRD model and codec. Computing this objective never updates their statistics.

Methods:

__init__(*[, include_prefix])
__new__(**kwargs)

__init__(*, include_prefix=False)

__new__(**kwargs)

embodichain_models.dynamics.nerd.make_nerd_batch(batch, codec, device)

Encode a raw NeRD sequence without changing its sampling contract.

Parameters:

• batch (dict[str, Any]) – Owned trajectory windows from DynamicsWindowDataset.

• codec (Any) – The matching physical state/contact/effort codec.

• device (str | device) – Device on which encoding and model execution take place.

Return type:

SupervisedBatch

Returns:

Encoded history inputs, physical increment targets and validity mask.

class embodichain_models.dynamics.nerd.NeRDWindowSource

Adapt an immutable NeRD window reader without owning sampling RNG.

Parameters:

• dataset (DynamicsWindowDataset) – Caller-owned window dataset, including optional resident cache.

• codec (NeRDStateCodec) – Physical codec matching that dataset.

• device (str | device) – Encoding device, preserving the reference preprocessing policy.

Methods:

__init__(dataset, codec[, device])
__new__(**kwargs)

__init__(dataset, codec, device='cpu')

__new__(**kwargs)

embodichain_models.dynamics.nerd.build_supervised_trainer(config)

Build a simulation-independent NeRD recipe for the shared trainer.

Parameters:

config (dict[str, Any]) – train-sl mapping. model is a NeRD model configuration; data.path identifies committed shards and optional data.statistics supplies shared frozen training statistics. objective accepts include_prefix and numerical_profile. optimizer accepts Adam rates, clipping and interval/offset LR scheduling. trainer uses SupervisedTrainerCfg.

Return type:

SupervisedTrainer

Returns:

Common trainer with a NeRD objective and owned window reader references. Its checkpoints use the common SL format. The NeRD showcase historical trainer continues reading and writing NeRD training checkpoints.

Network and normalization

Ordered input widths and Transformer settings define the network schema. forward predicts every history position; evaluate selects the last valid position for each row. Causal attention preserves valid prefix outputs under right padding; the objective owns the padded loss mask. Feature conversion follows the network’s floating dtype without detaching input gradients.

NeRDModelCfg	Network architecture and ordered input schema for NeRD.
NeRDModel	Causal NeRD network with immutable named inputs and portable statistics.
FeatureNormalizer	Frozen mean/variance buffers with the exact official sqrt(var + 1e-5).

class embodichain_models.dynamics.nerd.model.NeRDModelCfg

Network architecture and ordered input schema for NeRD.

The default architecture matches the official Cartpole checkpoint; input_dims and output_dim must be supplied. Ant uses n_embd=384 and its contact schema. attention_backend="math" persists an explicit SDPA choice in portable exports; default preserves Torch’s backend selection.

Methods:

__init__([input_dims, output_dim, n_embd, ...])
__new__(**kwargs)

__init__(input_dims=<factory>, output_dim=<factory>, n_embd=<factory>, n_layer=<factory>, n_head=<factory>, block_size=<factory>, history_length=<factory>, head_hidden_dims=<factory>, bias=<factory>, dropout=<factory>, normalize_input=<factory>, normalize_output=<factory>, output_tanh=<factory>, attention_backend=<factory>)

__new__(**kwargs)

class embodichain_models.dynamics.nerd.model.NeRDModel

Causal NeRD network with immutable named inputs and portable statistics.

Parameters:

cfg (NeRDModelCfg) – Explicit input layout and Transformer configuration.

All inputs have shape (batch, time, width). Right-padding is permitted: causal attention preserves prefix outputs, and the caller masks padded loss. Features follow the network’s floating dtype on their existing device; dtype conversion preserves input gradients independently of normalization. evaluate gathers the final valid prefix token independently per row. Neither training nor evaluation implicitly detaches gradients.

Methods:

__init__(cfg)	Initialize internal Module state, shared by both nn.Module and ScriptModule.
__new__(**kwargs)
evaluate(inputs[, valid_lengths])	Predict the last valid token, shape (batch, 1, output_dim).
forward(inputs)	Predict every sequence token in physical target units.
set_normalization(input_stats, output_mean, ...)	Set frozen statistics; missing features and invalid variances fail early.

__init__(cfg)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

__new__(**kwargs)

evaluate(inputs, valid_lengths=None)

Predict the last valid token, shape (batch, 1, output_dim).

Return type:

Tensor

forward(inputs)

Predict every sequence token in physical target units.

Return type:

Tensor

set_normalization(input_stats, output_mean, output_var)

Set frozen statistics; missing features and invalid variances fail early.

Return type:

None

class embodichain_models.dynamics.nerd.model.FeatureNormalizer

Frozen mean/variance buffers with the exact official sqrt(var + 1e-5).

Methods:

__init__(width)	Initialize internal Module state, shared by both nn.Module and ScriptModule.
__new__(**kwargs)
forward(value[, inverse])	Normalize or denormalize without modifying the caller's tensor.

__init__(width)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

__new__(**kwargs)

forward(value, inverse=False)

Normalize or denormalize without modifying the caller’s tensor.

Return type:

Tensor

Portable checkpoints

Exports contain tensor weights and JSON deployment metadata, including a weight checksum, numerical profile and attention backend. Loading restores the saved floating precision and returns an evaluation-mode model. Legacy Python-object conversion requires explicit trust and accepts the supported official Transformer architecture.

export_model(path, model, manifest)	Write tensor weights and JSON schema without executable Python objects.
load_model(path[, device])	Verify and load an exported model; returned model is in evaluation mode.
convert_legacy_checkpoint(checkpoint, ...[, ...])	Convert a trusted official Transformer checkpoint into a portable export.

embodichain_models.dynamics.nerd.checkpoint.export_model(path, model, manifest)

Write tensor weights and JSON schema without executable Python objects.

Parameters:

• path (str | Path) – Output directory, which must not already contain an export.

• model (NeRDModel) – Model whose parameters and statistics are exported on CPU.

• manifest (dict[str, Any]) – Deployment metadata, including codec and timestep contracts.

Return type:

None

embodichain_models.dynamics.nerd.checkpoint.load_model(path, device='cpu')

Verify and load an exported model; returned model is in evaluation mode.

Return type:

tuple[NeRDModel, dict[str, Any]]

embodichain_models.dynamics.nerd.checkpoint.convert_legacy_checkpoint(checkpoint, destination, manifest, *, trust_pickle=False)

Convert a trusted official Transformer checkpoint into a portable export.

Parameters:

• checkpoint (str | Path) – Official model.pt, containing a serialized Python model.

• destination (str | Path) – New export directory.

• manifest (dict[str, Any]) – Deployment metadata validated by the eventual provider.

• trust_pickle (bool) – Must be explicitly true: legacy pickle can execute code.

Return type:

None

Only the official identity encoder/ReLU head/Transformer architecture is accepted. Arbitrary architectures must be migrated explicitly.

Physical state, contact and history encoding

The Cartpole and Ant codec converts native Newton COM velocities to NeRD’s angular-first spatial twist and expresses relative targets in the current body frame. Quaternion helpers use xyzw ordering. Contact features use fixed plane slots; history append and reset return new tensors while retaining gradients through unaffected values.

NeRDStateCodec	Encode Newton child-COM velocities as the official NeRD spatial twist.
quat_multiply(a, b)	Multiply broadcastable xyzw quaternions in official NeRD FP32 order.
quat_rotate(q, v)	Rotate vectors by unit xyzw quaternions in official NeRD FP32 order.
quat_inverse(q)	Return the inverse of unit xyzw quaternions.
plane_contact_features(points_world, ...[, ...])	Project fixed slots onto n dot x = offset; normals must be unit length.
append_history(history, lengths, features)	Append one token to left-aligned fixed capacity histories without mutation.
reset_history(history, lengths, mask)	Functionally clear selected rows while retaining other rows' gradients.

class embodichain_models.dynamics.nerd.codec.NeRDStateCodec

Encode Newton child-COM velocities as the official NeRD spatial twist.

Parameters:

• robot (str) – cartpole or ant; fixes q/qd layout to 2/2 or 15/14.

• base_com (Sequence[float]) – Root body’s COM offset in its local body frame.

• continuous_q_indices (Sequence[int] | None) – Coordinates wrapped to [-pi, pi). Defaults to the Cartpole pole coordinate and no Ant coordinates.

State tensors may have any leading batch/time dimensions. Ant q uses xyz, xyzw, eight joint coordinates. Native qd uses world COM linear velocity, world angular velocity, eight joint speeds. The legacy spatial twist uses angular velocity then linear velocity referenced to the world origin.

Methods:

__init__([robot, base_com, continuous_q_indices])
__new__(**kwargs)
decode(q, qd, prediction, *[, current_embedding])	Decode a relative prediction into native Newton q/qd without detaching.
encode(q, qd, effort[, contacts, ...])	Encode inputs at t; contact points/normals are world-frame flattened slots.
legacy_to_native(q, velocity)	Convert angular-first world-origin velocity to native COM velocity.
native_to_legacy(q, qd)	Convert native qd to angular-first world-origin spatial velocity.
targets(q, qd, next_q, next_qd, *[, ...])	Encode relative next state in the frame at t, including quaternion delta.

__init__(robot='cartpole', base_com=(0.0, 0.0, 0.0), continuous_q_indices=None)

__new__(**kwargs)

decode(q, qd, prediction, *, current_embedding=None)

Decode a relative prediction into native Newton q/qd without detaching.

Parameters:

• q (Tensor) – Current native joint positions.

• qd (Tensor) – Current native joint velocities.

• prediction (Tensor) – Relative network prediction in physical target units.

• current_embedding (Tensor | None) – Optional unnormalized states_embedding from encode() for the same state, shared with model inputs.

Return type:

tuple[Tensor, Tensor]

Returns:

Next native joint positions and velocities.

encode(q, qd, effort, contacts=None, gravity_dir=None, *, legacy_velocity=None)

Encode inputs at t; contact points/normals are world-frame flattened slots.

Parameters:

• q (Tensor) – Current native joint positions.

• qd (Tensor) – Native joint velocities, using COM velocity for the Ant free base.

• effort (Tensor) – Generalized efforts, or eight actuated efforts for Ant.

• contacts (Mapping[str, Tensor] | None) – Fixed ordered world-frame contact features for Ant.

• gravity_dir (Tensor | None) – Optional world-frame unit gravity direction.

• legacy_velocity (Tensor | None) – Optional original NeRD velocity in the same world frame as q. It must match qd shape, dtype and device. This preserves imported FP32 values without a COM roundtrip; omitting it retains native velocity conversion. It is an independent input and is never modified or detached.

Return type:

dict[str, Tensor]

Returns:

Named unnormalized model features, retaining their input graphs.

legacy_to_native(q, velocity)

Convert angular-first world-origin velocity to native COM velocity.

Return type:

Tensor

native_to_legacy(q, qd)

Convert native qd to angular-first world-origin spatial velocity.

Return type:

Tensor

targets(q, qd, next_q, next_qd, *, current_embedding=None, legacy_velocity=None, next_legacy_velocity=None)

Encode relative next state in the frame at t, including quaternion delta.

Parameters:

• q (Tensor) – Current native joint positions.

• qd (Tensor) – Current native joint velocities.

• next_q (Tensor) – Next native joint positions.

• next_qd (Tensor) – Next native joint velocities.

• current_embedding (Tensor | None) – Optional unnormalized states_embedding from encode() on exactly the same current state. Reusing it preserves the reference trainer’s single current-frame transform.

• legacy_velocity (Tensor | None) – Optional original NeRD velocity at t, in the same world frame as q and matching qd shape/dtype/device. Omission retains native conversion.

• next_legacy_velocity (Tensor | None) – Optional original NeRD velocity at t+1, in the same world frame as next_q and matching next_qd. It is transformed into the frame at t like the native path.

Return type:

Tensor

Returns:

Relative position/orientation and velocity targets.

embodichain_models.dynamics.nerd.codec.quat_multiply(a, b)

Multiply broadcastable xyzw quaternions in official NeRD FP32 order.

Algebraically equivalent Hamilton-product formulas round differently. The identity orientation in Ant’s body-frame features is sensitive to these differences after normalization, so preserve the reference’s operation order.

Return type:

Tensor

embodichain_models.dynamics.nerd.codec.quat_rotate(q, v)

Rotate vectors by unit xyzw quaternions in official NeRD FP32 order.

Return type:

Tensor

embodichain_models.dynamics.nerd.codec.quat_inverse(q)

Return the inverse of unit xyzw quaternions.

Return type:

Tensor

embodichain_models.dynamics.nerd.codec.plane_contact_features(points_world, thickness0, plane_normal, plane_offset=0.0, thickness1=0.0)

Project fixed slots onto n dot x = offset; normals must be unit length.

Parameters:

• points_world (Tensor) – Candidate center points, shape (…, slots, 3).

• thickness0 (Tensor) – Shape radius plus margin, shape (…, slots).

• plane_normal (Tensor) – Unit normal, broadcastable to candidate points.

• plane_offset (Tensor | float) – Plane displacement along the normal.

• thickness1 (Tensor | float) – Plane margin.

Return type:

dict[str, Tensor]

Returns:

Official contact fields; depth is center-to-plane distance, before thickness.

embodichain_models.dynamics.nerd.codec.append_history(history, lengths, features)

Append one token to left-aligned fixed capacity histories without mutation.

Full rows slide left. Short rows retain position indices starting at zero. Storage is O(batch * capacity * feature_width), with no host synchronization.

Return type:

tuple[dict[str, Tensor], Tensor]

embodichain_models.dynamics.nerd.codec.reset_history(history, lengths, mask)

Functionally clear selected rows while retaining other rows’ gradients.

Return type:

tuple[dict[str, Tensor], Tensor]

Committed trajectories and windows

The single-writer HDF5 format stores T+1 states and T transitions. Only manifest-committed shards are visible to readers. Windows respect trajectory and split boundaries; optional startup windows carry right-padding masks. Readers bound open shard handles and offer an explicitly budgeted resident cache. Operational generation progress is excluded from the data fingerprint.

DynamicsShardWriter	Write trajectories directly to disk; only complete shards are visible.
DynamicsWindowDataset	Read bounded windows with O(trajectories) indexing and LRU file handles.
trajectory_split(trajectory_id[, seed])	Assign a stable trajectory-level 80/10/10 split independent of order.
load_manifest(root, *[, verify])	Read the committed dataset snapshot, optionally checking shard hashes.
dataset_fingerprint(manifest)	Hash immutable data identity, excluding operational generation progress.

class embodichain_models.dynamics.nerd.data.DynamicsShardWriter

Write trajectories directly to disk; only complete shards are visible.

This is a single-writer API. Each caller-owned trajectory is written before add_trajectory returns; no unbounded staging queue is retained. Resume ignores files absent from the manifest, including an interrupted final shard.

Methods:

__init__(root, metadata[, ...])
__new__(**kwargs)
add_trajectory(trajectory_id, states, ...[, ...])	Append one finite trajectory, with T+1 states and T transitions.
close()	Commit the final partial shard and close the writer.
commit()	Publish pending data and optional generation progress atomically.

Attributes:

committed_ids

IDs persisted in the current manifest, excluding pending writes.

__init__(root, metadata, trajectories_per_shard=64, *, resume=False, split_seed=0)

__new__(**kwargs)

add_trajectory(trajectory_id, states, transitions, *, split=None, metadata=None)

Append one finite trajectory, with T+1 states and T transitions.

Return type:

None

close()

Commit the final partial shard and close the writer.

Return type:

None

commit()

Publish pending data and optional generation progress atomically.

A caller may maintain manifest["generation_progress"] alongside pending trajectories. It is published with each automatic shard commit; a rejection-only update also publishes when no shard is pending.

Return type:

None

property committed_ids: frozenset[str]

IDs persisted in the current manifest, excluding pending writes.

class embodichain_models.dynamics.nerd.data.DynamicsWindowDataset

Read bounded windows with O(trajectories) indexing and LRU file handles.

Default full windows match the official sequence dataset. include_prefix also samples startup windows, padded on the right with a validity mask. No window ever spans a trajectory, reset or split boundary.

Methods:

__init__(root, split, history_length, *[, ...])
__new__(**kwargs)
close()	Release resident fields and file handles; returned batches stay owned.
enable_device_cache(device, max_bytes)	Cache raw transitions within an explicit single-process byte budget.
get_batch(indices)	Read an owned batch in caller order using bounded grouped HDF5 reads.

__init__(root, split, history_length, *, include_prefix=False, max_open_shards=4, read_page_steps=256, verify=False)

__new__(**kwargs)

close()

Release resident fields and file handles; returned batches stay owned.

Return type:

None

enable_device_cache(device, max_bytes)

Cache raw transitions within an explicit single-process byte budget.

Parameters:

• device (str | device) – CPU or CUDA device used by this learner.

• max_bytes (int) – Maximum resident raw-field allocation. Insufficient budgets fail before allocating; no data are dropped or quantized.

Return type:

int

Returns:

Actual resident field bytes. The original shard fingerprint and statistical/window weighting remain unchanged.

get_batch(indices)

Read an owned batch in caller order using bounded grouped HDF5 reads.

Requests are grouped by shard, trajectory and a fixed time page. The largest source read is read_page_steps + history_length frames; no trajectory/data cache grows with dataset size. Storage consists of the output batch, O(batch size) indices, one bounded source block and a per-field batched gather at a time. Repeated requests occupy independent rows in the returned tensors. Sampling and prefetch state are not touched by this synchronous method.

Parameters:

indices (Sequence[int] | ndarray | Tensor) – Nonempty CPU integer window indices. Negative indices have the same meaning as __getitem__.

Return type:

dict[str, Any]

Returns:

The same nested tensors/metadata as PyTorch’s default collation of individual windows, including right-padded startup masks.

embodichain_models.dynamics.nerd.data.trajectory_split(trajectory_id, seed=0)

Assign a stable trajectory-level 80/10/10 split independent of order.

Return type:

str

embodichain_models.dynamics.nerd.data.load_manifest(root, *, verify=False)

Read the committed dataset snapshot, optionally checking shard hashes.

Return type:

dict[str, Any]

embodichain_models.dynamics.nerd.data.dataset_fingerprint(manifest)

Hash immutable data identity, excluding operational generation progress.

Manifests written before generation progress was included retain their original fingerprint. Rejection-only commits do not change model inputs.

Return type:

str

Encoding, statistics and evaluation

Raw windows become named feature histories, physical increment targets and a valid-position mask. Statistics use only the training split and weight each valid overlapping-window occurrence; artifacts bind the dataset fingerprint. The normalized MSE uses frozen target variance. Offline evaluation reports valid-position prediction and decoded state errors without selecting on test.

encode_batch(batch, codec, device)	Encode a collated raw trajectory window through the shared codec.
normalized_mse(prediction, target, variance)	Official MSE over normalized physical targets, excluding padded steps.
compute_statistics(dataset, codec[, batch_size])	Compute bounded training-only, overlapping-window-weighted moments.
save_statistics(path, statistics)	Write a tensor-only statistics artifact usable by both trainers.
load_statistics(path)	Load a tensor-only shared reference statistics artifact.
evaluate_model(model, codec, dataset[, ...])	Evaluate all valid held-out positions without model selection on test.

embodichain_models.dynamics.nerd.integration.encode_batch(batch, codec, device)

Encode a collated raw trajectory window through the shared codec.

Return type:

tuple[dict[str, Tensor], Tensor, Tensor]

embodichain_models.dynamics.nerd.integration.normalized_mse(prediction, target, variance, mask=None)

Official MSE over normalized physical targets, excluding padded steps.

Return type:

Tensor

embodichain_models.dynamics.nerd.integration.compute_statistics(dataset, codec, batch_size=64)

Compute bounded training-only, overlapping-window-weighted moments.

Statistical weighting is explicit. Official comparison runs can instead load the exact shared reference artifact; no claim of matching an upstream running estimator’s epsilon initialization is implied by these population moments.

Return type:

dict[str, Any]

embodichain_models.dynamics.nerd.integration.save_statistics(path, statistics)

Write a tensor-only statistics artifact usable by both trainers.

Return type:

None

embodichain_models.dynamics.nerd.integration.load_statistics(path)

Load a tensor-only shared reference statistics artifact.

Return type:

dict[str, Any]

embodichain_models.dynamics.nerd.integration.evaluate_model(model, codec, dataset, batch_size=64, device='cpu', max_batches=None)

Evaluate all valid held-out positions without model selection on test.

Return type:

dict[str, float | int | str]

Shared supervised objective

The adapter preserves the raw reader’s sampling contract when creating a SupervisedBatch. Its objective returns a graph-connected normalized state-increment loss and the number of valid positions used for aggregation. Computing the loss does not update normalization statistics.

NeRDObjective	Preserve the NeRD normalized physical-target loss.
make_nerd_batch(batch, codec, device)	Encode a raw NeRD sequence without changing its sampling contract.

class embodichain_models.dynamics.nerd.supervised.NeRDObjective

Preserve the NeRD normalized physical-target loss.

Parameters:

include_prefix (bool) – Include valid startup positions using the batch mask. False retains the reference full-window MSE reduction order.

Input and target normalization remain owned by the NeRD model and codec. Computing this objective never updates their statistics.

Methods:

__init__(*[, include_prefix])
__new__(**kwargs)
compute_loss(model, batch)	Compute the normalized state-increment error without detaching it.

__init__(*, include_prefix=False)

__new__(**kwargs)

compute_loss(model, batch)

Compute the normalized state-increment error without detaching it.

Parameters:

• model (Module) – NeRD model with its frozen normalization buffers.

• batch (SupervisedBatch) – Encoded inputs, physical targets and valid-position mask.

Return type:

LossResult

Returns:

Graph-connected loss for one shared supervised update.

embodichain_models.dynamics.nerd.supervised.make_nerd_batch(batch, codec, device)

Encode a raw NeRD sequence without changing its sampling contract.

Parameters:

• batch (dict[str, Any]) – Owned trajectory windows from DynamicsWindowDataset.

• codec (Any) – The matching physical state/contact/effort codec.

• device (str | device) – Device on which encoding and model execution take place.

Return type:

SupervisedBatch

Returns:

Encoded history inputs, physical increment targets and validity mask.

Offline SL recipe

NeRDWindowSource uses the trainer’s explicit CPU sampling generator and encodes copied windows on the requested device. The factory combines this source, the physical objective, Adam scheduling and held-out evaluation with the common SL trainer. Its checkpoints use the SL format; historical NeRD training checkpoints remain owned by the showcase’s historical trainer.

Use the installable factory path embodichain_models.dynamics.nerd.sl_recipe:build_supervised_trainer in an embodichain train-sl configuration. It is available independently of the working directory and example-runner bootstrap.

NeRDWindowSource	Adapt an immutable NeRD window reader without owning sampling RNG.
build_supervised_trainer(config)	Build a simulation-independent NeRD recipe for the shared trainer.

class embodichain_models.dynamics.nerd.sl_recipe.NeRDWindowSource

Adapt an immutable NeRD window reader without owning sampling RNG.

Parameters:

• dataset (DynamicsWindowDataset) – Caller-owned window dataset, including optional resident cache.

• codec (NeRDStateCodec) – Physical codec matching that dataset.

• device (str | device) – Encoding device, preserving the reference preprocessing policy.

Methods:

__init__(dataset, codec[, device])
__new__(**kwargs)
sample(batch_size, *, generator)	Sample copied windows using an explicit CPU generator.

__init__(dataset, codec, device='cpu')

__new__(**kwargs)

sample(batch_size, *, generator)

Sample copied windows using an explicit CPU generator.

Parameters:

• batch_size (int) – Positive number of windows sampled with replacement.

• generator (Generator) – Caller-owned sampling generator, committed after updates.

Return type:

SupervisedBatch

Returns:

Encoded inputs, physical targets and valid-position mask.

embodichain_models.dynamics.nerd.sl_recipe.build_supervised_trainer(config)

Build a simulation-independent NeRD recipe for the shared trainer.

Parameters:

config (dict[str, Any]) – train-sl mapping. model is a NeRD model configuration; data.path identifies committed shards and optional data.statistics supplies shared frozen training statistics. objective accepts include_prefix and numerical_profile. optimizer accepts Adam rates, clipping and interval/offset LR scheduling. trainer uses SupervisedTrainerCfg.

Return type:

SupervisedTrainer

Returns:

Common trainer with a NeRD objective and owned window reader references. Its checkpoints use the common SL format. The NeRD showcase historical trainer continues reading and writing NeRD training checkpoints.