dragon-torch 1.1.0__py3-none-any.whl

dragon_torch/__init__.py

@@ -0,0 +1,7 @@
+"""dragon-torch: PyTorch deep learning research framework, focused on RL."""
+from importlib.metadata import PackageNotFoundError, version
+
+try:
+    __version__ = version('dragon-torch')
+except PackageNotFoundError:  # running from a source checkout that isn't installed
+    __version__ = '0.0.0'

dragon_torch/common/defaults.py

@@ -0,0 +1,6 @@
+import os
+numpy_default_dtype: str = os.getenv('DEFAULT_NUMPY_DTYPE')
+
+
+def identity(x):
+    return x

dragon_torch/common/torch/tree.py

@@ -0,0 +1,35 @@
+"""Tiny pytree helpers for observations that may be arrays, dicts or tuples."""
+from typing import Any, Callable, Sequence
+import torch
+
+type Tree = Any
+
+
+def tree_map(fn: Callable[[Any], Any], tree: Tree) -> Tree:
+    if isinstance(tree, dict):
+        return {k: tree_map(fn, v) for k, v in tree.items()}
+    if isinstance(tree, tuple):
+        return tuple(tree_map(fn, v) for v in tree)
+    return fn(tree)
+
+
+def tree_stack(trees: Sequence[Tree], dim: int = 0) -> Tree:
+    """Stack a list of same-structured trees of tensors leaf-wise."""
+    first = trees[0]
+    if isinstance(first, dict):
+        return {k: tree_stack([t[k] for t in trees], dim) for k in first}
+    if isinstance(first, tuple):
+        return tuple(tree_stack([t[i] for t in trees], dim) for i in range(len(first)))
+    return torch.stack(list(trees), dim)
+
+
+def tree_set(tree: Tree, index: int, value: Tree) -> None:
+    """In place: tree[index] = value, leaf-wise (for batched numpy / torch leaves)."""
+    if isinstance(tree, dict):
+        for k in tree:
+            tree_set(tree[k], index, value[k])
+    elif isinstance(tree, tuple):
+        for leaf, v in zip(tree, value):
+            tree_set(leaf, index, v)
+    else:
+        tree[index] = value

dragon_torch/common/torch/types.py

@@ -0,0 +1,7 @@
+from typing import Union, Literal
+
+type DeviceLike = Union[Literal['cpu'], Literal['cuda']]
+type TorchFloatKind = Union[Literal['float'], Literal['float16'],
+                            Literal['float32'], Literal['float64']]
+type TorchIntegerKind = Union[Literal['int'], Literal['int8'], Literal['int16'],
+                              Literal['int64'], Literal['int32']]

dragon_torch/common/types.py

@@ -0,0 +1,22 @@
+from dragon_torch.common.torch.types import DeviceLike, TorchFloatKind, TorchIntegerKind
+import os
+from pathlib import Path
+from dotenv import load_dotenv
+from typing import Union, Literal, List
+from .defaults import numpy_default_dtype
+load_dotenv('./.env')
+
+
+# ------- Global ------- #
+type PathLike = Union[Path, str]
+
+# ------- DatasetDefinition ------- #
+type BackendTypes = Union[Literal['jax'],
+                          Literal['numpy'], Literal['pandas'], Literal['torch']]
+type DatasetSplitType = Union[Literal['train'],
+                              Literal['test'], Literal['valid']]
+
+
+# ------- Torch Types ------- #
+type TorchNumeric = Union[TorchIntegerKind, TorchFloatKind]
+type FileExtensionLike = List[str | None]

dragon_torch/common/utils.py

@@ -0,0 +1,6 @@
+from pathlib import Path
+from .types import PathLike
+
+
+def build_path(x: PathLike) -> Path:
+    return x if isinstance(x, Path) else Path(x)

dragon_torch/data/RL/__init__.py

@@ -0,0 +1,3 @@
+from .collector import ActOutput, Actor, EnvCollector, Transition, ValueActor
+from .episodes import EpisodeSource
+from .rollout import RolloutSource, compute_gae

dragon_torch/data/RL/collector.py

