xwm.data¶

Batch streams and a synthetic controllable world: a sprite the actions move, plus distractors that move on their own. Enough to train and plan on CPU.

Data: batch streams and a synthetic controllable world.

Modules:

Name	Description
`batching`	Turning arrays into batch streams.
`synthetic`	A small controllable world, for tests and runnable examples.

Classes:

Name	Description
`SpriteState`	World state: a controlled agent plus `n_distractors` random walkers.
`SpriteWorld`	A 2-D world with one action-controlled sprite and some uncontrolled ones.

Functions:

Name	Description
`clip_windows`	Cut `(T, ...)` frames into overlapping clips of `window` frames.
`iter_batches`	Iterate mini-batches over a dict of equally-long arrays.
`random_actions`	`(n_sequences, length, action_dim)` smoothly correlated random actions.
`sprite_images`	Generate a still-image dataset: `(n, 3, size, size)` plus positions.
`sprite_sequences`	Generate an action-labelled video dataset.

SpriteState ¶

Bases: NamedTuple

World state: a controlled agent plus n_distractors random walkers.

SpriteWorld ¶

SpriteWorld(size: int = 32, *, n_distractors: int = 2, radius: float = 0.1, dt: float = 1.0, damping: float = 0.5, action_scale: float = 0.15, distractor_speed: float = 0.05)

Bases: Module

A 2-D world with one action-controlled sprite and some uncontrolled ones.

Parameters:

Name	Type	Description	Default
`size`	`int`	rendered resolution (square).	`32`
`n_distractors`	`int`	uncontrollable sprites performing a random walk.	`2`
`radius`	`float`	sprite radius in normalised units.	`0.1`
`dt`	`float`	integration step.	`1.0`
`damping`	`float`	velocity decay per step; `1.0` is frictionless.	`0.5`
`action_scale`	`float`	acceleration applied per unit of action.	`0.15`
`distractor_speed`	`float`	random-walk step size for the distractors.	`0.05`

The defaults make one step of full action displace the agent by about 1.5 radii. That matters for evaluation: with slower dynamics, a single step barely changes the image, "predict no change" becomes a near-optimal one-step baseline, and a latent dynamics model looks worthless at short horizons for reasons that have nothing to do with the model.

Methods:

Name	Description
`step`	Advance the world. `action` is a 2-D acceleration in `[-1, 1]^2`.
`render`	Render one state to `(3, size, size)` in `[0, 1]`.
`rollout`	Run `(T, 2)` actions from a random start.
`observe`	Just the frames from `rollout`.

Source code in xwm/data/synthetic.py

def __init__(
    self,
    size: int = 32,
    *,
    n_distractors: int = 2,
    radius: float = 0.10,
    dt: float = 1.0,
    damping: float = 0.5,
    action_scale: float = 0.15,
    distractor_speed: float = 0.05,
):
    self.size = size
    self.n_distractors = n_distractors
    self.radius = radius
    self.dt = dt
    self.damping = damping
    self.action_scale = action_scale
    self.distractor_speed = distractor_speed

step ¶

step(state: SpriteState, action: Array, *, key: PRNGKey | None = None) -> SpriteState

Advance the world. action is a 2-D acceleration in [-1, 1]^2.

Positions reflect off the walls, which keeps trajectories bounded without the discontinuity of wrapping.

Source code in xwm/data/synthetic.py

def step(self, state: SpriteState, action: Array, *, key: PRNGKey | None = None) -> SpriteState:
    """Advance the world. ``action`` is a 2-D acceleration in ``[-1, 1]^2``.

    Positions reflect off the walls, which keeps trajectories bounded
    without the discontinuity of wrapping.
    """
    vel = self.damping * state.vel + self.action_scale * jnp.clip(action, -1.0, 1.0)
    pos = state.pos + self.dt * vel
    # Reflect at the boundary and flip the corresponding velocity component.
    below, above = pos < 0.0, pos > 1.0
    pos = jnp.where(below, -pos, jnp.where(above, 2.0 - pos, pos))
    vel = jnp.where(below | above, -vel, vel)
    if key is None or self.n_distractors == 0:
        distractors = state.distractors
    else:
        walk = self.distractor_speed * jr.normal(key, state.distractors.shape)
        distractors = jnp.clip(state.distractors + walk, 0.0, 1.0)
    return SpriteState(pos=pos, vel=vel, distractors=distractors)

render ¶

