benchmarl.environments.TaskClass

class TaskClass(name: str, config: Dict[str, Any])

Bases: ABC

The class associated with an environment.

This class contains the logic on how to construct tasks for a specific environment.

The TaskClass.get_env_fun() is the core method of this class, which will define the logic on how to construct a torchrl.EnvBase given TaskClass.config() and TaskClass.name().

This methods, algonside all other abstract ones, is what users need to implement to introiduce new tasks.

Parameters:

• name (str) – The name of the task

• config (Dict[str, Any]) – The configuration of the task

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

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

abstract supports_continuous_actions() → bool

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

abstract supports_discrete_actions() → bool

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

abstract max_steps(env: EnvBase) → int

The maximum number of steps allowed in an evaluation rollout.

Parameters:

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

abstract has_render(env: EnvBase) → bool

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

Parameters:

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

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

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

abstract observation_spec(env: EnvBase) → Composite

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([]))

abstract info_spec(env: EnvBase) → Composite | None

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

abstract state_spec(env: EnvBase) → Composite | None

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

abstract action_spec(env: EnvBase) → Composite

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

abstract action_mask_spec(env: EnvBase) → Composite | None

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

abstract static env_name() → str

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

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

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

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]

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]

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) → Tensor

The render callback function for the enviornment.

This function is called at every step during evaluation to provide pixels for rendering.

Parameters:

• experiment (Experiment) – The Benchmarl experiment.

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

• data (TensorDictBase) – the current rollout data from the enviornment.

Returns:

The torch.Tensor containing the pixels