Dict

class Dict(*args, init_value=None, constrain_fn=None, sampling_order=None, **kwargs)[source]

Bases: Dict

Extended dictionary space with ordered sampling and nested support.

This class extends gymnasium.spaces.Dict to add state management, constraint validation, and explicit sampling order control. It composes multiple spaces into a hierarchical structure where dependencies between variables can be handled through ordered sampling.

Parameters:

  • *args – Positional arguments passed to gymnasium.spaces.Dict.

  • init_value (dict, optional) – Initial values for the space. If None, derived from init_value of contained spaces.

  • constrain_fn (callable, optional) – Function that returns True if the complete dictionary satisfies custom constraints.

  • sampling_order (list, optional) – Explicit order for sampling keys. If None, uses insertion order. Missing keys are appended.

  • **kwargs – Additional keyword arguments passed to Dict.

Variables:

  • init_value (dict) – Initial values for all contained spaces.

  • value (dict) – Current values of all contained spaces.

  • constrain_fn (callable) – Constraint validation function.

  • sampling_order (set) – Set of dotted paths for all variables in order.

Example

Create a nested space with sampling order dependencies:

from stable_worldmodel import spaces
import numpy as np

config = spaces.Dict(
    {
        "difficulty": spaces.Discrete(n=3, init_value=0),
        "world": spaces.Dict(
            {
                "width": spaces.Discrete(n=100, init_value=50),
                "height": spaces.Discrete(n=100, init_value=50),
            }
        ),
        "player_pos": spaces.Box(
            low=np.array([0, 0]),
            high=np.array([99, 99]),
            init_value=np.array([25, 25]),
        ),
    },
    sampling_order=["difficulty", "world", "player_pos"],
)

# Sample respects order
state = config.sample()

Note

Sampling order is crucial when variables have dependencies. For example, sample world size before sampling positions within it. Nested Dict spaces recursively apply their own sampling orders.

Accessing values in constraint functions: When implementing constrain_fn for Dict spaces, always use self.value['key']['key2'] instead of self['key']['key2'].value. The .value property recursively builds the complete value dictionary from the top level down, ensuring all nested values are up-to-date and correctly structured. Direct subspace access with .value only retrieves that specific subspace’s value without the full context.

Note that direct subspace access (e.g., self['key'].value) is perfectly fine for regular operations outside of constraint functions, such as reading individual subspace values or debugging. The recommendation to use top-level .value applies specifically to constraint functions where you need the complete, consistent state of all nested spaces.

Example of proper constraint function usage:

# Example: In a class with Dict space attribute
class Environment:
    def __init__(self):
        self.config_space = spaces.Dict({...})

    def validate_config(self):
        # ✓ CORRECT: Access via .value at top level
        values = self.config_space.value
        return values["player_pos"][0] < values["world"]["width"]

    def validate_wrong(self):
        # ✗ AVOID: Direct subspace access
        return (
            self.config_space["player_pos"].value[0]
            < self.config_space["world"]["width"].value
        )
__init__(*args, init_value=None, constrain_fn=None, sampling_order=None, **kwargs)[source]

Initialize a Dict space with state tracking and sampling order.

Parameters:

  • *args – Positional arguments passed to gymnasium.spaces.Dict.

  • init_value (dict, optional) – Initial values for the space. If None, derived from init_value of contained spaces. Defaults to None.

  • constrain_fn (callable, optional) – Function that takes a dict and returns True if the complete dictionary satisfies custom constraints. Defaults to None.

  • sampling_order (list, optional) – Explicit order for sampling keys. If None, uses insertion order. Missing keys are appended with warning. Defaults to None.

  • **kwargs – Keyword arguments passed to gymnasium.spaces.Dict.

Raises:

ValueError – If sampling_order contains keys not present in spaces.

check(debug=False)[source]

Validate all contained spaces’ current values.

Checks each contained space using its check() method if available, or falls back to contains(value). Optionally logs warnings for failed checks.

Parameters:

debug (bool, optional) – If True, logs warnings for spaces that fail validation. Defaults to False.

Returns:

True if all contained spaces have valid values, False otherwise.

Return type:

bool

contains(x) → bool[source]

Check if value is a valid member of this space.

Validates that x is a dictionary containing all required keys with values that satisfy each subspace’s constraints and the overall constraint function.

Parameters:

x – The value to check.

Returns:

True if x is a valid dict with all keys present, all values

within subspace bounds, and satisfies the constraint function. False otherwise.

Return type:

bool

property init_value

Initial values for all contained spaces.

Constructs initial value dictionary from contained spaces’ init_value properties. Falls back to sampling if a space lacks init_value.

Returns:

Dictionary mapping space keys to their initial values.

Return type:

dict

Type:

dict

names()[source]

Return all space keys including nested ones.

Returns:

A list of all keys in the Dict space, with nested keys using dot notation.

For example, a nested dict with key “a” containing subspace “b” would produce “a.b”.

Return type:

list

reset()[source]

Reset all contained spaces to their initial values.

Calls reset() on all contained spaces that have a reset method, then sets this space’s value to init_value.

sample(*args, max_tries=1000, warn_after_s=5.0, set_value=True, **kwargs)[source]

Sample a random element from the Dict space.

Samples each subspace in the sampling order and ensures the result satisfies any constraint functions. Uses rejection sampling if constraints are present.

Parameters:

  • *args – Positional arguments passed to each subspace’s sample method.

  • max_tries (int, optional) – Maximum number of rejection sampling attempts. Defaults to 1000.

  • warn_after_s (float, optional) – Issue a warning if sampling takes longer than this many seconds. Set to None to disable warnings. Defaults to 5.0.

  • set_value (bool, optional) – Whether to set the internal value to the sampled value. Defaults to True.

  • **kwargs – Additional keyword arguments passed to each subspace’s sample method.

Returns:

A dictionary with keys matching the space definition and values sampled

from their respective subspaces.

Return type:

dict

Raises:

RuntimeError – If a valid sample is not found within max_tries attempts.

property sampling_order

Set of dotted paths for all variables in sampling order.

Returns:

Set of strings representing dotted paths (e.g., ‘parent.child.key’)

for all variables including nested Dict spaces.

Return type:

set

Type:

set

update(keys)[source]

Update specific keys in the Dict space by resampling them.

Samples new values for the specified keys while maintaining the sampling order. Uses dot notation for nested keys (e.g., “a.b” for nested dict).

Parameters:

keys (container) – A container (list, set, etc.) of key names to resample. Keys should use dot notation for nested spaces.

Raises:

  • ValueError – If a specified key is not found in the Dict space.

  • AssertionError – If the updated values violate the space constraints.

property value

Current values of all contained spaces.

Constructs value dictionary from contained spaces’ value properties.

Returns:

Dictionary mapping space keys to their current values.

Return type:

dict

Raises:

ValueError – If a contained space does not have a value property.

Type:

dict