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 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 thetorchrl.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]