PyTorch

Optional — requires pip install "pybvh-ml[torch]". See the PyTorch Integration guide for the training-loop context, the set_epoch contract, and the length semantics.

Datasets

datasets

PyTorch Dataset classes for motion capture data.

EpochState()

Shared-memory epoch counter for DataLoader-worker visibility.

Pairs with :func:rng_for: it supplies the epoch term that makes a seeded sample's draw change from one epoch to the next. Public because a Dataset that isn't a :class:MotionDataset subclass needs exactly this to honor the set_epoch contract — hold one, call :meth:set from set_epoch, and read :attr:current in __getitem__.

The epoch lives in a multiprocessing.Value so that set_epoch() in the main process is observed by DataLoader workers — including persistent ones (persistent_workers=True), which are created once and never re-receive the dataset. Workers inherit the shared handle when the DataLoader passes the dataset as Process args, which works under both fork and spawn start methods.

The Value is built from an explicit spawn context rather than the process default, and that choice is load-bearing: a fork-context lock is an anonymous semaphore, unlinked at creation, whose handle is meaningless in a spawn-started child — it unpickles without complaint and segfaults the worker on first use. A spawn-context lock is named, so it survives both inheritance (fork) and reopen-by-name (spawn). Linux defaults to fork, so the mismatch is reachable with a plain DataLoader(..., multiprocessing_context="spawn").

-1 is the never-set sentinel (replaces a separate boolean — one would live in whichever process wrote it and read False in every forked worker; :attr:is_set reads the shared sentinel instead). Deliberately no __getstate__/__setstate__: swapping the Value for a plain int during pickling would silently break sharing under spawn (worker creation uses the same pickle machinery). The cost is that holders cannot be copy.deepcopy-ed or torch.save-ed directly — shared ctypes only travel via process inheritance.

Note that a spawn DataLoader additionally pickles the whole dataset, so everything it holds must be picklable — in particular a :class:~pybvh_ml.AugmentationPipeline whose kwargs are lambda callables cannot cross a spawn boundary; use module-level functions there.

current: int property

Current epoch (0 when never set), with no warning.

For read-only callers: the warn-once budget in :meth:effective belongs to the training path, and a diagnostic read must not spend it and mask the real warning later.

is_set: bool property

Whether :meth:set has been called yet.

:attr:current answers "which epoch do I augment as", and for that question "never set" and "epoch 0" are the same answer, deliberately. They are not the same fact, though, and code that needs the fact — a trainer hook claiming epoch 0 only if nothing has claimed it, a test asserting the hook ran — was left reaching for the private _raw().

Reads the shared value, so it is worker-visible like everything else here, and costs no warn-once budget. It is not atomic with :meth:set: if not state.is_set: state.set(0) is a check-then-act, so run it in the main process before the DataLoader starts its workers — which is where an epoch-0 claim has to happen anyway, since workers that fork beforehand carry the unset state into their first batches.

MotionDataset(clips: list[dict], labels: np.ndarray | None = None, target_length: int | None = None, augmentation: AugmentationPipeline | None = None, seed: int | None = None, *, center_root: bool = False, temporal: str = 'pad', layout: str = 'flat', streams: tuple[str, ...] = DEFAULT_STREAMS, position_centering: str | None = None, source_repr: str | None = None, target_repr: str | None = None, euler_orders: list[str] | None = None, names: Sequence[str] | None = None)

Bases: Dataset

Dataset that loads preprocessed motion clips.

Designed to work with the output of :func:pybvh_ml.preprocessing.load_preprocessed.

Parameters:

