cyreal.sources

  1"""Streaming dataset sources inspired by Grain's paradigm."""
  2from __future__ import annotations
  3
  4from dataclasses import dataclass
  5from typing import Any, Callable, Literal, Protocol, TypeVar
  6
  7import jax
  8import jax.numpy as jnp
  9import numpy as np
 10from jax import tree_util
 11from jax.experimental import io_callback
 12
 13PyTree = Any
 14StateT = TypeVar("StateT")
 15
 16
 17class Source(Protocol[StateT]):
 18    """Interface for stateful, JIT-friendly data streams.
 19
 20    A source exposes fixed-length epochs (`steps_per_epoch`) and provides
 21    methods for initializing internal state and iteratively producing samples
 22    plus boolean masks that denote whether a sample is valid (e.g. under
 23    padding).
 24    """
 25
 26    steps_per_epoch: int
 27    """Number of items emitted per epoch."""
 28
 29    def init_state(self, key: jax.Array) -> StateT:
 30        """Return an initial state for the source.
 31
 32        Args:
 33            key: Optional PRNG key used for randomized behavior such as
 34                shuffling. Implementations should fall back to a default key
 35                when ``None`` is provided.
 36
 37        Returns:
 38            Backend-specific state object that must be passed to ``next``.
 39        """
 40        ...
 41
 42    def next(self, state: StateT) -> tuple[PyTree, jax.Array, StateT]:
 43        """Advance the stream and return the next value.
 44
 45        Args:
 46            state: Previously-initialized source state.
 47
 48        Returns:
 49            Tuple ``(value, mask, new_state)`` where ``value`` is a PyTree of
 50            arrays, ``mask`` is a boolean array indicating whether the sample
 51            is valid, and ``new_state`` should be provided to the next call.
 52        """
 53        ...
 54
 55    def element_spec(self) -> PyTree:
 56        """PyTree of :class:`jax.ShapeDtypeStruct` describing emitted samples."""
 57        ...
 58
 59
 60@jax.tree_util.register_pytree_node_class
 61@dataclass
 62class _ArraySourceState:
 63    indices: jax.Array
 64    mask: jax.Array
 65    position: jax.Array
 66    key: jax.Array
 67    epoch: jax.Array
 68
 69    def tree_flatten(self):
 70        return (self.indices, self.mask, self.position, self.key, self.epoch), None
 71
 72    @classmethod
 73    def tree_unflatten(cls, aux_data, children):
 74        indices, mask, position, key, epoch = children
 75        return cls(indices=indices, mask=mask, position=position, key=key, epoch=epoch)
 76
 77
 78@jax.tree_util.register_pytree_node_class
 79@dataclass
 80class _DiskSourceState:
 81    indices: jax.Array
 82    position: jax.Array
 83    key: jax.Array
 84    epoch: jax.Array
 85    buffer: PyTree
 86    buffer_pos: jax.Array
 87    buffer_count: jax.Array
 88
 89    def tree_flatten(self):
 90        buffer_leaves, buffer_def = tree_util.tree_flatten(self.buffer)
 91        children = (
 92            self.indices,
 93            self.position,
 94            self.key,
 95            self.epoch,
 96            *buffer_leaves,
 97            self.buffer_pos,
 98            self.buffer_count,
 99        )
