# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
#
from __future__ import annotations
import abc
import importlib
import warnings
from abc import abstractmethod
from enum import Enum
from pathlib import Path
from typing import Any, Callable, Dict, List, Optional, Type
from tensordict import TensorDictBase
from torch import Tensor
from torchrl.data import Composite
from torchrl.envs import EnvBase, RewardSum, Transform
from benchmarl.utils import _read_yaml_config, DEVICE_TYPING
def _type_check_task_config(
environemnt_name: str,
task_name: str,
config: Dict[str, Any],
warn_on_missing_dataclass: bool = True,
):
task_config_class = _get_task_config_class(environemnt_name, task_name)
if task_config_class is not None:
return task_config_class(**config).__dict__
else:
if warn_on_missing_dataclass:
warnings.warn(
"TaskConfig python dataclass not found, task is being loaded without type checks"
)
return config
def _get_task_config_class(environemnt_name: str, task_name: str):
try:
module = importlib.import_module(
f"{'.'.join(__name__.split('.')[:-1])}.{environemnt_name}.{task_name}"
)
return module.TaskConfig
except ModuleNotFoundError:
return None
[docs]
class TaskClass(abc.ABC):
"""
The class associated with an environment.
This class contains the logic on how to construct tasks for a specific environment.
The :meth:`TaskClass.get_env_fun` is the core method of this class, which will define the logic on how to construct
a :class:`torchrl.EnvBase` given :meth:`TaskClass.config` and :meth:`TaskClass.name`.
This methods, algonside all other abstract ones, is what users need to implement to introiduce new tasks.
Args:
name (str): The name of the task
config (Dict[str, Any]): The configuration of the task
"""
def __init__(self, name: str, config: Dict[str, Any]):
self.name = name
if config is None:
config = {}
self.config = config
[docs]
@abstractmethod
def get_env_fun(
self,
num_envs: int,
continuous_actions: bool,
seed: Optional[int],
device: DEVICE_TYPING,
) -> Callable[[], EnvBase]:
"""
This function is used to obtain a TorchRL object from the enum Task.
Args:
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 :class:`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 :class:`torchrl.envs.EnvBase` object
"""
raise NotImplementedError
[docs]
@abstractmethod
def supports_continuous_actions(self) -> bool:
"""
Return true if your task supports continuous actions.
If true, self.get_env_fun might be called with continuous_actions=True
"""
raise NotImplementedError
[docs]
@abstractmethod
def supports_discrete_actions(self) -> bool:
"""
Return true if your task supports discrete actions.
If true, self.get_env_fun might be called with continuous_actions=False
"""
raise NotImplementedError
[docs]
@abstractmethod
def max_steps(self, env: EnvBase) -> int:
"""
The maximum number of steps allowed in an evaluation rollout.
Args:
env (EnvBase): An environment created via self.get_env_fun
"""
raise NotImplementedError
[docs]
@abstractmethod
def has_render(self, env: EnvBase) -> bool:
"""
If env.render() should be called on the environment
Args:
env (EnvBase): An environment created via self.get_env_fun
"""
raise NotImplementedError
[docs]
@abstractmethod
def group_map(self, 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.
Args:
env (EnvBase): An environment created via self.get_env_fun
"""
raise NotImplementedError
[docs]
@abstractmethod
def observation_spec(self, env: EnvBase) -> Composite:
"""
A spec for the observation.
Must be a Composite with as many entries as needed nested under the ``group_name`` key.
Args:
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([]))
"""
raise NotImplementedError
[docs]
@abstractmethod
def info_spec(self, env: EnvBase) -> Optional[Composite]:
"""
A spec for the info.
If provided, must be a Composite with one (group_name, "info") entry per group (this entry can be composite).
Args:
env (EnvBase): An environment created via self.get_env_fun
"""
raise NotImplementedError
[docs]
@abstractmethod
def state_spec(self, env: EnvBase) -> Optional[Composite]:
"""
A spec for the state.
If provided, must be a Composite with one entry.
Args:
env (EnvBase): An environment created via self.get_env_fun
"""
raise NotImplementedError
[docs]
@abstractmethod
def action_spec(self, env: EnvBase) -> Composite:
"""
A spec for the action.
If provided, must be a Composite with one (group_name, "action") entry per group.
Args:
env (EnvBase): An environment created via self.get_env_fun
"""
raise NotImplementedError
[docs]
@abstractmethod
def action_mask_spec(self, env: EnvBase) -> Optional[Composite]:
"""
A spec for the action mask.
If provided, must be a Composite with one (group_name, "action_mask") entry per group.
Args:
env (EnvBase): An environment created via self.get_env_fun
"""
raise NotImplementedError
[docs]
@staticmethod
@abstractmethod
def env_name() -> str:
"""
The name of the environment in the ``benchmarl/conf/task`` folder
"""
raise NotImplementedError
[docs]
@staticmethod
def 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.
Args:
batch (TensorDictBase): the batch obtained from collection.
"""
return {}
[docs]
@staticmethod
def 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.
Args:
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 :class:`torch.Tensor` containing the pixels
"""
try:
return env.render(mode="rgb_array")
except TypeError:
return env.render()
def __repr__(self):
cls_name = self.__class__.__name__
return f"{cls_name}.{self.name}: (config={self.config})"
def __str__(self):
return self.__repr__()
def __eq__(self, other):
if not isinstance(other, self.__class__):
return False
return self.name == other.name and self.config == other.config
[docs]
class Task(Enum):
"""Enum of tasks in an environment.
Tasks are enums, one enum for each environment.
Each enviornment usually contains multiple tasks.
Tasks are used just to enumerate the available tasks, to convert a :class:`Task` into its corresponding instantiation,
you can call :meth:`Task.get_from_yaml` which will load the task config form yaml into the associated :class:`TaskClass`.
Each enum member can also be converted to a :class:`TaskClass` by calling :meth:`Task.get_task`, (which by default behaves like
:meth:`Task.get_from_yaml`) or by calling ``get_task(config={...})``, providing your own config.
Each new environment should inherit from :class:`Task` and instantiate its members as:
TASK_1 = None
TASK_2 = None
...
Tasks configs are loaded from ``benchmarl/conf/task``.
"""
[docs]
@staticmethod
def associated_class() -> Type[TaskClass]:
"""
The associated task class
"""
raise NotImplementedError
[docs]
@classmethod
def env_name(cls) -> str:
"""
The name of the environment in the ``benchmarl/conf/task`` folder
"""
return cls.associated_class().env_name()
[docs]
def get_task(self, config: Optional[Dict[str, Any]] = None) -> TaskClass:
"""
Get the :class:`TaskClass` object associated with this enum element by passing it the task name and config.
If no config is given, it will be loaded from ``benchmarl/conf/task/self.env_name()/self.name`` using :meth:`Task.get_from_yaml`.
Args:
config (dict): Optional configuration of the task.
If not provided, the default configuration will be loaded from yaml.
Returns:
The :class:`TaskClass` object for the task.
"""
if config is None:
return self.get_from_yaml()
return self.associated_class()(name=self.name, config=config)
def __new__(cls, *args, **kwargs):
value = len(cls.__members__) + 1
obj = object.__new__(cls)
obj._value_ = value
return obj
@staticmethod
def _load_from_yaml(name: str) -> Dict[str, Any]:
yaml_path = Path(__file__).parent.parent / "conf" / "task" / f"{name}.yaml"
return _read_yaml_config(str(yaml_path.resolve()))
[docs]
def get_from_yaml(self, path: Optional[str] = None) -> TaskClass:
"""
Load the task configuration from yaml
Args:
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 :class:`TaskClass` with the loaded config
"""
environment_name = self.env_name()
task_name = self.name.lower()
full_name = str(Path(environment_name) / Path(task_name))
if path is None:
config = Task._load_from_yaml(full_name)
else:
config = _read_yaml_config(path)
config = _type_check_task_config(environment_name, task_name, config)
return self.get_task(config=config)
######################
# Deprecated functions
######################
@property
def config(self):
raise ValueError(
"Task.config is deprecated, use Task.get_task().config instead"
)
def update_config(self, config: Dict[str, Any]) -> Task:
raise ValueError(
"Task.update_config is deprecated please use Task.get_task().config.update() instead"
)
def supports_continuous_actions(self) -> bool:
raise ValueError(
"Called function is deprecated is deprecated, please use Task.get_task().function() instead"
)
def supports_discrete_actions(self) -> bool:
raise ValueError(
"Called function is deprecated is deprecated, please use Task.get_task().function() instead"
)
def max_steps(self, env: EnvBase) -> int:
raise ValueError(
"Called function is deprecated is deprecated, please use Task.get_task().function() instead"
)
def has_render(self, env: EnvBase) -> bool:
raise ValueError(
"Called function is deprecated is deprecated, please use Task.get_task().function() instead"
)
def group_map(self, env: EnvBase) -> Dict[str, List[str]]:
raise ValueError(
"Called function is deprecated is deprecated, please use Task.get_task().function() instead"
)
def observation_spec(self, env: EnvBase) -> Composite:
raise ValueError(
"Called function is deprecated is deprecated, please use Task.get_task().function() instead"
)
def info_spec(self, env: EnvBase) -> Optional[Composite]:
raise ValueError(
"Called function is deprecated is deprecated, please use Task.get_task().function() instead"
)
def state_spec(self, env: EnvBase) -> Optional[Composite]:
raise ValueError(
"Called function is deprecated is deprecated, please use Task.get_task().function() instead"
)
def action_spec(self, env: EnvBase) -> Composite:
raise ValueError(
"Called function is deprecated is deprecated, please use Task.get_task().function() instead"
)
def action_mask_spec(self, env: EnvBase) -> Optional[Composite]:
raise ValueError(
"Called function is deprecated is deprecated, please use Task.get_task().function() instead"
)
@staticmethod
def log_info(batch: TensorDictBase) -> Dict[str, float]:
raise ValueError(
"Called function is deprecated is deprecated, please use Task.get_task().function() instead"
)
def get_reward_sum_transform(self, env: EnvBase) -> Transform:
raise ValueError(
"Called function is deprecated is deprecated, please use Task.get_task().function() instead"
)
def get_env_transforms(self, env: EnvBase) -> List[Transform]:
raise ValueError(
"Called function is deprecated is deprecated, please use Task.get_task().function() instead"
)
def get_replay_buffer_transforms(self, env: EnvBase, group: str) -> List[Transform]:
raise ValueError(
"Called function is deprecated is deprecated, please use Task.get_task().function() instead"
)
@staticmethod
def render_callback(experiment, env: EnvBase, data: TensorDictBase):
raise ValueError(
"Called function is deprecated is deprecated, please use Task.get_task().function() instead"
)