@@ -0,0 +1,169 @@
+"""Env interaction shared by every RL data source.
+
+EnvCollector steps a vectorised env with any `Actor` and returns `Transition`s
+that already carry the *true* successor obs (info['final_obs'] substituted on
+episode end). On-policy sources (RolloutSource) and off-policy sources (a
+replay source storing (obs, action, reward, next_obs, terminated)) both build
+on it, so autoreset, normalisation, episode stats and step counting live here once.
+"""
+import time
+from dataclasses import dataclass, field
+from typing import Any, Optional, Protocol, runtime_checkable
+import numpy as np
+import torch
+from gymnasium import spaces
+from dragon_torch.common.torch.tree import tree_map, tree_set
+from dragon_torch.data.envs.config import DragonEnvConfig, make_vec_env
+from dragon_torch.data.envs.normalize import ObsNormalizer, RewardScaler
+from dragon_torch.data.envs.spaces import action_to_env
+
+
+@dataclass
+class ActOutput:
+    action: torch.Tensor  # batched, in the policy's own space (what buffers store)
+    extras: dict[str, torch.Tensor] = field(default_factory=dict)  # e.g. log_prob, value
+
+
+@runtime_checkable
+class Actor(Protocol):
+    """What sources need from a model. Off-policy models implement exploration
+    (e.g. epsilon-greedy) inside `act`."""
+
+    def act(self, obs: Any, deterministic: bool = False) -> ActOutput: ...
+
+
+@runtime_checkable
+class ValueActor(Actor, Protocol):
+    """Actor with a state-value estimate, needed for GAE bootstrapping."""
+
+    def value(self, obs: Any) -> torch.Tensor: ...
+
+
+@dataclass
+class Transition:
+    """One vector step. obs / next_obs are (normalised) numpy batches."""
+    obs: Any
+    action: torch.Tensor
+    reward: np.ndarray       # float32, scaled if reward normalisation is on
+    next_obs: Any            # true successor: final obs where an episode ended
+    terminated: np.ndarray
+    truncated: np.ndarray
+    extras: dict[str, torch.Tensor]
+
+    @property
+    def done(self) -> np.ndarray:
+        return self.terminated | self.truncated
+
+
+class EnvCollector:
+    def __init__(self, env_cfg: DragonEnvConfig, normalize_obs: bool = False,
+                 normalize_reward: bool = False, gamma: float = 0.99,
+                 update_stats: bool = True):
+        self.env_cfg = env_cfg
+        self.env = make_vec_env(env_cfg)
+        self.n_envs = self.env.num_envs
+        self.observation_space = self.env.single_observation_space
+        self.action_space = self.env.single_action_space
+        if normalize_obs and not isinstance(self.observation_space, spaces.Box):
+            raise TypeError(f"normalize_obs needs a Box observation space, got "
+                            f"{self.observation_space}; normalise inside your encoder instead")
+        self.obs_normalizer: Optional[ObsNormalizer] = \
+            ObsNormalizer(self.observation_space.shape) if normalize_obs else None
+        self.reward_scaler: Optional[RewardScaler] = \
+            RewardScaler(self.n_envs, gamma) if normalize_reward else None
+        # False for evaluation: use (shared) stats without changing them
+        self.update_stats = update_stats
+        self.global_step = 0
+        self._obs: Any = None
+        self._ep_return = np.zeros(self.n_envs)
+        self._ep_length = np.zeros(self.n_envs, np.int64)
+        self._finished: list[tuple[float, int]] = []
+        self._clock = (time.perf_counter(), 0)
+
+    # ------- obs ------- #
+    def _normalize(self, raw: Any) -> Any:
+        return raw if self.obs_normalizer is None else self.obs_normalizer.normalize(raw)
+
+    def _observe(self, raw: Any) -> Any:
+        if self.obs_normalizer is not None and self.update_stats:
+            self.obs_normalizer.update(raw)
+        return self._normalize(raw)
+
+    def reset(self, seed: Optional[int] = None) -> Any:
+        raw, _ = self.env.reset(seed=self.env_cfg.seed if seed is None else seed)
+        self._obs = self._observe(raw)
+        self._ep_return[:] = 0
+        self._ep_length[:] = 0
+        return self._obs
+
+    @property
+    def obs(self) -> Any:
+        """Current (normalised) obs batch; resets the env on first use."""
+        if self._obs is None:
+            self.reset()
+        return self._obs
+
+    # ------- stepping ------- #
+    def step(self, actor: Actor, deterministic: bool = False) -> Transition:
+        obs = self.obs
+        out = actor.act(obs, deterministic=deterministic)
+        raw_next, reward, terminated, truncated, info = self.env.step(
+            action_to_env(out.action, self.action_space))
+        done = terminated | truncated
+
+        next_policy_obs = self._observe(raw_next)  # reset obs where done
+        next_obs = next_policy_obs
+        if done.any():
+            true_next = tree_map(np.copy, raw_next)
+            for i in np.flatnonzero(done):
+                tree_set(true_next, i, info['final_obs'][i])
+            next_obs = self._normalize(true_next)
+
+        self._ep_return += reward
+        self._ep_length += 1
+        for i in np.flatnonzero(done):
+            self._finished.append((float(self._ep_return[i]), int(self._ep_length[i])))
+            self._ep_return[i], self._ep_length[i] = 0, 0
+
+        train_reward = reward if self.reward_scaler is None else \
+            self.reward_scaler(reward, done, update=self.update_stats)
+        self.global_step += self.n_envs
+        self._obs = next_policy_obs
+        return Transition(obs=obs, action=out.action, reward=np.asarray(train_reward, np.float32),
+                          next_obs=next_obs, terminated=terminated, truncated=truncated,
+                          extras=out.extras)
+
+    # ------- stats / state ------- #
+    def drain_episodes(self) -> list[tuple[float, int]]:
+        """(raw return, length) of episodes finished since the last drain."""
+        finished, self._finished = self._finished, []
+        return finished
+
+    def metrics(self) -> dict[str, float]:
+        now, steps = time.perf_counter(), self.global_step
+        then, steps_then = self._clock
+        self._clock = (now, steps)
+        record = {'global_step': steps, 'steps_per_sec': (steps - steps_then) / max(now - then, 1e-9)}
+        episodes = self.drain_episodes()
+        if episodes:
+            record['episode_return'] = float(np.mean([r for r, _ in episodes]))
+            record['episode_length'] = float(np.mean([l for _, l in episodes]))
+            record['episodes'] = len(episodes)
+        return record
+
+    def state_dict(self) -> dict[str, Any]:
+        # the env itself can't be checkpointed; a resumed run starts fresh episodes
+        return {'global_step': self.global_step,
+                'obs_normalizer': None if self.obs_normalizer is None else self.obs_normalizer.state_dict(),
+                'reward_scaler': None if self.reward_scaler is None else self.reward_scaler.state_dict()}
+
+    def load_state_dict(self, state: dict[str, Any]) -> None:
+        self.global_step = state['global_step']
+        self._clock = (time.perf_counter(), self.global_step)
+        for name in ('obs_normalizer', 'reward_scaler'):
+            target = getattr(self, name)
+            if target is not None and state.get(name) is not None:
+                target.load_state_dict(state[name])
+
+    def close(self) -> None:
+        self.env.close()