100        return children, buffer_def
101
102    @classmethod
103    def tree_unflatten(cls, buffer_def, children):
104        indices, position, key, epoch, *rest = children
105        buffer_leaf_count = buffer_def.num_leaves if buffer_def is not None else 0
106        buffer_leaves = rest[:buffer_leaf_count]
107        buffer = (
108            tree_util.tree_unflatten(buffer_def, buffer_leaves)
109            if buffer_def is not None
110            else None
111        )
112        buffer_pos, buffer_count = rest[buffer_leaf_count:]
113        return cls(
114            indices=indices,
115            position=position,
116            key=key,
117            epoch=epoch,
118            buffer=buffer,
119            buffer_pos=buffer_pos,
120            buffer_count=buffer_count,
121        )
122
123
124@dataclass
125class ArraySource(Source[_ArraySourceState]):
126    """Sample-level stream over an in-memory PyTree of arrays.
127
128    Args:
129        data: PyTree whose leaves are arrays with a leading sample dimension.
130        ordering: Either ``"sequential"`` or ``"shuffle"``.
131    """
132
133    data: PyTree
134    ordering: Literal["sequential", "shuffle"] = "shuffle"
135
136    def __post_init__(self) -> None:
137        leaves, self._treedef = tree_util.tree_flatten(self.data)
138        if not leaves:
139            raise ValueError("Data tree must contain at least one array.")
140
141        first = leaves[0]
142        self._num_samples = int(first.shape[0])
143        for leaf in leaves[1:]:
144            if leaf.shape[0] != self._num_samples:
145                raise ValueError("All leaves must share the leading dimension.")
146        if self._num_samples == 0:
147            raise ValueError("Dataset cannot be empty.")
148
149        self.steps_per_epoch = self._num_samples
150        self._mask_template = jnp.ones(self._num_samples, dtype=bool)
151        self._element_spec = tree_util.tree_map(
152            lambda leaf: jax.ShapeDtypeStruct(shape=leaf.shape[1:], dtype=leaf.dtype),
153            self.data,
154        )
155
156    @property
157    def num_samples(self) -> int:
158        return self._num_samples
159
160    def element_spec(self) -> PyTree:
161        """Shape/dtype metadata describing samples produced by the source."""
162        return self._element_spec
163
164    def _build_epoch_indices(self, key: jax.Array) -> tuple[jax.Array, jax.Array]:
165        base = jnp.arange(self._num_samples)
166        if self.ordering == "shuffle":
167            base = jax.random.permutation(key, base)
168        elif self.ordering != "sequential":
169            raise ValueError(f"Unknown ordering '{self.ordering}'.")
170
171        return base, self._mask_template
172
173    def init_state(self, key: jax.Array) -> _ArraySourceState:
174        """Create the initial iteration state.
175
176        Args:
177            key: Optional PRNG key. Defaults to ``jax.random.PRNGKey(0)`` when
178                omitted.
179        """
180        key, perm_key = jax.random.split(key)
181        indices, mask = self._build_epoch_indices(perm_key)
182        position = jnp.array(0, dtype=jnp.int32)
183        epoch = jnp.array(0, dtype=jnp.int32)
184        return _ArraySourceState(indices=indices, mask=mask, position=position, key=key, epoch=epoch)
185
186    def next(self, state: _ArraySourceState) -> tuple[PyTree, jax.Array, _ArraySourceState]:
187        """Return the next sample (with mask) and the advanced state."""
188        index = jax.lax.dynamic_index_in_dim(state.indices, state.position, axis=0, keepdims=False)
189        mask_value = jax.lax.dynamic_index_in_dim(state.mask, state.position, axis=0, keepdims=False)
190        sample = tree_util.tree_map(
191            lambda arr: jax.lax.dynamic_index_in_dim(arr, index, axis=0, keepdims=False),
192            self.data,
193        )
194
195        next_position = state.position + 1
196
197        def _reset_epoch(_: None):
198            new_key, perm_key = jax.random.split(state.key)
199            indices, mask = self._build_epoch_indices(perm_key)
200            return _ArraySourceState(
201                indices=indices,
202                mask=mask,
203                position=jnp.array(0, dtype=jnp.int32),
204                key=new_key,
205                epoch=state.epoch + 1,
206            )
207
208        def _advance(_: None):
209            return _ArraySourceState(
210                indices=state.indices,
211                mask=state.mask,
212                position=next_position,
213                key=state.key,
214                epoch=state.epoch,
215            )
216
217        need_reset = next_position >= state.indices.shape[0]
218        new_state = jax.lax.cond(need_reset, _reset_epoch, _advance, operand=None)
219        return sample, mask_value, new_state
220
221
222@dataclass
223class DiskSource(Source[_DiskSourceState]):
224    """Sample-level stream that loads items via a Python callback (disk, RPC, etc.).
225
226    This is slow, only use this if your dataset will not fit in system memory.
227    """
228
229    length: int
230    """Number of samples in the dataset."""
231    sample_fn: Callable[[int], PyTree]
232    """Python callable that takes an integer index and returns a PyTree of arrays."""
233    sample_spec: PyTree | None = None
234    """Optional PyTree of `jax.ShapeDtypeStruct` describing the shape and dtype of samples."""
235    ordering: Literal["sequential", "shuffle"] = "shuffle"
236    """Sample ordering strategy, either 'sequential' or 'shuffle'. The shuffling occurs over the entire dataset, not within the prefetch buffer."""
237    prefetch_size: int = 64
238    """Number of samples to prefetch into a JAX array buffer. Set this larger to achieve better throughput at the cost of more memory usage."""
239
240    def __post_init__(self) -> None:
241        if self.length <= 0:
242            raise ValueError("Dataset cannot be empty.")
243        if self.prefetch_size <= 0:
244            raise ValueError("prefetch_size must be positive.")
245
246        if self.sample_spec is None:
247            example = self.sample_fn(0)
248
249            def _to_spec(leaf):
250                arr = np.asarray(leaf)
251                return jax.ShapeDtypeStruct(shape=arr.shape, dtype=arr.dtype)
252
253            self.sample_spec = tree_util.tree_map(_to_spec, example)
254
255        leaves = tree_util.tree_leaves(self.sample_spec)
256        if not leaves:
257            raise ValueError("element_spec must include at least one leaf.")
258        for leaf in leaves:
259            if not isinstance(leaf, jax.ShapeDtypeStruct):
260                raise TypeError("element_spec leaves must be jax.ShapeDtypeStruct instances.")
261
262        self._num_samples = int(self.length)
263        self.steps_per_epoch = self._num_samples
264        self._element_spec = self.sample_spec
265        self.prefetch_size = int(self.prefetch_size)
266
267        def _zeros(spec: jax.ShapeDtypeStruct):
268            return np.zeros(spec.shape, dtype=np.dtype(spec.dtype))
269
270        def _buffer_shape(spec: jax.ShapeDtypeStruct):
271            return jax.ShapeDtypeStruct(
272                shape=(self.prefetch_size, *spec.shape),
273                dtype=spec.dtype,
274            )
275
276        self._zero_sample = tree_util.tree_map(_zeros, self.sample_spec)
277        self._chunk_spec = tree_util.tree_map(_buffer_shape, self.sample_spec)
278        self._buffer_template = tree_util.tree_map(
279            lambda spec: jnp.zeros((self.prefetch_size, *spec.shape), dtype=spec.dtype),
280            self.sample_spec,
281        )
282
283    def element_spec(self) -> PyTree:
284        """Shape/dtype metadata describing samples produced by the source."""
285        return self._element_spec
286
287    def _build_epoch_indices(self, key: jax.Array) -> jax.Array:
288        base = jnp.arange(self._num_samples)
289        if self.ordering == "shuffle":
290            base = jax.random.permutation(key, base)
291        elif self.ordering != "sequential":
292            raise ValueError(f"Unknown ordering '{self.ordering}'.")
293        return base
294
295    def init_state(self, key: jax.Array) -> _DiskSourceState:
296        """Build the starting state, optionally seeding randomness with ``key``."""
297        key, perm_key = jax.random.split(key)
298        indices = self._build_epoch_indices(perm_key)
299        position = jnp.array(0, dtype=jnp.int32)
300        epoch = jnp.array(0, dtype=jnp.int32)
301        return _DiskSourceState(
302            indices=indices,
303            position=position,
304            key=key,
305            epoch=epoch,
306            buffer=self._buffer_template,
307            buffer_pos=jnp.array(0, dtype=jnp.int32),
308            buffer_count=jnp.array(0, dtype=jnp.int32),
309        )
310
311    def _chunk_callback(self, indices: np.ndarray, mask: np.ndarray) -> PyTree:
312        idx_array = np.asarray(indices, dtype=np.int64)
313        mask_array = np.asarray(mask, dtype=bool)
314        samples: list[PyTree] = []
315        for keep, idx in zip(mask_array, idx_array):
316            if keep:
317                samples.append(self.sample_fn(int(idx)))
318            else:
319                samples.append(self._zero_sample)
320        return tree_util.tree_map(lambda *xs: np.stack(xs, axis=0), *samples)
321
322    def _maybe_reset_epoch(self, state: _DiskSourceState) -> _DiskSourceState:
323        def _reset(state: _DiskSourceState):
324            new_key, perm_key = jax.random.split(state.key)
325            indices = self._build_epoch_indices(perm_key)
326            return _DiskSourceState(
327                indices=indices,
328                position=jnp.array(0, dtype=jnp.int32),
329                key=new_key,
330                epoch=state.epoch + 1,
331                buffer=self._buffer_template,
332                buffer_pos=jnp.array(0, dtype=jnp.int32),
333                buffer_count=jnp.array(0, dtype=jnp.int32),
334            )
335
336        return jax.lax.cond(state.position >= self._num_samples, _reset, lambda s: s, state)
337
338    def _maybe_refill_buffer(self, state: _DiskSourceState) -> _DiskSourceState:
339        def _needs(state: _DiskSourceState):
340            return jnp.logical_or(state.buffer_count == 0, state.buffer_pos >= state.buffer_count)
341
342        def _refill(state: _DiskSourceState):
343            refreshed = self._maybe_reset_epoch(state)
344            remaining = self._num_samples - refreshed.position
345            chunk = jnp.minimum(remaining, self.prefetch_size)
346            chunk = jnp.maximum(chunk, 0)
347            chunk = chunk.astype(jnp.int32)
348            offsets = jnp.arange(self.prefetch_size, dtype=jnp.int32)
349            gather_positions = jnp.minimum(
350                refreshed.position + offsets,
351                refreshed.indices.shape[0] - 1,
352            )
353            chunk_indices = jax.vmap(
354                lambda idx: jax.lax.dynamic_index_in_dim(
355                    refreshed.indices, idx, axis=0, keepdims=False
356                )
357            )(gather_positions)
358            valid_mask = offsets < chunk
359            buffer = io_callback(self._chunk_callback, self._chunk_spec, chunk_indices, valid_mask)
360            new_position = refreshed.position + chunk
361            return _DiskSourceState(
362                indices=refreshed.indices,
363                position=new_position,
364                key=refreshed.key,
365                epoch=refreshed.epoch,
366                buffer=buffer,
367                buffer_pos=jnp.array(0, dtype=jnp.int32),
368                buffer_count=chunk,
369            )
370
371        return jax.lax.cond(_needs(state), _refill, lambda s: s, state)
372
373    def next(self, state: _DiskSourceState) -> tuple[PyTree, jax.Array, _DiskSourceState]:
374        """Return buffered sample, all-True mask, and updated state."""
375        state = self._maybe_refill_buffer(state)
376        sample = tree_util.tree_map(
377            lambda buf: jax.lax.dynamic_index_in_dim(
378                buf, state.buffer_pos, axis=0, keepdims=False
379            ),
380            state.buffer,
381        )
382        mask_value = jnp.array(True, dtype=bool)
383        new_state = _DiskSourceState(
384            indices=state.indices,
385            position=state.position,
386            key=state.key,
387            epoch=state.epoch,
388            buffer=state.buffer,
389            buffer_pos=state.buffer_pos + 1,
390            buffer_count=state.buffer_count,
391        )
392        return sample, mask_value, new_state
393
394
395
396
397@jax.tree_util.register_pytree_node_class
398@dataclass
399class GymnaxSourceState:
400    env_state: PyTree
401    obs: PyTree
402    key: jax.Array
403    step: jax.Array
404    epoch: jax.Array
405    policy_state: PyTree | None = None
406    new_episode: jax.Array | None = None
407
408    def tree_flatten(self):
409        return (
410            self.env_state,
411            self.obs,
412            self.key,
413            self.step,
414            self.epoch,
415            self.policy_state,
416            self.new_episode,
417        ), None
418
419    @classmethod
420    def tree_unflatten(cls, aux_data, children):
421        env_state, obs, key, step, epoch, policy_state, new_episode = children
422        return cls(
423            env_state=env_state,
424            obs=obs,
425            key=key,
426            step=step,
427            epoch=epoch,
428            policy_state=policy_state,
429            new_episode=new_episode,
430        )
431
432
433@dataclass
434class GymnaxSource(Source[GymnaxSourceState]):
435    """Stream transitions by rolling out a Gymnax environment with a policy.
436    
437    Useful for reinforcement learning.
438
439    Args:
440        env: Gymnax environment instance.
441        env_params: Parameters to pass to the environment's reset and step functions.
442        policy_step_fn: Callable that takes (observation, policy_state, new_episode, key) and
443            returns (action, new_policy_state).
444        policy_state_template: Example PyTree carrying everything required by
445            ``policy_step_fn`` (for example, policy parameters and recurrent
446            carries). This template is used only to infer the element spec; callers
447            are responsible for injecting a real policy state into the loader
448            state before calling ``next``.
449        steps_per_epoch: Number of environment steps per epoch for a single environment.
450    """
451
452    env: Any
453    env_params: Any
454    policy_step_fn: Callable[[PyTree, PyTree, jax.Array, jax.Array], tuple[PyTree, PyTree]]
455    policy_state_template: PyTree | None = None
456    steps_per_epoch: int = 1024
457
458    def __post_init__(self) -> None:
459        if self.steps_per_epoch <= 0:
460            raise ValueError("steps_per_epoch must be positive.")
461        if self.policy_state_template is None:
462            raise ValueError("GymnaxSource requires a policy_state_template for shape inference.")
463
464        def _sample(key, policy_state):
465            obs, env_state = self.env.reset(key, self.env_params)
466            action, next_policy_state = self.policy_step_fn(
467                obs,
468                policy_state,
469                jnp.array(True, dtype=jnp.bool_),
470                key,
471            )
472            next_obs, _, reward, done, info = self.env.step(
473                key,
474                env_state,
475                action,
476                self.env_params,
477            )
478            transition = {
479                "state": obs,
480                "action": action,
481                "reward": reward,
482                "next_state": next_obs,
483                "done": done,
484                "info": info,
485            }
486            return transition, next_policy_state
487
488        shaped, _ = jax.eval_shape(_sample, jax.random.PRNGKey(0), self.policy_state_template)
489        self._element_spec = tree_util.tree_map(
490            lambda arr: jax.ShapeDtypeStruct(shape=arr.shape, dtype=arr.dtype), shaped
491        )
492        self.policy_state_template = None
493
494    def element_spec(self) -> PyTree:
495        """Shape/dtype metadata describing Gymnax transitions."""
496        return self._element_spec
497
498    def init_state(self, key: jax.Array) -> GymnaxSourceState:
499        """Return RNG-seeded environment + policy state for iteration."""
500        key, env_key = jax.random.split(key)
501        obs, env_state = self.env.reset(env_key, self.env_params)
502        return GymnaxSourceState(
503            env_state=env_state,
504            obs=obs,
505            key=key,
506            step=jnp.array(0, dtype=jnp.int32),
507            epoch=jnp.array(0, dtype=jnp.int32),
508            policy_state=None,
509            new_episode=jnp.array(True, dtype=jnp.bool_),
510        )
511
512    def next(self, state: GymnaxSourceState) -> tuple[PyTree, jax.Array, GymnaxSourceState]:
513        """Roll the environment forward one step and emit a transition."""
514        key, policy_key, step_key, done_reset_key, epoch_reset_key = jax.random.split(state.key, 5)
515
516        if state.policy_state is None:
517            raise ValueError(
518                "GymnaxSource state is missing `policy_state`; set it explicitly before calling `next`."
519            )
520        policy_state = state.policy_state
521
522        if state.new_episode is None:
523            raise ValueError("GymnaxSource state is missing `new_episode` flag.")
524
525        action, updated_policy_state = self.policy_step_fn(
526            state.obs,
527            policy_state,
528            state.new_episode,
529            policy_key,
530        )
531        next_obs, next_env_state, reward, done, info = self.env.step(
532            step_key,
533            state.env_state,
534            action,
535            self.env_params,
536        )
537
538        transition = {
539            "state": state.obs,
540            "action": action,
541            "reward": reward,
542            "next_state": next_obs,
543            "done": done,
544            "info": info,
545        }
546        mask = jnp.array(True, dtype=bool)
547
548        done_flag = jnp.asarray(done, dtype=bool)
549        done_flag = jnp.reshape(done_flag, ())
550        reset_obs, reset_env_state = self.env.reset(done_reset_key, self.env_params)
551
552        cont_obs, cont_env_state = jax.lax.cond(
553            done_flag,
554            lambda _: (reset_obs, reset_env_state),
555            lambda _: (next_obs, next_env_state),
556            operand=None,
557        )
558
559        next_step = state.step + 1
560        need_epoch_reset = next_step >= self.steps_per_epoch
561
562        def _reset_epoch(_: None):
563            epoch_obs, epoch_env_state = self.env.reset(epoch_reset_key, self.env_params)
564            return GymnaxSourceState(
565                env_state=epoch_env_state,
566                obs=epoch_obs,
567                key=key,
568                step=jnp.array(0, dtype=jnp.int32),
569                epoch=state.epoch + 1,
570                policy_state=updated_policy_state,
571                new_episode=jnp.array(True, dtype=jnp.bool_),
572            )
573
574        def _continue(_: None):
575            return GymnaxSourceState(
576                env_state=cont_env_state,
577                obs=cont_obs,
578                key=key,
579                step=next_step,
580                epoch=state.epoch,
581                policy_state=updated_policy_state,
582                new_episode=done_flag,
583            )
584
585        new_state = jax.lax.cond(need_epoch_reset, _reset_epoch, _continue, operand=None)
586        return transition, mask, new_state