opentau.envs.utils

This module contains utility functions for environments.

Functions

add_envs_task(env, observation)	Adds task feature to the observation dict with respect to the first environment attribute.
add_eval_metadata(observation, cfg)	Inject pi07 inference-time metadata from cfg.env.metadata.
add_subgoal_images(observation, generator)	Inject subgoal{k} / subgoal_is_pad keys from a generator.
are_all_envs_same_type(env)	Checks if all environments in a vectorized environment are of the same type.
check_env_attributes_and_types(env)	Checks if all environments in a vectorized environment have 'task_description' or 'task' attributes.
close_envs(-> None  -> None  -> None)	Close a single environment, a list of environments, or a dictionary of environments.
preprocess_observation(np_observations, cfg)	Convert environment observation to OpenTau format observation.

opentau.envs.utils.add_envs_task(env: VectorEnv, observation: dict[str, Any]) → dict[str, Any]

Adds task feature to the observation dict with respect to the first environment attribute.

Parameters:

• env – A vectorized Gym environment (SyncVectorEnv or AsyncVectorEnv).

• observation – A dictionary of observations from the vectorized environment, which will be modified in place.

Returns:

The updated observation dictionary with the ‘prompt’ key added.

opentau.envs.utils.add_eval_metadata(observation: dict[str, Any], cfg: TrainPipelineConfig) → dict[str, Any]

Inject pi07 inference-time metadata from cfg.env.metadata.

Numeric fields (speed / quality / mistake) are broadcast as (num_envs,) tensors with a parallel {key}_is_pad=False flag, using the same dtypes the dataset emits at training time (torch.long for speed / quality, torch.bool for mistake). String fields (robot_type / control_mode) are broadcast as list[str] of length num_envs. Fields set to None on the config are skipped, so the corresponding batch key stays absent and the policy’s prepare_metadata pad default kicks in.

fps is the env’s stepping frequency (EnvConfig.fps, e.g. 20 for LIBERO) and is broadcast as a torch.long tensor with fps_is_pad=False whenever meta.emit_fps is True (the default). Set emit_fps=False on EnvMetadataConfig to skip it — useful when resuming a checkpoint trained without fps conditioning.

cfg.env is guaranteed non-None here — eval_policy dereferences cfg.env.type before the rollout loop starts.

Parameters:

• observation – Observation dict produced by preprocess_observation (must contain state to source the batch size and device).

• cfg – Training/eval pipeline config; reads cfg.env.metadata and cfg.env.fps.

Returns:

The observation dict, mutated in place.

opentau.envs.utils.add_subgoal_images(observation: dict[str, Any], generator: SubgoalImageGenerator | None) → dict[str, Any]

Inject subgoal{k} / subgoal_is_pad keys from a generator.

Mirrors add_envs_task() / add_eval_metadata(): the generator inspects observation (which already carries camera{k}, state, img_is_pad, prompt, and any pi07 metadata keys) and returns the subgoal batch keys the policy’s prepare_subgoal_images consumes.

The caller is expected to have invoked generator.start_episode(prompts) once at the top of the rollout — the per-step __call__ only reads the cached choice. When generator is None, this is a no-op and the policy’s missing-key fallback (zero subgoals, mask=False) takes over.

Parameters:

• observation – Observation dict mutated in place via .update(...).

• generator – Subgoal image generator, or None to skip injection.

Returns:

The observation dict, mutated in place.

opentau.envs.utils.are_all_envs_same_type(env: VectorEnv) → bool

Checks if all environments in a vectorized environment are of the same type.

Parameters:

env – A vectorized Gym environment (SyncVectorEnv or AsyncVectorEnv).

Returns:

True if all environments are of the same type, False otherwise.

opentau.envs.utils.check_env_attributes_and_types(env: VectorEnv) → None

Checks if all environments in a vectorized environment have ‘task_description’ or ‘task’ attributes. A warning will be raised if any environment is missing these attributes.

Parameters:

env – A vectorized Gym environment (SyncVectorEnv or AsyncVectorEnv).

Raises:

ValueError – If the environment is not a SyncVectorEnv or AsyncVectorEnv.

opentau.envs.utils.close_envs(obj: Any) → None

opentau.envs.utils.close_envs(env: Mapping) → None

opentau.envs.utils.close_envs(envs: Sequence) → None

opentau.envs.utils.close_envs(env: Env) → None

opentau.envs.utils.close_envs(env: None) → None

Close a single environment, a list of environments, or a dictionary of environments.

opentau.envs.utils.preprocess_observation(np_observations: dict, cfg: TrainPipelineConfig) → dict[str, Tensor]

Convert environment observation to OpenTau format observation. :param np_observations: Dictionary of observation batches from a Gym vector environment. :param cfg: Training configuration that contains max_state_dim, num_cams, resolution, etc.

Returns:

Dictionary of observation batches with keys renamed to OpenTau format and values as tensors.