DeepMind BSuite#

The BSuite backend adapts DeepMind’s Behaviour Suite for Reinforcement Learning (BSuite) diagnostic tasks to TorchWM’s image-first interface. BSuite tasks provide compact dm_env observations and mostly discrete actions with known optimal reward profiles, making them useful for detecting fundamental issues in world-model training (e.g. credit assignment, memory, exploration).

Install: pip install bsuite

Main API#

from torchwm import make_bsuite_env, list_available_bsuite_ids

env = make_bsuite_env("catch/0", seed=42, size=(64, 64))
obs = env.reset()
print(obs["image"].shape)  # (3, 64, 64)

# List all installed BSuite task IDs
print(list_available_bsuite_ids()[:5])

When BSuite is not installed, list_available_bsuite_ids() returns example IDs: bandit/0, cartpole/0, catch/0, deep_sea/0, discounting_chain/0, memory_len/0, mnist/0, mountain_car/0, umbrella_chain/0.

Seed determinism#

BSuiteImageEnv accepts a seed parameter at construction. The seed controls the internal RNG for action-space sampling in _BSuiteDiscreteActionSpace. BSuite itself does not expose a native seed API (its randomness is encoded in the bsuite_id), so reset(seed=...) reseeds only the wrapper’s action-space RNG:

env = make_bsuite_env("catch/0", seed=10)
action_a = env.action_space.sample()  # deterministic for seed=10

env.reset(seed=20)
action_b = env.action_space.sample()  # deterministic for seed=20

Observations#

BSuiteImageEnv converts BSuite’s compact dm_env observations into synthetic RGB images:

  • Native BSuite observations are flattened and normalised to [0, 1].

  • The flattened array is arranged onto a square canvas and repeated to 3 channels.

  • The canvas is resized to (H, W) using bilinear interpolation.

  • The result is returned as {"image": uint8 array with shape (3, H, W)}.

When include_state=True, observations also include a "state" key with the flattened raw observation vector.

env = make_bsuite_env("catch/0", include_state=True)
obs = env.reset()
assert "image" in obs and "state" in obs

Actions#

BSuite tasks expose discrete actions. BSuiteImageEnv maps them to a continuous one-hot Box of shape (n,) with values in [-1, 1]. The selected action is the index of the largest value. info["action"] stores the one-hot vector; info["executed_action"] stores the integer index.

Info contract#

Key

Always?

Description

discount

Yes

float32 scalar from the BSuite TimeStep

bsuite_id

Yes

Task identifier string (e.g. "catch/0")

action

Yes

One-hot action vector

executed_action

Yes

Integer index sent to the BSuite environment

vector_observation

Yes

Flattened raw observation array

terminated

Yes

bool — True when BSuite TimeStep.last()

truncated

Yes

Always False (BSuite does not expose time limits)

Troubleshooting#

  • ImportError: No module named bsuite: install the optional dependency with pip install bsuite.

  • "catch/0" not recognised: verify the task ID exists in your BSuite installation with list_available_bsuite_ids().

  • Synthetic images are uninformative: BSuite observations are small numeric vectors tiled onto a canvas; the synthetic image preserves all information but may look unfamiliar. Use include_state=True to access the raw vector alongside the image.