Name	Type	Description	Default
clips	list of dict	Each dict must have root_pos (F, 3) and joint_rot (F, J, C), plus joint_pos (F, J, 3) or node_pos (F, N, 3) when the dataset was preprocessed with include_positions=True. The pre-0.5.0 key joint_data is still read, so clip dicts carried over from that era work unchanged.	required
labels	array - like or None	Per-clip integer labels. Must cover every clip when given.	None
names	sequence of str or None	Per-clip identity, surfaced as item["name"] and collated into batch["names"]. Anything doing per-clip work rather than reporting a dataset-level mean needs it: writing per-prediction rows, routing a clip to the fold whose model never saw its performer, keying saved activations. :meth:from_preprocessed fills it from the stored filenames, which is where it should come from. The convention is the filename stem — what :func:~pybvh_ml.preprocessing.preprocess_directory records and what :class:OnTheFlyDataset reports — not the full path, so the same clip carries the same name whether it is read from a preprocessed file or straight from BVH. The cost is that two identically named files from different directories collide; pass your own disambiguated strings if a corpus does that. None (default) omits the key entirely rather than substituting indices, so a downstream can tell "no identity was provided" from a real name.	None
target_length	int or None	If given, standardize all clips to this length using temporal. The length reported by __getitem__ is the number of valid frames actually present in the returned tensor, so padded frames are excluded.	None
temporal	('pad', 'crop', 'resample', 'resample_deterministic')	How target_length is reached. These are genuinely different operations, not styles: "pad" (default) and "crop" keep a fixed window of the clip — truncating from the end and from the center respectively, zero-padding when the clip is shorter — while the two resample modes keep the whole arc of the clip at a fixed frame budget, sampling target_length frame indices spread across its full duration (:func:~pybvh_ml.sequences.uniform_temporal_sample). Resample when the shape of the whole clip carries the signal and clips vary in duration; crop or pad when a fixed-duration window does. "resample" draws a random offset within each segment (a temporal augmentation, and it consumes the sample's rng); "resample_deterministic" takes each segment's first frame, so a clip yields the same frames on every read — the evaluation counterpart. Both report length == target_length: every returned frame is real data.	"pad"
layout	('flat', 'ctv', 'tvc')	Tensor layout of the returned data. "flat" (default) gives (T, D) via :func:~pybvh_ml.packing.pack_to_flat; "ctv" and "tvc" give the graph layouts (C, T, V) / (T, V, C) that GCN and skeleton-transformer models consume. Only "flat" works with :func:~pybvh_ml.torch.collate_motion_batch, which pads a time-major axis 0; the graph layouts are fixed-size by construction (they pair with target_length) so they stack with :func:torch.utils.data.default_collate.	"flat"
streams	tuple of str	Which streams the packed tensor carries, and in what order — forwarded to the packer, so the vocabulary and shape rules are :func:~pybvh_ml.pack_to_ctv's. Default ("root_pos", "joint_rot"); ("joint_pos",) with layout="ctv" is the ST-GCN input, (3, T, J). Still one data tensor with explicit streams rather than a second tensor, so the batch contract does not depend on preprocessing flags. The derived streams ("joint_vel", "joint_acc" and the node-space pair) are computed before temporal is applied, so a difference is always taken across consecutive frames of the augmented clip and then subsampled along with its base — never differenced across padded or resampled frames. That ordering is available only here; the same three lines written in a collate_fn run on this method's output and cannot reproduce it.	DEFAULT_STREAMS
position_centering	('world', 'skeleton', 'first')	Frame the clips' position arrays are in, recorded on every :class:~pybvh_ml.MotionArrays this dataset mints. Storage metadata alone is not enough: the steps that depend on it (root-position noise, the FK refresh, center_root=True packing) only ever see the container. :meth:from_preprocessed fills it from the loaded dataset, which is where it should come from. Required — not None — when streams includes a derived stream: the frame decides whether joint_vel is world velocity or velocity relative to the root, so an undeclared one is refused at construction rather than guessed on the first sample.	"world"
source_repr	str or None	Convert each clip's joint_data from source_repr to target_repr before augmentation, so a dataset stored in one representation can train in another without a second preprocessing pass. Both are required together; target_repr=None (default) returns the stored representation untouched. :meth:from_preprocessed fills source_repr from the dataset metadata, which is where it should come from — restating it at the call site is how it ends up wrong.	None
target_repr	str or None	Convert each clip's joint_data from source_repr to target_repr before augmentation, so a dataset stored in one representation can train in another without a second preprocessing pass. Both are required together; target_repr=None (default) returns the stored representation untouched. :meth:from_preprocessed fills source_repr from the dataset metadata, which is where it should come from — restating it at the call site is how it ends up wrong.	None
euler_orders	list of str or None	Per-joint Euler orders, required when either end of the conversion is "euler". From skeleton_info["euler_orders"].	None
augmentation	AugmentationPipeline or None	Applied on-the-fly during __getitem__, after any target_repr conversion — so the pipeline's declared representation is target_repr when one is set.	None
center_root	bool	If True, subtract each clip's first-frame root position in __getitem__. Default False — clips from :func:~pybvh_ml.preprocessing.load_preprocessed are already centered when the dataset was saved with center_root=True (check the loaded center_root metadata); set True for hand-built raw clip dicts, mirroring :class:OnTheFlyDataset.	False
seed	int or None	Base seed for reproducible augmentation. When set, combined with the current epoch (see :meth:set_epoch) and the sample index into a SeedSequence so each (seed, epoch, idx) triple produces a distinct but reproducible stream. Set None for fresh OS entropy each call.	None

Notes

Per-epoch augmentation variety: call dataset.set_epoch(epoch) at the start of each training epoch so the seeded augmentation changes across epochs — same contract as :class:torch.utils.data.distributed.DistributedSampler. The epoch lives in shared memory, so this works with num_workers > 0 including persistent_workers=True. When seed is set and set_epoch is never called, every epoch sees the same augmentation per sample index (useful for debugging, harmful for training dynamics).

Pickling: because of the shared-memory epoch, instances cannot be copy.deepcopy-ed or torch.save-ed directly — shared state only travels via process inheritance (which is exactly how the DataLoader hands the dataset to its workers).

epoch: int property

Epoch this dataset augments as — 0 when none was set.

epoch_is_set: bool property

Whether :meth:set_epoch has been called.

Distinct from epoch == 0, which is also what an unset dataset reports. The question matters under a training framework that builds its DataLoader — and forks its workers — before the hook you put set_epoch in: if not ds.epoch_is_set: ds.set_epoch(0) in an earlier hook claims the epoch before any worker inherits the unset state. Reads the shared counter, so the answer is the same in every worker.

from_preprocessed(loaded: dict, **kwargs) -> 'MotionDataset' classmethod

Build a dataset from a :func:~pybvh_ml.preprocessing.load_preprocessed result.

Wires the metadata that would otherwise be restated by hand at the call site: the clips, labels and filenames (as names), the stored representation (as source_repr, which target_repr conversion needs), skeleton_info["euler_orders"], and the stored position_centering. center_root defaults to False because the stored arrays already reflect the choice made at preprocessing time — centering again here would be a second, unrecorded transform.

Parameters:

Name	Type	Description	Default
loaded	dict	The dict returned by :func:~pybvh_ml.preprocessing.load_preprocessed.	required
**kwargs		Forwarded to :class:MotionDataset; anything passed here overrides what the metadata supplies.	{}

Raises:

Type	Description
ValueError	If streams= names a stream the clips do not carry — the message says to preprocess with include_positions=True.

Examples:

>>> loaded = load_preprocessed("train.npz")          # stored as euler
>>> ds = MotionDataset.from_preprocessed(
...     loaded, target_repr="6d", layout="ctv",
...     temporal="resample", target_length=64, seed=0)

>>> loaded = load_preprocessed("keypoints.npz")   # include_positions
>>> ds = MotionDataset.from_preprocessed(
...     loaded, layout="ctv", streams=("joint_pos",),
...     temporal="resample", target_length=64, seed=0)
>>> ds[0]["data"].shape                # (3, 64, J) — into ST-GCN

set_epoch(epoch: int) -> None

Set the current epoch for per-epoch reproducible augmentation.

Mirrors :meth:torch.utils.data.distributed.DistributedSampler.set_epoch; reaches DataLoader workers (persistent ones included) via shared memory.

explain_augmentation(idx: int, *, epoch: int | None = None) -> list[dict]

Report what the augmentation did to sample idx.

Re-runs this sample's augmentation on the same (seed, epoch, idx) rng the loader used, so the records describe the draw that actually ran rather than a fresh one. Their layout is the pipeline's return_params format: {"name", "applied", "params"} per step.

Parameters:

Name	Type	Description	Default
idx	int	Sample index; negative indexing works as in __getitem__.	required
epoch	int	Epoch to replay. Defaults to the dataset's current epoch — pass it explicitly to ask about an earlier one.	None

Returns:

Type	Description
list of dict	One record per configured augmentation step, or [] when the dataset has no augmentation.

Raises:

Type	Description
ValueError	If the dataset was built without a seed. Unseeded draws come from fresh OS entropy and cannot be reconstructed; answering with a new draw would describe an augmentation that never ran.

Notes

The replay is truthful only while its inputs are unchanged: the same pipeline (same steps, probabilities and ranges) over the same clip arrays. Rebuild the pipeline differently and the records describe a run that no longer exists.

OnTheFlyDataset(bvh_paths: list[str | Path], representation: str | None = '6d', target_length: int | None = None, augmentation: AugmentationPipeline | None = None, center_root: bool = True, label_fn: Callable[[str], int] | None = None, seed: int | None = None, *, world_up: str = 'auto', lr_mapping: dict[str, str] | None = None, temporal: str = 'pad', layout: str = 'flat', streams: tuple[str, ...] = DEFAULT_STREAMS, include_positions: bool = False, position_space: str = 'joint', position_centering: str = 'world')

Bases: Dataset

Dataset that loads BVH files on-the-fly for maximum augmentation variety.

Slower than :class:MotionDataset but avoids pre-extracting arrays, so every epoch sees freshly augmented data.

Parameters:

Name	Type	Description	Default
bvh_paths	list of str or Path	Paths to BVH files (coerced to :class:~pathlib.Path).	required
representation	str or None	Rotation representation for joint data. Extraction happens per clip, so this is already the "target" representation — :class:MotionDataset's source_repr / target_repr pair has no counterpart here. None extracts no rotations, which requires include_positions=True.	'6d'
include_positions	bool	Also extract positions per clip. This class calls :func:~pybvh_ml.preprocessing.extract_repr per item rather than reading preprocessed clips, so include_positions / position_space / position_centering live here rather than in a preprocessing step. One FK pass per clip regardless of how many representations are requested — pybvh caches world-frame FK on the Bvh and invalidates it on motion writes.	False
position_space	('joint', 'node')	Which index space to extract — see :func:~pybvh_ml.preprocess_directory.	"joint"
position_centering	('world', 'skeleton', 'first')	Which frame to extract them in, recorded on every :class:~pybvh_ml.MotionArrays this dataset mints.	"world"
streams	tuple of str	Which streams the packed tensor carries — see :class:MotionDataset.	DEFAULT_STREAMS
target_length	int or None	If given, standardize to this length using temporal. The reported length counts only the valid frames present in the returned tensor (see :class:MotionDataset).	None
temporal	('pad', 'crop', 'resample', 'resample_deterministic')	How target_length is reached — see :class:MotionDataset.	"pad"
layout	('flat', 'ctv', 'tvc')	Tensor layout of the returned data — see :class:MotionDataset.	"flat"
augmentation	AugmentationPipeline or None		None
center_root	bool	If True (default), subtract each clip's first-frame root position after extraction.	True
label_fn	callable or None	label_fn(filename_stem) -> int.	None

Notes

Every item carries name — the source file's stem — with no parameter to enable it, because unlike :class:MotionDataset this class has the paths and identity is never a guess. The stem rather than the full path is what makes it the same identity a preprocessed dataset reports (filenames), so a clip keeps its name across both paths; it is also what label_fn receives. world_up : str Forwarded to :func:pybvh.read_bvh_file per clip. "auto" (default) auto-detects; pass "+y" etc. to override — same semantics as :func:~pybvh_ml.preprocessing.preprocess_directory. lr_mapping : dict or None Forwarded to :func:pybvh.read_bvh_file. Explicit left/right joint pair mapping for uniform dataset conventions. seed : int or None See :class:MotionDataset for seeding semantics. Call :meth:set_epoch at the start of each epoch for reproducible per-epoch variety — reaches DataLoader workers (persistent ones included) via shared memory; see :class:MotionDataset for the pickling caveat.

epoch: int property

Epoch this dataset augments as — 0 when none was set.

epoch_is_set: bool property

Whether :meth:set_epoch has been called.

Distinct from epoch == 0, which is also what an unset dataset reports. The question matters under a training framework that builds its DataLoader — and forks its workers — before the hook you put set_epoch in: if not ds.epoch_is_set: ds.set_epoch(0) in an earlier hook claims the epoch before any worker inherits the unset state. Reads the shared counter, so the answer is the same in every worker.

set_epoch(epoch: int) -> None

Set the current epoch for reproducible per-epoch augmentation.

explain_augmentation(idx: int, *, epoch: int | None = None) -> list[dict]

Report what the augmentation did to sample idx.

Re-runs this sample's augmentation on the same (seed, epoch, idx) rng the loader used, so the records describe the draw that actually ran rather than a fresh one. Their layout is the pipeline's return_params format: {"name", "applied", "params"} per step. The source file is re-read, so this costs a parse per call.

Parameters:

Name	Type	Description	Default
idx	int	Sample index; negative indexing works as in __getitem__.	required
epoch	int	Epoch to replay. Defaults to the dataset's current epoch — pass it explicitly to ask about an earlier one.	None

Returns:

Type	Description
list of dict	One record per configured augmentation step, or [] when the dataset has no augmentation.

Raises:

Type	Description
ValueError	If the dataset was built without a seed. Unseeded draws come from fresh OS entropy and cannot be reconstructed; answering with a new draw would describe an augmentation that never ran.

Notes

The replay is truthful only while its inputs are unchanged: the same pipeline (same steps, probabilities and ranges) over the same source file. Edit the BVH on disk and the records describe a run that no longer exists.

rng_for(seed: int | None, epoch: int, idx: int) -> np.random.Generator

Build the per-sample generator for one (seed, epoch, idx) triple.

The seeding scheme both Dataset classes use, exposed because any Dataset needs it — not only subclasses of the two shipped here. A SeedSequence([seed, epoch, idx]) makes each sample's stream independent of the order samples are drawn in and of which worker draws them, so a clip augments identically whether it lands in worker 0 or worker 3, and shuffling doesn't change what any sample receives.

Parameters:

Name	Type	Description	Default
seed	int or None	Base seed. None returns a generator seeded from fresh OS entropy — reproducibility is off, and the (epoch, idx) pair is ignored.	required
epoch	int	Current epoch; see :class:EpochState for propagating it to DataLoader workers.	required
idx	int	Sample index. Must be non-negative — resolve Python negative indexing before calling.	required

Returns:

Type	Description
Generator

Examples:

>>> from pybvh_ml.torch import EpochState, rng_for
>>> class MyFeeder(torch.utils.data.Dataset):
...     def __init__(self, seed=0):
...         self.seed = seed
...         self.epoch_state = EpochState()
...     def set_epoch(self, epoch):
...         self.epoch_state.set(epoch)
...     def __getitem__(self, idx):
...         rng = rng_for(self.seed, self.epoch_state.current, idx)
...         ...

Collate

collate

Collate function for variable-length motion sequences.

collate_motion_batch(batch: list[dict]) -> dict[str, torch.Tensor]

Collate variable-length motion clips into a padded batch.

Parameters:

Name	Type	Description	Default
batch	list of dict	Each dict must have data — a 2-D (T, D) tensor, the flat layout — and length (int), the number of valid frames in data, i.e. length <= data.shape[0] with any frames beyond it being padding (the contract :class:~pybvh_ml.torch.MotionDataset and :class:~pybvh_ml.torch.OnTheFlyDataset provide under their default layout="flat"). Optionally label (int) and name (str).	required

Returns:

Type

Description

dict

data : (B, T_max, D) float tensor, zero-padded. lengths : (B,) long tensor of valid frame counts. mask : (B, T_max) bool tensor (True = valid frame). labels : (B,) long tensor (if labels present). names : list of B strings, in batch order (if names present). A list rather than a tensor — strings have no tensor form — which is what :func:torch.utils.data.default_collate also does with them, so a batch means the same thing under either collate.

Raises:

Type	Description
ValueError	If any data is not 2-D. Padding is time-major on axis 0, which the graph layouts don't satisfy: (C, T, V) puts channels there, so padding and masking would silently run along the wrong axis. Those layouts are fixed-size by construction — stack them with :func:torch.utils.data.default_collate instead.