---

name: embodied-dataflow description: "Build and troubleshoot DreamerV3 Embodied Env, Agent, Driver, Replay, Stream, wrapper, and custom environment dataflow contracts." disable-model-invocation: true metadata: disco-role: operating license: MIT

Embodied Dataflow for DreamerV3

Use this sub-skill when a task is about the Embodied runtime layer that feeds DreamerV3: environment adapters, custom Env contracts, Agent.policy() data shapes, Driver stepping/callbacks, replay insertion/sampling, streams, and wrappers. It is intentionally safe: it documents and checks data contracts, not training quality.

Route here for

• Building or validating a custom embodied.Env adapter.
• Explaining why act_space must include reset, why observations need reward, is_first, is_last, and is_terminal, or why log/ keys do not reach the policy.
• Wiring Driver(...).on_step(...) callbacks into replay with stable worker ids.
• Debugging replay samples, selectors, chunk save/load, stream composition, and sequence contiguity.
• Choosing and ordering wrappers such as NormalizeAction, ClipAction, UnifyDtypes, CheckSpaces, DiscretizeAction, ResizeImage, TimeLimit, and RestartOnException.
• Using the safe dummy environment as a contract example before attempting optional Atari, Crafter, DMC, DMLab, Gym, Minecraft, ProcGen, PinPad, or BSuite integrations.

Route elsewhere

• CLI flags, configs.yaml, run-loop selection, logging/checkpoints, and train smoke commands belong to sibling sub-skill train-configure.
• JAX, Ninjax, DreamerV3 model internals, optimizer state, sharding, and neural module debugging belong to sibling sub-skill jax-models.
• Installation, CUDA/system packages, optional environment dependency setup, Docker, plotting, score files, and metrics operations belong to sibling sub-skill results-ops.

Fast operating path

1. Read references/api-contracts.md before implementing or reviewing any Env, Agent, Driver, Replay, or Stream code. It contains the exact key, shape, callback, and sample/update contracts.

2. Read references/env-integration.md when adding a built-in or custom environment adapter, selecting wrappers, or deciding which optional dependency boundary applies.

3. Read references/replay-and-streams.md when replay sampling is empty, samples appear non-contiguous, priorities are updated, chunks are saved/loaded, or stream combinators are used.

4. Read references/troubleshooting.md when an error mentions reset, CheckSpaces, dtype/shape bounds, parallel workers, replay waiting, optional environment imports, or stream assertions.

5. Run scripts/check_embodied_contracts.py as a local contract smoke before training. It supports:

   python scripts/check_embodied_contracts.py --help
   python scripts/check_embodied_contracts.py --all
   python scripts/check_embodied_contracts.py --mode env
   python scripts/check_embodied_contracts.py --mode replay
   python scripts/check_embodied_contracts.py --mode driver
   python scripts/check_embodied_contracts.py --mode env --factory my_pkg.envs:make_env
   

   The checker uses a tiny dummy-like environment by default and can validate a user-provided zero-argument environment factory. It does not train.

Common recipes

Custom environment preflight

1. Implement obs_space, act_space, step(action), and close().
2. Include act_space['reset'] = elements.Space(bool) even if the wrapped backend resets internally.
3. Return reward, is_first, is_last, and is_terminal on every step.
4. Make step({'reset': True, ...}) return the first observation of a new episode with is_first=True, is_last=False, and reward normally 0.0.
5. Wrap with UnifyDtypes and then CheckSpaces during debugging.
6. Run the bundled checker in env mode before connecting a policy or replay.

Driver to replay wiring

import embodied
from embodied.envs import dummy

def make_env():
  return dummy.Dummy('disc', length=100)

env0 = make_env()
agent = embodied.RandomAgent(env0.obs_space, env0.act_space)
env0.close()

replay = embodied.Replay(length=16, capacity=1000)
driver = embodied.Driver([make_env, make_env], parallel=False)
driver.reset(agent.init_policy)
driver.on_step(lambda tran, worker: replay.add(tran, worker=worker))
driver(agent.policy, steps=64)
batch = replay.sample(batch=4, mode='train')
assert batch['is_first'].shape[:2] == (4, 16)

Keep the callback worker id when adding to replay. Worker ids keep sequences from separate environment streams from crossing.

Native verification candidates owned by this sub-skill

Use these as final native candidates after the whole DreamerV3 skill is integrated and its verification plan allows native tests:

• embodied/tests/test_driver.py::TestDriver::test_episode_length: driver reset/episode-length semantics on the dummy environment.
• embodied/tests/test_replay.py::TestReplay::test_sample_single: replay single-worker contiguous sequence sampling.
• Optional if resources allow: embodied/tests/test_parallel.py, because it exercises parallel run-loop dataflow and save/load behavior but is longer and depends on additional runtime services.

Do not run those native tests while merely using this runtime sub-skill; use the bundled checker for quick contract preflight.

Known limits

• Optional environment adapters can require external packages, ROMs, system libraries, display/Mujoco setup, or Java/MineRL resources. This sub-skill only defines the dataflow contract and adapter boundaries; setup belongs to results-ops.
• The checker validates contract shape and safe replay/driver mechanics. It does not prove training performance, JAX model correctness, or distributed run-loop behavior.
• For environments with unusual constructor arguments, expose a small zero-argument factory and pass it via --factory module:callable to the checker.