render(state: SpriteState) -> Array

Render one state to (3, size, size) in [0, 1].

The agent occupies the red channel and the distractors the green one, so a probe can tell trivially whether a representation kept the controllable content, the uncontrollable content, or both.

Source code in xwm/data/synthetic.py

def render(self, state: SpriteState) -> Array:
    """Render one state to ``(3, size, size)`` in ``[0, 1]``.

    The agent occupies the red channel and the distractors the green one, so
    a probe can tell trivially whether a representation kept the
    controllable content, the uncontrollable content, or both.
    """
    agent = self._blob(state.pos)
    if self.n_distractors:
        distractors = jnp.max(jax.vmap(self._blob)(state.distractors), axis=0)
    else:
        distractors = jnp.zeros_like(agent)
    grid = (jnp.arange(self.size) + 0.5) / self.size
    background = 0.15 * (grid[:, None] + grid[None, :]) / 2.0
    return jnp.clip(jnp.stack([agent, distractors, background + 0.0 * agent]), 0.0, 1.0)

rollout ¶

rollout(key: PRNGKey, actions: Array) -> tuple[Array, SpriteState]

Run (T, 2) actions from a random start.

Returns (frames, states) with frames of shape (T + 1, 3, size, size) -- one more frame than actions, since the initial observation precedes the first action.

Source code in xwm/data/synthetic.py

def rollout(self, key: PRNGKey, actions: Array) -> tuple[Array, SpriteState]:
    """Run ``(T, 2)`` actions from a random start.

    Returns ``(frames, states)`` with ``frames`` of shape
    ``(T + 1, 3, size, size)`` -- one more frame than actions, since the
    initial observation precedes the first action.
    """
    k_reset, k_steps = jr.split(key)
    state = self.reset(k_reset)
    keys = jr.split(k_steps, actions.shape[0])

    def advance(s, inputs):
        a, k = inputs
        s = self.step(s, a, key=k)
        return s, s

    final, states = jax.lax.scan(advance, state, (actions, keys))
    all_states = jax.tree_util.tree_map(
        lambda first, rest: jnp.concatenate([first[None], rest]), state, states
    )
    return jax.vmap(self.render)(all_states), all_states

observe ¶

observe(key: PRNGKey, actions: Array) -> Array

Just the frames from rollout.

Source code in xwm/data/synthetic.py

def observe(self, key: PRNGKey, actions: Array) -> Array:
    """Just the frames from :meth:`rollout`."""
    return self.rollout(key, actions)[0]

clip_windows ¶

clip_windows(video: Array, window: int, *, stride: int = 1) -> Array

Cut (T, ...) frames into overlapping clips of window frames.

Returns (n_windows, window, ...). Video encoders take fixed-length clips, so this is how a long sequence is fed to one.

Source code in xwm/data/batching.py

def clip_windows(video: Array, window: int, *, stride: int = 1) -> Array:
    """Cut ``(T, ...)`` frames into overlapping clips of ``window`` frames.

    Returns ``(n_windows, window, ...)``. Video encoders take fixed-length
    clips, so this is how a long sequence is fed to one.
    """
    t = video.shape[0]
    if window > t:
        raise ValueError(f"window={window} exceeds sequence length {t}")
    starts = jnp.arange(0, t - window + 1, stride)
    return jnp.stack([video[s : s + window] for s in starts])

iter_batches ¶

iter_batches(data: dict[str, Array], batch_size: int, *, key: PRNGKey | None = None, shuffle: bool = True, drop_last: bool = True, epochs: int | None = 1) -> Iterator[Batch]

Iterate mini-batches over a dict of equally-long arrays.

Parameters:

Name	Type	Description	Default
`data`	`dict[str, Array]`	field name -> array with a common leading axis.	required
`epochs`	`int \| None`	passes over the data; `None` repeats forever, which is what `xwm.training.Trainer.fit` wants when driven by `steps`.	`1`

Source code in xwm/data/batching.py