dragon_torch/data/RL/episodes.py

@@ -0,0 +1,47 @@
+from typing import Any, Iterator, Optional
+from dragon_torch.data.envs.config import DragonEnvConfig
+from dragon_torch.data.RL.collector import Actor, EnvCollector
+
+
+class EpisodeSource:
+    """Evaluation source: iterating runs the bound model until n_episodes finish,
+    yielding {'return', 'length'} per episode (raw, unscaled rewards).
+
+    Every iteration re-seeds the env with env_cfg.seed, so successive evaluations
+    of the same policy are comparable. On bind it borrows the training source's
+    observation normaliser (read-only) so the policy sees the inputs it trained on.
+    """
+
+    def __init__(self, env_cfg: DragonEnvConfig, n_episodes: int = 10,
+                 deterministic: bool = True, max_steps: int = 1_000_000):
+        self.collector = EnvCollector(env_cfg, update_stats=False)
+        self.n_episodes, self.deterministic, self.max_steps = n_episodes, deterministic, max_steps
+        self.actor: Optional[Actor] = None
+
+    def bind(self, model: Actor) -> None:
+        if not isinstance(model, Actor):
+            raise TypeError(f"EpisodeSource needs a model with act(), got {type(model)}")
+        self.actor = model
+        train_collector = getattr(getattr(model, 'train_source', None), 'collector', None)
+        self.collector.obs_normalizer = getattr(train_collector, 'obs_normalizer', None)
+
+    def __iter__(self) -> Iterator[dict[str, Any]]:
+        if self.actor is None:
+            raise RuntimeError("EpisodeSource is unbound; evaluate() binds the model")
+        self.collector.reset()
+        self.collector.drain_episodes()
+        yielded, steps = 0, 0
+        while yielded < self.n_episodes:
+            if steps >= self.max_steps:
+                raise RuntimeError(f"{self.n_episodes} episodes did not finish within "
+                                   f"{self.max_steps} vector steps; add a TimeLimit wrapper")
+            self.collector.step(self.actor, deterministic=self.deterministic)
+            steps += 1
+            for ret, length in self.collector.drain_episodes():
+                yield {'return': ret, 'length': length}
+                yielded += 1
+                if yielded == self.n_episodes:
+                    return
+
+    def close(self) -> None:
+        self.collector.close()

