ObsExtraction

class fluidgym.wrappers.obs_extraction.ObsExtraction(env: FluidEnvLike, keys: list[str])

Bases: FluidWrapper

A wrapper that extracts specific observations from the observation dictionary.

It extracts only the observations specified in the keys list.

Parameters:

• env (FluidEnvLike) – The environment to wrap.

• keys (list[str] | None) – The list of keys to extract from the observation dictionary.

property observation_space: Dict

The observation space of the environment.

reset(seed: int | None = None, randomize: bool | None = None) → tuple[dict[str, Tensor], dict[str, Tensor]]

Resets the environment to an initial internal state, returning an initial observation and info.

Parameters:

• seed (int | None) – The seed to use for random number generation. If None, the current seed is used.

• randomize (bool | None) – Whether to randomize the initial state. If None, the default behavior is used.

Returns:

A tuple containing the initial observation and an info dictionary.

Return type:

tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]

step(action: Tensor) → tuple[dict[str, Tensor], Tensor, bool, bool, dict[str, Tensor]]

Run one timestep of the environment’s dynamics using the agent actions.

When the end of an episode is reached (terminated or truncated), it is necessary to call reset() to reset this environment’s state for the next episode.

Parameters:

action (torch.Tensor) – The action to take.

Returns:

• tuple[

• dict[str, torch.Tensor], torch.Tensor, bool, bool, dict[str, torch.Tensor]] – A tuple containing the observation, reward, terminated flag, truncated flag, and info dictionary.