def iter_batches(
    data: dict[str, Array],
    batch_size: int,
    *,
    key: PRNGKey | None = None,
    shuffle: bool = True,
    drop_last: bool = True,
    epochs: int | None = 1,
) -> Iterator[Batch]:
    """Iterate mini-batches over a dict of equally-long arrays.

    Args:
        data: field name -> array with a common leading axis.
        epochs: passes over the data; ``None`` repeats forever, which is what
            :meth:`xwm.training.Trainer.fit` wants when driven by ``steps``.
    """
    sizes = {k: v.shape[0] for k, v in data.items()}
    if len(set(sizes.values())) != 1:
        raise ValueError(f"fields disagree on length: {sizes}")
    n = next(iter(sizes.values()))
    if batch_size > n:
        raise ValueError(f"batch_size={batch_size} exceeds dataset size {n}")

    epoch = 0
    while epochs is None or epoch < epochs:
        order = jnp.arange(n)
        if shuffle:
            if key is None:
                raise ValueError("shuffle=True needs a key")
            order = jr.permutation(jr.fold_in(key, epoch), n)
        limit = n - n % batch_size if drop_last else n
        for start in range(0, limit, batch_size):
            idx = order[start : start + batch_size]
            yield {k: v[idx] for k, v in data.items()}
        epoch += 1

random_actions ¶

random_actions(key: PRNGKey, n_sequences: int, length: int, action_dim: int = 2, *, smoothness: float = 0.7) -> Array

(n_sequences, length, action_dim) smoothly correlated random actions.

White-noise actions make an almost unlearnable dataset -- the agent jitters in place and no action has visible consequences. Temporally correlated actions (an AR(1) process, smoothness being the correlation) produce trajectories that actually go somewhere.

Source code in xwm/data/synthetic.py

def random_actions(
    key: PRNGKey,
    n_sequences: int,
    length: int,
    action_dim: int = 2,
    *,
    smoothness: float = 0.7,
) -> Array:
    """``(n_sequences, length, action_dim)`` smoothly correlated random actions.

    White-noise actions make an almost unlearnable dataset -- the agent jitters
    in place and no action has visible consequences. Temporally correlated
    actions (an AR(1) process, ``smoothness`` being the correlation) produce
    trajectories that actually go somewhere.
    """
    noise = jr.uniform(key, (n_sequences, length, action_dim), minval=-1.0, maxval=1.0)

    def smooth(carry, x):
        carry = smoothness * carry + (1.0 - smoothness) * x
        return carry, carry

    _, out = jax.lax.scan(smooth, noise[:, 0], noise.transpose(1, 0, 2))
    return out.transpose(1, 0, 2)

sprite_images ¶

sprite_images(key: PRNGKey, n_images: int, *, world: SpriteWorld | None = None) -> dict[str, Array]

Generate a still-image dataset: (n, 3, size, size) plus positions.

Source code in xwm/data/synthetic.py

def sprite_images(
    key: PRNGKey,
    n_images: int,
    *,
    world: SpriteWorld | None = None,
) -> dict[str, Array]:
    """Generate a still-image dataset: ``(n, 3, size, size)`` plus positions."""
    world = world or SpriteWorld()
    states = jax.vmap(world.reset)(jr.split(key, n_images))
    return {"image": jax.vmap(world.render)(states), "position": states.pos}

sprite_sequences ¶

sprite_sequences(key: PRNGKey, n_sequences: int, length: int, *, world: SpriteWorld | None = None, smoothness: float = 0.7) -> dict[str, Array]

Generate an action-labelled video dataset.

Returns a dict with:

video: (n, length, 3, size, size)
action: (n, length - 1, 2) -- action[i, t] joins frames t and t + 1, the convention xwm.action.ActionWorldModel expects.
position: (n, length, 2) ground-truth agent position, for probes.

Source code in xwm/data/synthetic.py

def sprite_sequences(
    key: PRNGKey,
    n_sequences: int,
    length: int,
    *,
    world: SpriteWorld | None = None,
    smoothness: float = 0.7,
) -> dict[str, Array]:
    """Generate an action-labelled video dataset.

    Returns a dict with:

    * ``video``: ``(n, length, 3, size, size)``
    * ``action``: ``(n, length - 1, 2)`` -- ``action[i, t]`` joins frames
      ``t`` and ``t + 1``, the convention :class:`xwm.action.ActionWorldModel`
      expects.
    * ``position``: ``(n, length, 2)`` ground-truth agent position, for probes.
    """
    world = world or SpriteWorld()
    k_act, k_env = jr.split(key)
    actions = random_actions(
        k_act, n_sequences, length - 1, world.action_dim, smoothness=smoothness
    )
    keys = jr.split(k_env, n_sequences)
    frames, states = jax.vmap(world.rollout)(keys, actions)
    return {"video": frames, "action": actions, "position": states.pos}