benchmarl.environments.Task

class Task(value)[source]

Bases: Enum

Task.

Tasks are enums, one enum for each environment. Each enum member has a config attribute that is a dictionary which can be loaded from .yaml files. You can also access and modify this attribute directly.

Each new environment should inherit from Task and instantiate its members as

TASK_1 = None TASK_2 = None …

Tasks configs are loaded from benchmarl/conf/environments

update_config(config: Dict[str, Any]) Task[source]

Updates the task config

Parameters:

config (dictionary) – The config to update in the task

Returns: The updated task

get_env_fun(num_envs: int, continuous_actions: bool, seed: int | None, device: device | str | int) Callable[[], EnvBase][source]

This function is used to obtain a TorchRL object from the enum Task.

Parameters:
  • num_envs (int) – The number of envs that should be in the batch_size of the returned env. In vectorized envs, this can be used to set the number of batched environments. If your environment is not vectorized, you can just ignore this, and it will be wrapped in a torchrl.envs.SerialEnv with num_envs automatically.

  • continuous_actions (bool) – Whether your environment should have continuous or discrete actions. If your environment does not support both, ignore this and refer to the supports_x_actions methods.

  • seed (optional, int) – The seed of your env

  • device (str) – the device of your env, you can pass this to any torchrl env constructor

Returns: a function that takes no arguments and returns a torchrl.envs.EnvBase object

supports_continuous_actions() bool[source]

Return true if your task supports continuous actions. If true, self.get_env_fun might be called with continuous_actions=True

supports_discrete_actions() bool[source]

Return true if your task supports discrete actions. If true, self.get_env_fun might be called with continuous_actions=False

max_steps(env: EnvBase) int[source]

The maximum number of steps allowed in an evaluation rollout.

Parameters:

env (EnvBase) – An environment created via self.get_env_fun

has_render(env: EnvBase) bool[source]

If env.render() should be called on the environment

Parameters:

env (EnvBase) – An environment created via self.get_env_fun

group_map(env: EnvBase) Dict[str, List[str]][source]

The group_map mapping agents groups to agent names. This should be reelected in the TensorDicts coming from the environment where agent data is supposed to be stacked according to this.

Parameters:

env (EnvBase) – An environment created via self.get_env_fun

observation_spec(env: EnvBase) Composite[source]

A spec for the observation. Must be a Composite with as many entries as needed nested under the group_name key.

Parameters:

env (EnvBase) – An environment created via self.get_env_fun

Examples

>>> print(task.observation_spec(env))
Composite(
    agents: Composite(
        observation: Composite(
            image: UnboundedDiscreteTensorSpec(
                shape=torch.Size([8, 88, 88, 3]),
                space=ContinuousBox(
                    low=Tensor(shape=torch.Size([8, 88, 88, 3]), device=cpu, dtype=torch.int64, contiguous=True),
                    high=Tensor(shape=torch.Size([8, 88, 88, 3]), device=cpu, dtype=torch.int64, contiguous=True)),
                device=cpu,
                dtype=torch.uint8,
                domain=discrete),
            array: Unbounded(
                shape=torch.Size([8, 3]),
                space=None,
                device=cpu,
                dtype=torch.float32,
                domain=continuous), device=cpu, shape=torch.Size([8])), device=cpu, shape=torch.Size([8])), device=cpu, shape=torch.Size([]))
info_spec(env: EnvBase) Composite | None[source]

A spec for the info. If provided, must be a Composite with one (group_name, “info”) entry per group (this entry can be composite).

Parameters:

env (EnvBase) – An environment created via self.get_env_fun

state_spec(env: EnvBase) Composite | None[source]

A spec for the state. If provided, must be a Composite with one entry.

Parameters:

env (EnvBase) – An environment created via self.get_env_fun

action_spec(env: EnvBase) Composite[source]

A spec for the action. If provided, must be a Composite with one (group_name, “action”) entry per group.

Parameters:

env (EnvBase) – An environment created via self.get_env_fun

action_mask_spec(env: EnvBase) Composite | None[source]

A spec for the action mask. If provided, must be a Composite with one (group_name, “action_mask”) entry per group.

Parameters:

env (EnvBase) – An environment created via self.get_env_fun

static env_name() str[source]

The name of the environment in the benchmarl/conf/task folder

static log_info(batch: TensorDictBase) Dict[str, float][source]

Return a str->float dict with extra items to log. This function has access to the collected batch and is optional.

Parameters:

batch (TensorDictBase) – the batch obtained from collection.

get_reward_sum_transform(env: EnvBase) Transform[source]

Returns the RewardSum transform for the environment

Parameters:

env (EnvBase) – An environment created via self.get_env_fun

get_env_transforms(env: EnvBase) List[Transform][source]

Returns a list of torchrl.envs.Transform to be applied to the env.

Parameters:

env (EnvBase) – An environment created via self.get_env_fun

get_replay_buffer_transforms(env: EnvBase, group: str) List[Transform][source]

Returns a list of torchrl.envs.Transform to be applied to the torchrl.data.ReplayBuffer of the specified group.

Parameters:
  • env (EnvBase) – An environment created via self.get_env_fun

  • group (str) – The agent group using the replay buffer

static render_callback(experiment, env: EnvBase, data: TensorDictBase)[source]
get_from_yaml(path: str | None = None) Task[source]

Load the task configuration from yaml

Parameters:

path (str, optional) – The full path of the yaml file to load from. If None, it will default to benchmarl/conf/task/self.env_name()/self.name

Returns: the task with the loaded config