diff --git a/agent_context/MAP.yaml b/agent_context/MAP.yaml index fe3fe96781..9f59109509 100644 --- a/agent_context/MAP.yaml +++ b/agent_context/MAP.yaml @@ -548,7 +548,7 @@ topics: - id: supervised-learning title: Supervised Learning aliases: [supervised learning, SL learning, 监督学习, 统一监督学习] - keywords: [train-sl, SupervisedBatch, LossResult, SupervisedObjective, SupervisedLearner, SupervisedTrainer, TensorBatchSource, distributed_run, torchrun, DDP, 单机多卡, 分布式监督学习] + keywords: [train-sl, SupervisedBatch, LossResult, SupervisedObjective, SupervisedLearner, SupervisedTrainer, TensorBatchSource, SequenceBatchSource, distributed_run, torchrun, DDP, 单机多卡, 分布式监督学习] paths: [topics/supervised-learning/supervised-learning.md] source_of_truth: - embodichain/learning/sl/ diff --git a/agent_context/topics/data-pipeline/data-pipeline.md b/agent_context/topics/data-pipeline/data-pipeline.md index 612ba903f2..359a735f25 100644 --- a/agent_context/topics/data-pipeline/data-pipeline.md +++ b/agent_context/topics/data-pipeline/data-pipeline.md @@ -49,6 +49,11 @@ requires explicit action semantics, step duration, field mapping and a resident memory limit; imports do not initialize simulation or import the SDK until load. The first backend samples resident numeric tensors after bounded SDK reads. +`learning/sl/sequence.py` bridges these sources through an explicit encoder; +the [SL context](../supervised-learning/supervised-learning.md) owns its identity +and recovery admission. Online lifecycle and delivery credits remain here. +Neither trainer RNG restoration nor encoding failure rewinds source quotas. + The demo's `lerobot.py` owns writing and sidecar/hash auditing around the public reader. It measures initial loading rather than per-batch disk reads and does not replace recording or offline-inspection tools. See the diff --git a/agent_context/topics/supervised-learning/supervised-learning.md b/agent_context/topics/supervised-learning/supervised-learning.md index 25641fe38d..f42c6819e3 100644 --- a/agent_context/topics/supervised-learning/supervised-learning.md +++ b/agent_context/topics/supervised-learning/supervised-learning.md @@ -14,6 +14,12 @@ operations remain in `compute`; source storage belongs to data adapters. update. Its optional numerical context covers forward, backward and clipping. - `TensorBatchSource` snapshots fixed CPU tensor trees and returns owned samples under the caller's explicit CPU generator. +- Optional `SequenceBatchSource` in `learning/sl/sequence.py` adapts a data-pipeline + `BatchSource` using an explicit versioned encoder. It preserves canonical masks + and provenance under reserved `metadata['sequence_source']`, excluding raw + observations/actions there. Encoders select causal inputs and supervision; + model families retain action/codec/objective semantics. Core SL imports do not + load this bridge, ODS or the SDK. See `examples/learning/sl_data_source/`. - `SupervisedTrainer` in `learning/sl/trainer.py` composes a source and learner, saving optimizer, RNG, sampling progress, source identity and recipe semantics. - `embodichain train-sl --config FILE [--resume FILE]` lazily resolves a @@ -41,6 +47,16 @@ replacement, so samples may overlap. 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. +The sequence bridge defaults to refusing exact resume. Explicitly immutable, +unbudgeted resident sources bind actual CPU content, schema, window, encoder +version and semantic metadata (including language); storage diagnostics do not. +Sampling checks declarations, quota and tensor versions; resume rehashes content +in bounded blocks. An optional source `validate_resume()` preflight runs inside +the collective checkpoint validation before any trainer state or output admission +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, a save uses the previously gathered committed-update snapshot without a new @@ -62,4 +78,5 @@ nonfinite update boundaries, owned samples, simulator-free imports and exact CPU recovery with dropout, optimizer and sampling state. Model-specific qualification belongs to each model's acceptance recipe. The distributed tests exercise weighted global updates and fixed-world-size recovery with CPU/Gloo; -the example is validated by launch smoke runs rather than example-specific tests. +sequence tests reuse real CPU source fixtures for masks, provenance, owned output, +quota and recovery admission. Examples use launch smoke checks, not tutorial tests. diff --git a/docs/source/api_reference/embodichain/embodichain.learning.sl.rst b/docs/source/api_reference/embodichain/embodichain.learning.sl.rst index 6251a7c10c..0d979cbfa1 100644 --- a/docs/source/api_reference/embodichain/embodichain.learning.sl.rst +++ b/docs/source/api_reference/embodichain/embodichain.learning.sl.rst @@ -102,6 +102,7 @@ embodichain.learning.sl SupervisedLearner SupervisedSource TensorBatchSource + SequenceBatchSource SupervisedTrainer SupervisedTrainerCfg distributed_run @@ -124,6 +125,9 @@ embodichain.learning.sl .. autoclass:: TensorBatchSource :members: +.. autoclass:: SequenceBatchSource + :members: + .. autoclass:: SupervisedTrainer :members: @@ -155,6 +159,39 @@ embodichain.learning.sl.learner .. autoclass:: SupervisedLearner :members: +embodichain.learning.sl.sequence +-------------------------------- + +.. currentmodule:: embodichain.learning.sl.sequence + +.. autosummary:: + + SequenceBatchSource + +.. autoclass:: SequenceBatchSource + :members: + +The optional bridge samples one canonical causal window and calls an explicit, +versioned encoder. It preserves original masks and provenance under the reserved +``metadata['sequence_source']`` namespace, excluding observations and actions +from that diagnostic context. The encoder deliberately selects measured inputs +and supervision; the adapter infers neither action meaning nor feature layout. +The runnable ``examples/learning/sl_data_source`` example demonstrates action +targets and causal state-transition targets using a tiny synthetic CPU corpus. + +Exact recovery is disabled by default. ``immutable=True`` requires an unchanged, +content-hashed resident ``TensorDataSource`` (including ``LeRobotDataSource``) +without sampling quotas. Schema, window, encoder version, language and semantic +metadata enter identity; storage paths and read/cache diagnostics do not. +Construction and resume hash CPU content using bounded blocks. Sampling checks +the immutable promise without rehashing the corpus at every update. + +Live or quota-managed sources can train and save snapshots, but exact +``SupervisedTrainer.resume()`` is rejected collectively before learning, RNG or +output admission state changes. Producer state and delivery credits are not +restored by the trainer. Start/stop, refill, waits and handling sampling failures +remain caller-owned; the bridge performs no retry or quota refund. + embodichain.learning.sl.sources ------------------------------- diff --git a/embodichain/learning/sl/__init__.py b/embodichain/learning/sl/__init__.py index 7bbe8974b3..468b0a89ec 100644 --- a/embodichain/learning/sl/__init__.py +++ b/embodichain/learning/sl/__init__.py @@ -18,12 +18,18 @@ from __future__ import annotations +from importlib import import_module +from typing import TYPE_CHECKING, Any + from .types import LossResult, SupervisedBatch, SupervisedObjective from .distributed import distributed_run from .learner import SupervisedLearner from .sources import SupervisedSource, TensorBatchSource from .trainer import SupervisedTrainer, SupervisedTrainerCfg +if TYPE_CHECKING: + from .sequence import SequenceBatchSource + __all__ = [ "LossResult", "SupervisedBatch", @@ -34,4 +40,17 @@ "SupervisedTrainer", "SupervisedTrainerCfg", "distributed_run", + "SequenceBatchSource", ] + + +def __getattr__(name: str) -> Any: + if name != "SequenceBatchSource": + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + value = getattr(import_module(f"{__name__}.sequence"), name) + globals()[name] = value + return value + + +def __dir__() -> list[str]: + return sorted(set(globals()) | set(__all__)) diff --git a/embodichain/learning/sl/sequence.py b/embodichain/learning/sl/sequence.py new file mode 100644 index 0000000000..e058b6c16a --- /dev/null +++ b/embodichain/learning/sl/sequence.py @@ -0,0 +1,261 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Optional causal-data bridge with explicit encoding and recovery capability.""" + +from __future__ import annotations + +from collections.abc import Callable, Mapping +import copy +import hashlib +import json +from typing import Any + +import torch +from tensordict import TensorDict + +from embodichain.data_pipeline.datasets.schema import data_schema_fingerprint +from embodichain.data_pipeline.datasets.sequence import BatchSource, SequenceSpecCfg +from embodichain.data_pipeline.datasets.sources import TensorDataSource + +from .types import SupervisedBatch + +__all__ = ["SequenceBatchSource"] + + +def _json(value: Any) -> str: + # Normalize JSON's integer mapping keys (notably language dictionaries) + # before sorting; storage diagnostics do not enter this representation. + return json.dumps( + json.loads(json.dumps(value, allow_nan=False)), + sort_keys=True, + separators=(",", ":"), + allow_nan=False, + ) + + +def _hash_tensor(digest: Any, value: torch.Tensor) -> None: + """Visit logical C-order bytes with at most 1 MiB of contiguous staging.""" + limit = 1024 * 1024 + if value.numel() * value.element_size() <= limit: + chunk = value.detach().contiguous().reshape(-1).view(torch.uint8) + digest.update(memoryview(chunk.numpy())) + return + row_bytes = value[0].numel() * value.element_size() + if row_bytes > limit: + for index in range(value.shape[0]): + _hash_tensor(digest, value[index]) + else: + count = max(1, limit // row_bytes) + for start in range(0, value.shape[0], count): + _hash_tensor(digest, value[start : start + count]) + + +class SequenceBatchSource: + """Encode one caller-owned causal source without managing its production. + + Args: + source: Live or resident ``BatchSource``. It owns sampling, continuity, + padding, delivery quotas and copied window storage. + sequence: Frozen history/future, segment and generation selection. + encoder: Explicit pure function receiving one owned canonical context + and returning ``SupervisedBatch``. Select measured inputs, targets + and objective masks deliberately; future observations are never + assigned to inputs by this adapter. Returned tensors must belong + to this context or newly allocated storage, not a reused cache. + encoder_version: Persisted identity of the encoding, including semantic + checks, units, feature order and any captured configuration. Callable + code is not inspected; the caller must version changes truthfully. + immutable: Explicit promise that resident data and semantic metadata + remain unchanged. Exact recovery additionally requires a + ``TensorDataSource`` (including resident ``LeRobotDataSource``) + without delivery quotas. The default refuses exact recovery. + + The reserved ``metadata['sequence_source']`` contains source identity, + schema, sequence and ``context`` (canonical validity/loss masks and all + provenance, excluding observations/actions). Optional ``storage`` is copied + for diagnostics and excluded from identity. Encoder metadata cannot use + this reserved key. Non-storage source metadata must be finite JSON data. + No model/action meaning is inferred from tensor widths. + + Construction hashes immutable CPU payloads once, without a corpus copy, + using bounded staging. Sampling checks tensor versions/structure and + semantic declarations; recovery rehashes content. Do not mutate resident + storage, including through aliases or untracked ``.data`` writes. Online + buffers and delivery credits are not checkpointed: live/budgeted runs can + train and save optimizer snapshots, but :meth:`validate_resume` rejects + exact continuation. Start/stop, refill, waits and error recovery remain + caller-owned. A delivered context is not refunded if encoding/update fails. + """ + + def __init__( + self, + source: BatchSource, + sequence: SequenceSpecCfg, + *, + encoder: Callable[[TensorDict], SupervisedBatch], + encoder_version: str, + immutable: bool = False, + ) -> None: + if not callable(encoder): + raise TypeError("Sequence encoding requires an explicit callable") + if not isinstance(encoder_version, str) or not encoder_version.strip(): + raise ValueError("encoder_version must be a nonempty persisted identity") + if type(immutable) is not bool: + raise ValueError("immutable must be an explicit boolean") + if not isinstance(source.source_id, str) or not source.source_id: + raise ValueError("Sequence source requires a nonempty source_id") + self.source = source + self._sequence = copy.deepcopy(sequence) + self._sequence.validate() + self._sequence._check() + self._encoder = encoder + self._encoder_version = encoder_version + self._immutable = immutable + self._schema = copy.deepcopy(dict(source.schema)) + self._declaration = self._describe() + self._content: str | None = None + self._versions: list[Any] = [] + if immutable: + if not isinstance(source, TensorDataSource): + raise ValueError( + "Immutable recovery requires a resident TensorDataSource" + ) + self._check_quota() + self._content = self._hash_content() + self._versions = self._tensor_versions() + self.source_id = ( + "sequence:" + + hashlib.sha256( + _json([self._declaration, self._content]).encode() + ).hexdigest() + ) + + def _describe(self) -> str: + metadata = getattr(self.source, "metadata", {}) + if not isinstance(metadata, Mapping): + raise ValueError("Sequence source metadata must be a mapping") + return _json( + { + "format": "supervised-sequence-v1", + "source_type": f"{type(self.source).__module__}.{type(self.source).__qualname__}", + "source_id": self.source.source_id, + "schema": data_schema_fingerprint(self.source.schema), + "sequence": self._sequence.to_dict(), + "encoder_version": self._encoder_version, + "immutable": self._immutable, + "metadata": { + key: value for key, value in metadata.items() if key != "storage" + }, + } + ) + + def _leaves(self) -> list[tuple[Any, torch.Tensor]]: + assert isinstance(self.source, TensorDataSource) + return sorted( + self.source.data.items(include_nested=True, leaves_only=True), + key=lambda item: (item[0],) if isinstance(item[0], str) else item[0], + ) + + def _hash_content(self) -> str: + assert isinstance(self.source, TensorDataSource) + digest = hashlib.sha256(_json(list(self.source.data.batch_size)).encode()) + for key, value in self._leaves(): + if value.device.type != "cpu" or value.layout != torch.strided: + raise ValueError("Immutable sequence data must be strided CPU tensors") + digest.update(_json([key, list(value.shape), str(value.dtype)]).encode()) + _hash_tensor(digest, value) + return digest.hexdigest() + + def _tensor_versions(self) -> list[Any]: + assert isinstance(self.source, TensorDataSource) + return [tuple(self.source.data.batch_size)] + [ + ( + key, + tuple(value.shape), + value.dtype, + value.stride(), + value.data_ptr(), + value._version, + ) + for key, value in self._leaves() + ] + + def _check_quota(self) -> None: + assert isinstance(self.source, TensorDataSource) + if self.source.get_sampling_budget_snapshot()["enabled"]: + raise ValueError("Immutable exact recovery requires an unbudgeted source") + + def _check_source(self) -> None: + if self._describe() != self._declaration: + raise ValueError( + "Sequence source identity, schema or semantic metadata changed" + ) + if self._immutable: + self._check_quota() + if self._tensor_versions() != self._versions: + raise ValueError("Immutable sequence source tensors changed") + + def sample(self, batch_size: int, *, generator: torch.Generator) -> SupervisedBatch: + """Deliver one encoded owned window batch without retries or refunds. + + Args: + batch_size: Number of contexts requested from the underlying source. + generator: Explicit caller-owned CPU sampling generator. + + Returns: + Encoded input/target/mask trees and aligned host-side source context. + """ + self._check_source() + context = self.source.sample_context( + batch_size, copy.deepcopy(self._sequence), generator=generator + ) + encoded = self._encoder(context) + if not isinstance(encoded, SupervisedBatch): + raise TypeError("Sequence encoder must return SupervisedBatch") + if "sequence_source" in encoded.metadata: + raise ValueError("Encoder metadata cannot use reserved sequence_source") + metadata = { + "source_id": self.source_id, + "schema": copy.deepcopy(self._schema), + "sequence": self._sequence.to_dict(), + "context": context.exclude("obs", "actions"), + } + storage = getattr(self.source, "metadata", {}).get("storage") + if storage is not None: + metadata["storage"] = copy.deepcopy(storage) + return SupervisedBatch( + encoded.inputs, + encoded.targets, + encoded.masks, + {**encoded.metadata, "sequence_source": metadata}, + ) + + def validate_resume(self) -> None: + """Reject non-reproducible source recovery before trainer state changes. + + Raises: + ValueError: Source is live/unverified/budgeted, its declarations or + resident payload changed, or its immutable promise was broken. + """ + if not self._immutable: + raise ValueError( + "Exact resume requires an explicitly immutable unbudgeted resident source; " + "live or quota-managed delivery state is not restored" + ) + self._check_source() + if self._hash_content() != self._content: + raise ValueError("Immutable sequence source content changed") diff --git a/embodichain/learning/sl/sources.py b/embodichain/learning/sl/sources.py index 1e2c3259d3..7d9559ab05 100644 --- a/embodichain/learning/sl/sources.py +++ b/embodichain/learning/sl/sources.py @@ -30,7 +30,13 @@ class SupervisedSource(Protocol): - """Owned batches drawn with an explicit caller-owned CPU sampling RNG.""" + """Owned batches drawn with an explicit caller-owned CPU sampling RNG. + + Sources with additional recovery constraints may provide + ``validate_resume() -> None``. The trainer calls this optional preflight + collectively before changing learning, RNG, sampler or output state. + Raising rejects exact recovery; saving a training snapshot remains allowed. + """ source_id: str diff --git a/embodichain/learning/sl/trainer.py b/embodichain/learning/sl/trainer.py index 51069a65aa..ede0a9f169 100644 --- a/embodichain/learning/sl/trainer.py +++ b/embodichain/learning/sl/trainer.py @@ -396,6 +396,12 @@ def restore() -> None: self._artifacts.adopt_resume(run_id) 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 + # recoverable merely by restoring the trainer's sampling generator. + validate_source = getattr(self.source, "validate_resume", None) + if validate_source is not None: + validate_source() distributed = state.get("distributed", {}) if distributed.get("world_size", 1) != self._context.world_size: raise ValueError( diff --git a/examples/README.md b/examples/README.md index 11c91c7d0a..d720762027 100644 --- a/examples/README.md +++ b/examples/README.md @@ -39,6 +39,10 @@ training and reference experiments. Reusable models and codecs are installed as `embodichain_models`; downloaded datasets, weights and raw reference assets are stored outside the repository and wheel. +The [SL data-source example](learning/sl_data_source/README.md) connects resident +causal windows to explicit action or state-transition encoders, and describes +caller-owned online sampling and the limits of exact recovery. + The [RWM fitting example](learning/rwm/README.md) qualifies controlled supervised state-transition fitting and state/action gradients. Its generated recipes use the installable `embodichain_models.dynamics.rwm` model family and shared SL diff --git a/examples/learning/sl_data_source/README.md b/examples/learning/sl_data_source/README.md new file mode 100644 index 0000000000..9392372cb2 --- /dev/null +++ b/examples/learning/sl_data_source/README.md @@ -0,0 +1,62 @@ +# Causal data sources and supervised learning + +This CPU-only example sends a resident `TensorDataSource` through +`SequenceBatchSource` into the common SL trainer. It needs no simulator, SDK, +dataset download or producer process. Two explicit encoders illustrate the +interface: + +```bash +embodichain run-example learning/sl_data_source --encoder bc --output outputs/sl_data_bc +embodichain run-example learning/sl_data_source --encoder transition --output outputs/sl_data_transition +``` + +The `bc` encoder uses two measured observations as inputs and two current/future +actions as targets. The transition encoder uses measured history and the current +drive as inputs; the next observation is a target only. Both consume the +source's padding/loss masks explicitly. This is a synthetic dimensionless drive +contract and a tiny MLP, not a qualified robot BC policy or NeRD effort recipe. +Recipes must validate real action representations, units, timing and frames; +position commands cannot be relabeled as generalized efforts. + +Every encoder is an explicit pure function returning `SupervisedBatch`. Give it +a persisted `encoder_version` that changes with its behavior or captured +configuration. Keep returned tensors in the current owned context or newly +allocated storage. Do not reuse mutable encoding buffers. The adapter reserves +`metadata['sequence_source']` for the schema, frozen window selection, identity +and original validity/loss masks plus provenance; that context omits observations +and actions. Only the encoder selects model inputs and targets. + +The example sets `immutable=True`: the resident CPU payload and semantic +metadata must never be mutated, and delivery quotas must stay disabled. +Construction hashes actual content in bounded blocks; exact resume also +revalidates content, schema, window and encoder identity. Reusing a handwritten +dataset ID does not authorize different data. Storage paths, cache files and +read-batch diagnostics do not change semantic identity. Language dictionaries +and non-storage metadata do. Use the same update budget and recipe when resuming +an interrupted run from its latest `last.pt` with `--resume FILE`. + +For an already-started live ODS engine, the caller can construct the same bridge: + +```python +from embodichain.data_pipeline.datasets import OnlineDataSource, SequenceSpecCfg +from embodichain.learning.sl import SequenceBatchSource + +source = OnlineDataSource(started_engine, schema=explicit_schema) +training_source = SequenceBatchSource( + source, + SequenceSpecCfg(history_length=2, prediction_horizon=2), + encoder=encode_measured_inputs_and_targets, + encoder_version="my-explicit-robot-recipe-v1", +) +# Pass training_source to SupervisedTrainer. The caller still owns the engine. +``` + +The caller owns start/stop, bounded waits, refill and producer recovery. The +bridge performs one sampling call without retries. Quotas charge successful +delivery, so an encoding or optimizer failure does not refund them. Live buffers +and budgeted resident sources can train and save model/optimizer snapshots, but +`trainer.resume()` rejects them before changing model, optimizer, sampler, RNG +or output admission state. Restoring RNG cannot restore buffer contents or +delivery credits. A fresh run may explicitly initialize model weights from a +trusted snapshot; that is not exact continuation. The same resident bridge also +accepts `LeRobotDataSource` after its caller-owned bounded SDK load. diff --git a/examples/learning/sl_data_source/__main__.py b/examples/learning/sl_data_source/__main__.py new file mode 100644 index 0000000000..93e4aeb27d --- /dev/null +++ b/examples/learning/sl_data_source/__main__.py @@ -0,0 +1,153 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Two explicit toy encoders from resident causal data to the common SL trainer.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +import torch +from tensordict import TensorDict + +from embodichain.data_pipeline.datasets import SequenceSpecCfg, TensorDataSource +from embodichain.learning.sl import ( + LossResult, + SequenceBatchSource, + SupervisedBatch, + SupervisedLearner, + SupervisedTrainer, + SupervisedTrainerCfg, +) + +__all__ = ["main"] + + +class _MaskedRegression: + def compute_loss( + self, model: torch.nn.Module, batch: SupervisedBatch + ) -> LossResult: + error = (model(batch.inputs) - batch.targets).square() + valid = batch.masks["target"].expand_as(error) + loss = error[valid].mean() + return LossResult(loss, weight=int(valid.sum())) + + +def main() -> None: + """Fit actions or a causal transition without a simulator or data worker.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--encoder", choices=("bc", "transition"), default="bc") + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--updates", type=int, default=20) + parser.add_argument("--resume", type=Path) + args = parser.parse_args() + generator = torch.Generator(device="cpu").manual_seed(17) + drive = torch.randn(8, 17, 1, generator=generator) + state = torch.zeros(8, 17, 2) + state[:, 0] = torch.randn(8, 2, generator=generator) + for time in range(16): + velocity = state[:, time, 1] + 0.1 * drive[:, time, 0] + state[:, time + 1] = torch.stack( + (state[:, time, 0] + 0.1 * velocity, velocity), -1 + ) + schema = { + "version": "toy-drive-v1", + "observation": { + "names": ["position", "velocity"], + "timing": "pre_action", + "dtype": "float32", + }, + "action": { + "names": ["drive"], + "semantics": "dimensionless_toy_drive", + "step_dt": 0.1, + }, + } + source = TensorDataSource( + TensorDict( + { + "obs": state, + "actions": drive, + "valid": torch.ones(8, 17, dtype=torch.bool), + }, + [8, 17], + ), + schema=schema, + source_id="synthetic-drive-v1", + ) + # Real recipes validate their own command/effort/pose, units and frames. + # This toy action contract is neither robot joint positions nor NeRD efforts. + if source.schema["action"]["semantics"] != "dimensionless_toy_drive": + raise ValueError("Toy recipe requires its explicit drive semantics") + sequence = SequenceSpecCfg(history_length=2, prediction_horizon=2) + + def encode_bc(context: TensorDict) -> SupervisedBatch: + # Measured history ends at index 1. Actions at 1 and 2 are supervision. + return SupervisedBatch( + context["obs"][:, :2].flatten(1), + context["actions"][:, 1:, 0], + {"target": context["loss_mask"][:, 1:]}, + ) + + def encode_transition(context: TensorDict) -> SupervisedBatch: + # Future observation 2 is a target only. The current drive is supplied + # as an input because this recipe predicts its next physical state. + return SupervisedBatch( + torch.cat((context["obs"][:, :2].flatten(1), context["actions"][:, 1]), -1), + context["obs"][:, 2], + {"target": (context["valid"][:, 1] & context["loss_mask"][:, 2])[:, None]}, + ) + + bridge = SequenceBatchSource( + source, + sequence, + encoder=encode_bc if args.encoder == "bc" else encode_transition, + encoder_version=f"toy-{args.encoder}-v1", + immutable=True, + ) + torch.manual_seed(29) + model = torch.nn.Sequential( + torch.nn.Linear(4 if args.encoder == "bc" else 5, 16), + torch.nn.Tanh(), + torch.nn.Linear(16, 2), + ) + learner = SupervisedLearner( + model, _MaskedRegression(), torch.optim.Adam(model.parameters(), lr=0.001) + ) + trainer = SupervisedTrainer( + learner, + bridge, + SupervisedTrainerCfg( + total_updates=args.updates, + batch_size=8, + output_dir=str(args.output), + checkpoint_interval=0, + ), + signature={ + "recipe": "sequence-source-toy-v1", + "encoder": args.encoder, + "action_semantics": schema["action"]["semantics"], + }, + ) + if args.resume is not None: + trainer.resume(args.resume) + print(json.dumps(trainer.fit(), allow_nan=False)) + + +if __name__ == "__main__": + main() diff --git a/tests/data_pipeline/test_lerobot_source.py b/tests/data_pipeline/test_lerobot_source.py index d9ca5a64c1..0630221d9d 100644 --- a/tests/data_pipeline/test_lerobot_source.py +++ b/tests/data_pipeline/test_lerobot_source.py @@ -165,6 +165,8 @@ def test_padding_capacity_cannot_exceed_resident_allowance(official_dataset): def test_fixed_capacity_and_read_chunking_preserve_numeric_results(official_dataset): + from embodichain.learning.sl import SequenceBatchSource, SupervisedBatch + mapping = LeRobotFieldMappingCfg(provenance_key="source.trace") first = _load( official_dataset, mapping=mapping, max_episode_steps=8, read_batch_size=1 @@ -178,6 +180,36 @@ def test_fixed_capacity_and_read_chunking_preserve_numeric_results(official_data ) assert first.source_id == second.source_id + def encode(context): + return SupervisedBatch( + context["obs"][:, :2], + context["actions"][:, 1:], + {"target": context["loss_mask"][:, 1:]}, + ) + + bridges = [ + SequenceBatchSource( + source, + SequenceSpecCfg(2, 4), + encoder=encode, + encoder_version="sdk-toy-actions-v1", + immutable=True, + ) + for source in (first, second) + ] + assert bridges[0].source_id == bridges[1].source_id + batches = [ + bridge.sample(16, generator=torch.Generator().manual_seed(13)) + for bridge in bridges + ] + torch.testing.assert_close(batches[0].inputs, batches[1].inputs) + torch.testing.assert_close(batches[0].targets, batches[1].targets) + for key in ("valid", "loss_mask", "frame_ids", "generations"): + torch.testing.assert_close( + batches[0].metadata["sequence_source"]["context"][key], + batches[1].metadata["sequence_source"]["context"][key], + ) + def test_ordered_feature_names_are_checked(official_dataset): schema = copy.deepcopy(_SCHEMA) diff --git a/tests/learning/sl/test_distributed.py b/tests/learning/sl/test_distributed.py index bf6dda86c6..a3b0d62e17 100644 --- a/tests/learning/sl/test_distributed.py +++ b/tests/learning/sl/test_distributed.py @@ -396,6 +396,29 @@ def build(output: Path) -> SupervisedTrainer: assert saved["learner"]["num_updates"] == 3 +def _source_resume_worker(rank: int, directory: Path) -> None: + from .test_sequence import _build_trainer + + original = _build_trainer(directory / "original") + original.fit(stop_after=2) + for boundary in ("unverified", "quota"): + candidate = _build_trainer( + directory / boundary, immutable=not (rank == 1 and boundary == "unverified") + ) + if rank == 1 and boundary == "quota": + candidate.source.source.configure_sampling_budget(5, 5) + before = _snapshot(candidate) + run_id = candidate._artifacts.run_id + with pytest.raises(RuntimeError, match="immutable|unbudgeted"): + candidate.resume(original.output / "last.pt") + _assert_tree_equal(_snapshot(candidate), before) + assert not candidate.output.exists() + assert not candidate._distributed_prepared + assert candidate._artifacts.run_id == run_id + assert not candidate._artifacts._claimed + dist.barrier() + + def _distributed_worker( rank: int, group_file: str, directory: str, scenario: str ) -> None: @@ -414,6 +437,8 @@ def _distributed_worker( _failure_worker(rank, Path(directory)) elif scenario == "artifacts": _artifact_worker(rank, Path(directory)) + elif scenario == "source-resume": + _source_resume_worker(rank, Path(directory)) else: _recovery_worker(rank, Path(directory), scenario) finally: @@ -482,6 +507,12 @@ def test_output_identity_and_history_admission_reaches_all_ranks( _spawn("artifacts", tmp_path) +def test_source_resume_preflight_rejection_reaches_all_ranks_without_mutation( + tmp_path: Path, +) -> None: + _spawn("source-resume", tmp_path) + + def test_torchrun_cli_initializes_cleans_up_and_resumes(tmp_path: Path) -> None: output = tmp_path / "run" config = { diff --git a/tests/learning/sl/test_sequence.py b/tests/learning/sl/test_sequence.py new file mode 100644 index 0000000000..1e128fbf33 --- /dev/null +++ b/tests/learning/sl/test_sequence.py @@ -0,0 +1,372 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Owned causal encoding and source-aware exact recovery without a producer.""" + +from __future__ import annotations + +import copy +import importlib.util +from pathlib import Path + +import pytest +import torch + +from embodichain.data_pipeline.datasets import ( + DataBudgetExhaustedError, + DataUnavailableError, + SequenceSpecCfg, + TensorDataSource, +) +from embodichain.learning.sl import ( + SequenceBatchSource, + SupervisedBatch, + SupervisedLearner, + SupervisedTrainer, + SupervisedTrainerCfg, +) +from .test_distributed import _assert_tree_equal, _snapshot, _Regression + +# The repository's tests/ tree is not a Python package. Load the existing +# simulator-free source fixture by its path without changing global sys.path. +_fixture_spec = importlib.util.spec_from_file_location( + "_sl_causal_source_fixtures", + Path(__file__).resolve().parents[2] / "data_pipeline/test_sources.py", +) +assert _fixture_spec is not None and _fixture_spec.loader is not None +_fixtures = importlib.util.module_from_spec(_fixture_spec) +_fixture_spec.loader.exec_module(_fixtures) +_data, _SCHEMA, paired_sources = ( + _fixtures._data, + _fixtures._SCHEMA, + _fixtures.paired_sources, +) + + +def _encode(context): + return SupervisedBatch( + context["obs"][:, :2], + context["actions"][:, 1:], + {"target": context["loss_mask"][:, 1:]}, + {"recipe": "explicit-history-to-actions"}, + ) + + +def _fit_encode(context): + return SupervisedBatch(context["obs"][:, 1], context["actions"][:, 1]) + + +def _resident(*, schema=None, data=None, metadata=None): + return TensorDataSource( + _data() if data is None else data, + schema=copy.deepcopy(_SCHEMA if schema is None else schema), + source_id="same-manual-dataset-id", + metadata=metadata, + ) + + +def _build_trainer( + path: Path, + *, + source=None, + encoder_version="toy-current-regression-v1", + sequence=None, + immutable=True, +): + torch.manual_seed(42) + model = torch.nn.Sequential( + torch.nn.Linear(2, 8), torch.nn.Dropout(0.2), torch.nn.Linear(8, 1) + ) + bridge = SequenceBatchSource( + _resident() if source is None else source, + SequenceSpecCfg(2, 2) if sequence is None else sequence, + encoder=_fit_encode, + encoder_version=encoder_version, + immutable=immutable, + ) + return SupervisedTrainer( + SupervisedLearner( + model, _Regression(), torch.optim.Adam(model.parameters(), lr=0.01) + ), + bridge, + SupervisedTrainerCfg( + total_updates=4, batch_size=3, output_dir=str(path), checkpoint_interval=0 + ), + signature={"semantics": "toy-regression-only"}, + ) + + +@pytest.mark.parametrize("mode,generations", [("episode", None), ("segment", (2,))]) +def test_live_and_resident_encoding_preserves_windows_masks_and_owned_provenance( + paired_sources, mode, generations +): + live, resident = paired_sources + sequence = SequenceSpecCfg(2, 4, mode, generations) + bridges = [ + SequenceBatchSource( + source, sequence, encoder=_encode, encoder_version="toy-actions-v1" + ) + for source in (live, resident) + ] + sequence.prediction_horizon = 1 + expected = resident.sample_context( + 32, + SequenceSpecCfg(2, 4, mode, generations), + generator=torch.Generator().manual_seed(91), + ) + outputs = [ + bridge.sample(32, generator=torch.Generator().manual_seed(91)) + for bridge in bridges + ] + for output in outputs: + torch.testing.assert_close(output.inputs, expected["obs"][:, :2]) + torch.testing.assert_close(output.targets, expected["actions"][:, 1:]) + torch.testing.assert_close(output.masks["target"], expected["loss_mask"][:, 1:]) + metadata = output.metadata["sequence_source"] + assert metadata["sequence"]["prediction_horizon"] == 4 + assert metadata["schema"] == _SCHEMA + assert "obs" not in metadata["context"] and "actions" not in metadata["context"] + for key in expected.exclude("obs", "actions").keys(): + torch.testing.assert_close(metadata["context"][key], expected[key]) + saved = outputs[0].inputs.clone() + bridges[0].sample(32, generator=torch.Generator().manual_seed(13)).inputs.fill_( + -1000 + ) + torch.testing.assert_close(outputs[0].inputs, saved) + outputs[0].inputs.fill_(-1000) + outputs[1].targets.fill_(-1000) + outputs[0].metadata["sequence_source"]["context"]["frame_ids"].fill_(-1000) + assert not (live.engine.shared_buffer["obs"] == -1000).any() + assert not (live.engine.shared_buffer["ods_frame_index"] == -1000).any() + assert not (resident.data["actions"] == -1000).any() + + +@pytest.mark.parametrize( + "error_type", [DataUnavailableError, DataBudgetExhaustedError, RuntimeError] +) +def test_source_errors_are_propagated_once_without_retry(monkeypatch, error_type): + source = _resident() + bridge = SequenceBatchSource( + source, SequenceSpecCfg(2, 2), encoder=_encode, encoder_version="v1" + ) + error = error_type("terminal delivery failure") + calls = [] + + def fail(*args, **kwargs): + calls.append(1) + raise error + + monkeypatch.setattr(source, "sample_context", fail) + with pytest.raises(error_type) as captured: + bridge.sample(2, generator=torch.Generator()) + assert captured.value is error and calls == [1] + + +@pytest.mark.parametrize("reserved", [False, True]) +def test_encoding_failure_does_not_refund_delivery_quota(reserved): + source = _resident() + source.configure_sampling_budget(2, 2) + calls = [] + + def encode(context): + calls.append(1) + if reserved: + return SupervisedBatch( + context["obs"], context["actions"], metadata={"sequence_source": {}} + ) + raise ValueError("encoding failed") + + bridge = SequenceBatchSource( + source, SequenceSpecCfg(2, 2), encoder=encode, encoder_version="failing-v1" + ) + with pytest.raises( + ValueError, match="reserved sequence_source" if reserved else "encoding failed" + ): + bridge.sample(3, generator=torch.Generator().manual_seed(5)) + assert calls == [1] + assert source.get_sampling_budget_snapshot()["total_delivered_chunks"] == 3 + with pytest.raises(DataBudgetExhaustedError): + bridge.sample(2, generator=torch.Generator().manual_seed(5)) + assert calls == [1] + assert source.get_sampling_budget_snapshot()["total_delivered_chunks"] == 3 + + +@pytest.mark.parametrize( + "changed", ["schema", "encoder", "sequence", "corpus", "language"] +) +def test_resume_identity_changes_fail_before_mutating_any_trainer_state( + tmp_path, changed +): + metadata = {"language": {"task": {0: "pick"}}} + original = _build_trainer( + tmp_path / "original", source=_resident(metadata=metadata) + ) + original.fit(stop_after=2) + schema, data = copy.deepcopy(_SCHEMA), _data() + kwargs = {} + if changed == "schema": + schema["action"]["semantics"] = "joint_position" + elif changed == "encoder": + kwargs["encoder_version"] = "changed-v2" + elif changed == "sequence": + kwargs["sequence"] = SequenceSpecCfg(2, 3) + elif changed == "corpus": + data["actions"][0, 0] += 1 + else: + metadata["language"]["task"][0] = "release" + candidate = _build_trainer( + tmp_path / "candidate", + source=_resident(schema=schema, data=data, metadata=metadata), + **kwargs, + ) + before = _snapshot(candidate) + run_id = candidate._artifacts.run_id + with pytest.raises(ValueError, match="identity"): + candidate.resume(original.output / "last.pt") + _assert_tree_equal(_snapshot(candidate), before) + assert not candidate.output.exists() + assert candidate._artifacts.run_id == run_id + assert not candidate._artifacts._claimed + + +def test_storage_diagnostics_can_change_without_changing_exact_recovery(tmp_path): + def source(storage): + return _resident( + metadata={"storage": storage, "language": {"task": {0: "pick"}}} + ) + + original = _build_trainer( + tmp_path / "original", + source=source({"root": "/old", "numeric_read_batch_size": 1}), + ) + original.fit(stop_after=2) + candidate = _build_trainer( + tmp_path / "candidate", + source=source({"root": "/moved", "numeric_read_batch_size": 19}), + ) + assert original.source.source_id == candidate.source.source_id + candidate.resume(original.output / "last.pt") + _assert_tree_equal(_snapshot(candidate), _snapshot(original)) + candidate.fit() + original.resume(original.output / "last.pt") + original.fit() + _assert_tree_equal(_snapshot(candidate), _snapshot(original)) + diagnostic = candidate.source.sample(1, generator=torch.Generator()).metadata[ + "sequence_source" + ]["storage"] + assert diagnostic["root"] == "/moved" + + +@pytest.mark.parametrize("kind", ["live", "quota", "unverified"]) +def test_live_quota_and_unverified_snapshots_reject_resume_without_mutation( + tmp_path, paired_sources, kind +): + live, _ = paired_sources + source = live if kind == "live" else _resident() + if kind == "quota": + source.configure_sampling_budget(10, 10) + trainer = _build_trainer(tmp_path / kind, source=source, immutable=False) + trainer.fit(stop_after=1) + before = _snapshot(trainer) + files = {path.name: path.read_bytes() for path in trainer.output.iterdir()} + quota = source.get_sampling_budget_snapshot() if kind == "quota" else None + with pytest.raises(ValueError, match="Exact resume.*immutable"): + trainer.resume(trainer.output / "last.pt") + _assert_tree_equal(_snapshot(trainer), before) + assert files == {path.name: path.read_bytes() for path in trainer.output.iterdir()} + if quota is not None: + assert source.get_sampling_budget_snapshot() == quota + + +@pytest.mark.parametrize( + "changed", ["quota", "source_id", "schema", "language", "tensor"] +) +def test_immutable_promise_is_checked_after_construction(changed): + source = _resident(metadata={"language": {"task": {0: "pick"}}}) + bridge = SequenceBatchSource( + source, + SequenceSpecCfg(2, 2), + encoder=_fit_encode, + encoder_version="v1", + immutable=True, + ) + if changed == "quota": + source.configure_sampling_budget(2, 2) + elif changed == "source_id": + source.source_id = "changed" + elif changed == "schema": + source.schema["action"]["step_dt"] = 0.1 + elif changed == "language": + source.metadata["language"]["task"][0] = "release" + else: + source.data["obs"][0, 0] += 1 + generator = torch.Generator().manual_seed(9) + before = generator.get_state().clone() + with pytest.raises(ValueError, match="changed|unbudgeted"): + bridge.sample(1, generator=generator) + torch.testing.assert_close(generator.get_state(), before) + with pytest.raises(ValueError, match="changed|unbudgeted"): + bridge.validate_resume() + + +def test_resume_rehash_detects_untracked_tensor_writes(): + source = _resident() + bridge = SequenceBatchSource( + source, + SequenceSpecCfg(2, 2), + encoder=_fit_encode, + encoder_version="v1", + immutable=True, + ) + source.data["actions"].data[0, 0] += 1 + with pytest.raises(ValueError, match="content changed"): + bridge.validate_resume() + + +def test_immutable_content_hash_handles_noncontiguous_large_cpu_features(): + from tensordict import TensorDict + + values = ( + torch.arange(2 * 3 * 100003, dtype=torch.float32) + .reshape(2, 100003, 3) + .transpose(1, 2) + ) + assert not values.is_contiguous() + data = TensorDict( + { + "obs": values, + "actions": torch.ones(2, 3, 1), + "valid": torch.ones(2, 3, dtype=torch.bool), + }, + [2, 3], + ) + schema = {"action": {"semantics": "toy", "step_dt": 0.1}} + + def wrap(payload): + return SequenceBatchSource( + TensorDataSource( + payload, schema=schema, source_id="large", copy_data=False + ), + SequenceSpecCfg(1, 1), + encoder=_fit_encode, + encoder_version="v1", + immutable=True, + ) + + first = wrap(data) + second = wrap(data.contiguous()) + assert first.source_id == second.source_id + first.validate_resume() diff --git a/tests/learning/sl/test_trainer.py b/tests/learning/sl/test_trainer.py index 97bab43e33..6c0fcbac9d 100644 --- a/tests/learning/sl/test_trainer.py +++ b/tests/learning/sl/test_trainer.py @@ -41,13 +41,38 @@ def compute_loss(self, model, batch): return LossResult((model(batch.inputs) - batch.targets).square().mean()) -def build(path): +def build(path, *, sequence=False): torch.manual_seed(42) model = torch.nn.Sequential( torch.nn.Linear(2, 8), torch.nn.Dropout(0.2), torch.nn.Linear(8, 1) ) x = torch.arange(40, dtype=torch.float32).view(20, 2) / 40 data = TensorBatchSource(SupervisedBatch(x, x.sum(-1, keepdim=True))) + if sequence: + from tensordict import TensorDict + from embodichain.data_pipeline.datasets import TensorDataSource, SequenceSpecCfg + from embodichain.learning.sl import SequenceBatchSource + + data = SequenceBatchSource( + TensorDataSource( + TensorDict( + { + "obs": x.reshape(4, 5, 2), + "actions": x.sum(-1, keepdim=True).reshape(4, 5, 1), + "valid": torch.ones(4, 5, dtype=torch.bool), + }, + [4, 5], + ), + schema={"action": {"semantics": "toy-sum-target", "step_dt": 0.1}}, + source_id="toy-resident", + ), + SequenceSpecCfg(2, 2), + encoder=lambda context: SupervisedBatch( + context["obs"][:, 1], context["actions"][:, 1] + ), + encoder_version="current-to-sum-v1", + immutable=True, + ) learner = SupervisedLearner( model, Regression(), torch.optim.Adam(model.parameters(), lr=0.01) ) @@ -59,19 +84,23 @@ def build(path): ) -def test_resume_reproduces_dropout_optimizer_and_sampling(tmp_path): - full = build(tmp_path / "full") +@pytest.mark.parametrize("sequence", [False, True]) +def test_resume_reproduces_dropout_optimizer_and_sampling(tmp_path, sequence): + full = build(tmp_path / "full", sequence=sequence) full.fit() expected = copy.deepcopy(full.learner.state_dict()) - split = build(tmp_path / "split") + split = build(tmp_path / "split", sequence=sequence) split.fit(stop_after=7) - resumed = build(tmp_path / "resumed") + resumed = build(tmp_path / "resumed", sequence=sequence) resumed.resume(tmp_path / "split/last.pt") resumed.fit() for key, value in expected["model"].items(): torch.testing.assert_close( value, resumed.learner.model.state_dict()[key], rtol=0, atol=0 ) + torch.testing.assert_close( + resumed.learner.state_dict()["optimizer"], expected["optimizer"], rtol=0, atol=0 + ) combined = (tmp_path / "split/metrics.jsonl").read_text() + ( tmp_path / "resumed/metrics.jsonl" ).read_text() @@ -113,7 +142,7 @@ def test_source_snapshots_and_sampled_batches_own_storage(): def test_sl_import_does_not_load_simulation(): - code = "import sys; import embodichain.learning.sl; assert not any(k.startswith(('dexsim', 'embodichain.lab', 'embodichain.learning.rl')) for k in sys.modules)" + code = "import sys; import embodichain.learning.sl; assert not any(k.startswith(('dexsim', 'embodichain.lab', 'embodichain.learning.rl', 'embodichain.data_pipeline', 'lerobot')) for k in sys.modules); from embodichain.learning.sl import SequenceBatchSource; assert not any(k.startswith(('dexsim', 'embodichain.lab', 'embodichain.data_pipeline.engine', 'lerobot')) for k in sys.modules)" subprocess.run([sys.executable, "-c", code], check=True)