dragon_torch/data/RL/rollout.py

@@ -0,0 +1,122 @@
+from typing import Any, Iterator, Optional
+import torch
+from dragon_torch.common.torch.tree import tree_map, tree_stack
+from dragon_torch.data.envs.config import DragonEnvConfig
+from dragon_torch.data.envs.spaces import obs_to_tensor
+from dragon_torch.data.RL.collector import EnvCollector, ValueActor
+
+
+def compute_gae(rewards: torch.Tensor, values: torch.Tensor, next_values: torch.Tensor,
+                terminated: torch.Tensor, dones: torch.Tensor, gamma: float,
+                gae_lambda: float) -> tuple[torch.Tensor, torch.Tensor]:
+    """Generalised advantage estimation over [T, N] tensors.
+
+    next_values[t] is V of the *true* successor of step t (the final obs when an
+    episode ended), so truncated episodes bootstrap while terminated ones don't,
+    and the advantage recursion is cut at every episode boundary.
+    Returns (advantages, returns).
+    """
+    advantages = torch.zeros_like(rewards)
+    running = torch.zeros_like(rewards[0])
+    for t in reversed(range(rewards.shape[0])):
+        delta = rewards[t] + gamma * next_values[t] * (1 - terminated[t].float()) - values[t]
+        running = delta + gamma * gae_lambda * (1 - dones[t].float()) * running
+        advantages[t] = running
+    return advantages, advantages + values
+
+
+class RolloutSource:
+    """On-policy data source for DragonRLModel.fit.
+
+    Each iteration (= one training epoch) collects n_steps x n_envs fresh
+    transitions with the bound model, computes GAE, then yields
+    update_epochs x n_minibatches shuffled minibatches shaped
+        {'obs': ..., 'targets': {'actions', 'log_prob', 'values', 'advantages', 'returns'}}.
+    PPO: update_epochs ~4-10, n_minibatches ~4. A2C: update_epochs=1, n_minibatches=1.
+    """
+
+    def __init__(self, env_cfg: DragonEnvConfig, n_steps: int = 128, gamma: float = 0.99,
+                 gae_lambda: float = 0.95, update_epochs: int = 4, n_minibatches: int = 4,
+                 normalize_obs: bool = False, normalize_reward: bool = False):
+        self.collector = EnvCollector(env_cfg, normalize_obs, normalize_reward, gamma)
+        if (n_steps * self.collector.n_envs) < n_minibatches:
+            raise ValueError("n_steps * n_envs must be >= n_minibatches")
+        self.n_steps, self.gamma, self.gae_lambda = n_steps, gamma, gae_lambda
+        self.update_epochs, self.n_minibatches = update_epochs, n_minibatches
+        self.actor: Optional[ValueActor] = None
+        self._rng = torch.Generator().manual_seed(env_cfg.seed)
+
+    def bind(self, model: ValueActor) -> None:
+        if not isinstance(model, ValueActor):
+            raise TypeError(f"RolloutSource needs a model with act() and value(), got {type(model)}")
+        self.actor = model
+
+    @property
+    def steps_per_epoch(self) -> int:
+        return self.n_steps * self.collector.n_envs
+
+    @property
+    def global_step(self) -> int:
+        return self.collector.global_step
+
+    # ------- collection ------- #
+    def collect(self) -> dict[str, Any]:
+        """One rollout as [T, N, ...] tensors, plus advantages / returns."""
+        if self.actor is None:
+            raise RuntimeError("RolloutSource is unbound; fit()/evaluate() bind the model")
+        obs, actions, rewards, terminated, dones, log_probs, values = [], [], [], [], [], [], []
+        truncation_values = []  # (t, mask, V(final obs)) for truncated envs
+        for t in range(self.n_steps):
+            tr = self.collector.step(self.actor)
+            missing = {'log_prob', 'value'} - set(tr.extras)
+            if missing:
+                raise KeyError(f"act() must return extras {sorted(missing)} for on-policy training")
+            obs.append(obs_to_tensor(tr.obs))
+            actions.append(tr.action)
+            rewards.append(torch.as_tensor(tr.reward))
+            terminated.append(torch.as_tensor(tr.terminated))
+            dones.append(torch.as_tensor(tr.done))
+            log_probs.append(tr.extras['log_prob'])
+            values.append(tr.extras['value'])
+            if tr.truncated.any():
+                mask = torch.as_tensor(tr.truncated)
+                truncation_values.append((t, mask, self.actor.value(tr.next_obs)))
+
+        values_t = torch.stack(values).float()
+        next_values = torch.cat([values_t[1:], self.actor.value(self.collector.obs)[None].float()])
+        for t, mask, v in truncation_values:
+            next_values[t] = torch.where(mask, v.float(), next_values[t])
+        rewards_t = torch.stack(rewards)
+        advantages, returns = compute_gae(rewards_t, values_t, next_values, torch.stack(terminated),
+                                          torch.stack(dones), self.gamma, self.gae_lambda)
+        return {'obs': tree_stack(obs), 'actions': torch.stack(actions),
+                'log_prob': torch.stack(log_probs), 'values': values_t,
+                'advantages': advantages, 'returns': returns, 'rewards': rewards_t}
+
+    def __iter__(self) -> Iterator[dict[str, Any]]:
+        data = self.collect()
+        flat = lambda x: x.reshape(-1, *x.shape[2:])
+        obs = tree_map(flat, data['obs'])
+        targets = {k: flat(data[k]) for k in ('actions', 'log_prob', 'values', 'advantages', 'returns')}
+        total = self.steps_per_epoch
+        size = total // self.n_minibatches
+        for _ in range(self.update_epochs):
+            perm = torch.randperm(total, generator=self._rng)
+            for i in range(self.n_minibatches):
+                idx = perm[i * size:(i + 1) * size]
+                yield {'obs': tree_map(lambda x: x[idx], obs),
+                       'targets': {k: v[idx] for k, v in targets.items()}}
+
+    # ------- hooks ------- #
+    def metrics(self) -> dict[str, float]:
+        return self.collector.metrics()
+
+    def state_dict(self) -> dict[str, Any]:
+        return {'collector': self.collector.state_dict(), 'rng': self._rng.get_state()}
+
+    def load_state_dict(self, state: dict[str, Any]) -> None:
+        self.collector.load_state_dict(state['collector'])
+        self._rng.set_state(state['rng'])
+
+    def close(self) -> None:
+        self.collector.close()

dragon_torch/data/dataloaders.py

@@ -0,0 +1 @@
+# TODO: WRITE FUNCTION TO CREATE DATALOADER + SOURCE + DATASET. EXTEND DATALOADER TO INCLUDE AN ACCELERATED DATALOADER BY DEFAULT

dragon_torch/data/datasets/aggregator.py

@@ -0,0 +1,120 @@
+from typing import Any, Callable, Iterator, Sequence, Union
+import numpy as np
+from .dragon import DragonDataset
+from copy import deepcopy
+from torch.utils.data import IterableDataset, get_worker_info
+
+type WeightKey = Union[int, str]
+type AggregatorBatch = list[list[Any]]
+
+STREAMING_ATTRIBUTES = ['stream_from_hf']
+
+
+class Aggregator(IterableDataset):
+    """Combines map-style datasets into weighted batches.
+
+    Every batch holds int(weight * batch_size) samples from each dataset, drawn
+    without replacement, so an epoch ends as soon as any weighted dataset runs
+    out. Datasets are indexed per sample: pass DragonDatasets without a batch_size.
+
+    Weights are keyed by position in `datasets` or by a dataset's `kind`.
+    Call set_epoch() between epochs to reshuffle.
+    """
+
+    def __init__(self, datasets: Sequence[DragonDataset], batch_size: int, weights: dict[WeightKey, float],
+                 shuffle: bool = True, seed: int = None, collate_fn: Callable[[AggregatorBatch], Any] = None):
+        if len(datasets) == 0:
+            raise ValueError("Aggregator needs at least one dataset")
+        if batch_size < 1:
+            raise ValueError(f"batch_size must be >= 1, got {batch_size}")
+        self.datasets = self.validate_datasets(list(datasets))
+        self.batch_size = batch_size
+        self.weights = self._resolve_weights(weights)
+        self.samples_per_dataset = [int(w * batch_size) for w in self.weights]
+        if sum(self.samples_per_dataset) == 0:
+            raise ValueError(
+                f"weights={self.weights} draw no samples at batch_size={batch_size}")
+        self.shuffle = shuffle
+        self.seed = seed
+        self.collate_fn = collate_fn
+        # len() can be a network call for hugging face sources, so only ask once
+        self._lengths = [len(ds) for ds in self.datasets]
+        self.set_epoch(0)
+
+    def validate_datasets(self, datasets: list[DragonDataset]) -> list[DragonDataset]:
+        if not issubclass(datasets[0].__class__, DragonDataset):
+            return datasets
+        result = []
+        for dataset in datasets:
+            if not hasattr(dataset, 'to_aggregator') and not hasattr(dataset, 'config'):
+                continue
+            if hasattr(dataset, 'to_aggregator'):
+                dataset: DragonDataset = dataset.to_aggregator()
+            _cfg = dataset.config
+            if hasattr(_cfg, 'stream_from_hf'):
+                setattr(deepcopy(_cfg), 'stream_from_hf', False)
+                setattr(dataset, 'config', _cfg)
+                print(
+                    f'MSG :: dataset={dataset.kind}, Setting stream_from_hf=False for ')
+                result.append(dataset)
+            else:
+                result.append(dataset)
+        return result
+
+    def _index_of(self, key: WeightKey) -> int:
+        if isinstance(key, int):
+            if not 0 <= key < len(self.datasets):
+                raise KeyError(
+                    f"Weight key {key} out of range for {len(self.datasets)} datasets")
+            return key
+        matches = [i for i, ds in enumerate(self.datasets)
+                   if getattr(ds, 'kind', None) == key]
+        if len(matches) != 1:
+            raise KeyError(
+                f"Weight key '{key}' matches {len(matches)} datasets by kind, expected 1")
+        return matches[0]
+
+    def _resolve_weights(self, weights: dict[WeightKey, float]) -> list[float]:
+        if len(weights) != len(self.datasets):
+            raise ValueError(
+                f"Expected {len(self.datasets)} weights, got {len(weights)}")
+        resolved = [None] * len(self.datasets)
+        for key, weight in weights.items():
+            i = self._index_of(key)
+            if resolved[i] is not None:
+                raise ValueError(f"Dataset {i} was given more than one weight")
+            if not 0 <= weight <= 1:
+                raise ValueError(
+                    f"Weight for dataset {i} must be in [0, 1], got {weight}")
+            resolved[i] = float(weight)
+        if sum(resolved) > 1 + 1e-9:
+            raise ValueError(f"Weights sum to {sum(resolved)}, must be <= 1")
+        return resolved
+
+    def set_epoch(self, epoch: int):
+        self.epoch = epoch
+        rng = np.random.default_rng(
+            None if self.seed is None else [self.seed, epoch])
+        self._orders = [rng.permutation(n) if self.shuffle else np.arange(n)
+                        for n in self._lengths]
+
+    def __len__(self) -> int:
+        return min(n // k for n, k in zip(self._lengths, self.samples_per_dataset) if k > 0)
+
+    def __getitem__(self, identifier: int) -> AggregatorBatch:
+        n_batches = len(self)
+        idx = identifier + n_batches if identifier < 0 else identifier
+        if not 0 <= idx < n_batches:
+            raise IndexError(
+                f"Batch {identifier} out of range for {n_batches} batches")
+        batch = [[ds[int(j)] for j in order[idx * k:(idx + 1) * k]]
+                 for ds, order, k in zip(self.datasets, self._orders, self.samples_per_dataset)]
+        return batch if self.collate_fn is None else self.collate_fn(batch)
+
+    def __iter__(self) -> Iterator[AggregatorBatch]:
+        # under a multi-worker DataLoader each worker takes every num_workers-th batch
+        worker = get_worker_info()
+        indices = range(len(self)) if worker is None else range(
+            worker.id, len(self), worker.num_workers)
+        for idx in indices:
+            yield self[idx]

dragon_torch/data/datasets/dragon.py

@@ -0,0 +1,74 @@
+from typing import TypeVar, Union, Literal, ClassVar, Generic, Any, Callable, Self
+from pydantic import BaseModel
+import os
+from abc import ABC, abstractmethod
+from uuid import uuid4
+from datasets import Dataset
+
+from dragon_torch.common.types import PathLike, Path
+from dragon_torch.data.sources import Source, SourceConfig, SourceKind, source_factory
+
+SourceType = TypeVar('SourceType', covariant=True)
+ItemDtype = TypeVar('BatcherType', covariant=True)
+
+
+class _Dataset(Dataset, ABC, Generic[SourceType, ItemDtype]):
+    kind: ClassVar[str]
+    data: ClassVar[Any] = None
+    _registry: ClassVar[dict[str, type["_Dataset"]]] = {}
+
+    def __init_subclass__(cls, **kw: Any) -> None:
+        super().__init_subclass__(**kw)
+        if hasattr(cls, "kind"):
+            _Dataset._registry[cls.kind] = cls
+        else:
+            _Dataset._registry[cls.__name__] = cls
+
+    def describe(self, title: str = "Dragon Dataset Base Class", logger=None, level='info') -> None:
+        if logger is not None:
+            recorder = getattr(logger, level)
+        else:
+            recorder = print
+
+        recorder(f"Dataset Title={title}")
+        for key, value in self.config.__dict__.items():
+            recorder(f"    {key.capitalize()}={value}")
+
+    def set_data(self, _d):
+        setattr(self, 'data', _d)
+
+    @abstractmethod
+    def __len__(self) -> int: ...
+
+    @abstractmethod
+    def __getitem__(self, identifier: Union[int, str]) -> ItemDtype: ...
+
+
+class DragonDataset(_Dataset[SourceType, ItemDtype], Generic[SourceType, ItemDtype]):
+    def __init__(self, source_kind: SourceKind, source_cfg: SourceConfig, batch_size: int = None):
+        self.source = source_factory(source_cfg, kind=source_kind)
+        self.kind = source_kind
+        self.config = source_cfg
+        self.batch_size = batch_size if batch_size is not None else getattr(
+            source_cfg, 'batch_size', None)
+
+    def post_process(self, x):
+        return x if not hasattr(self.config, 'post_process') else self.config.post_process(x)
+
+    def load(self):
+        setattr(self, 'data', self.source.read())
+
+    def get_dataset(self) -> ItemDtype:
+        return self.data if not self.data is None else self.source.read()
+
+    def to_aggregator(self) -> Self:
+        setattr(self, 'batch_size', None)
+        return self
+
+    def __len__(self) -> int:
+        if self.batch_size is None:
+            return len(self.source)
+        return len(self.source) // self.batch_size
+
+    def __getitem__(self, identifier: Union[str, int]) -> ItemDtype:
+        return self.post_process(self.get_dataset().__getitem__(identifier))