torchrl 0.11.0__cp314-cp314t-macosx_11_0_arm64.whl

torchrl/envs/custom/tictactoeenv.py

@@ -0,0 +1,288 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+from __future__ import annotations
+
+import torch
+from tensordict import TensorDict, TensorDictBase
+from torchrl.data.tensor_specs import Categorical, Composite, Unbounded
+from torchrl.envs.common import EnvBase
+
+
+class TicTacToeEnv(EnvBase):
+    """A Tic-Tac-Toe implementation.
+
+    Keyword Args:
+        single_player (bool, optional): whether one or two players have to be
+            accounted for. ``single_player=True`` means that ``"player1"`` is
+            playing randomly. If ``False`` (default), at each turn,
+            one of the two players has to play.
+        device (torch.device, optional): the device where to put the tensors.
+            Defaults to ``None`` (default device).
+
+    The environment is stateless. To run it across multiple batches, call
+
+        >>> env.reset(TensorDict(batch_size=desired_batch_size))
+
+    If the ``"mask"`` entry is present, ``rand_action`` takes it into account to
+    generate the next action. Any policy executed on this env should take this
+    mask into account, as well as the turn of the player (stored in the ``"turn"``
+    output entry).
+
+    Specs:
+        >>> print(env.specs)
+        Composite(
+            output_spec: Composite(
+                full_observation_spec: Composite(
+                    board: Categorical(
+                        shape=torch.Size([3, 3]),
+                        space=DiscreteBox(n=2),
+                        dtype=torch.int32,
+                        domain=discrete),
+                    turn: Categorical(
+                        shape=torch.Size([1]),
+                        space=DiscreteBox(n=2),
+                        dtype=torch.int32,
+                        domain=discrete),
+                    mask: Categorical(
+                        shape=torch.Size([9]),
+                        space=DiscreteBox(n=2),
+                        dtype=torch.bool,
+                        domain=discrete),
+                    shape=torch.Size([])),
+                full_reward_spec: Composite(
+                    player0: Composite(
+                        reward: UnboundedContinuous(
+                            shape=torch.Size([1]),
+                            space=ContinuousBox(
+                                low=Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.float32, contiguous=True),
+                                high=Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.float32, contiguous=True)),
+                            dtype=torch.float32,
+                            domain=continuous),
+                        shape=torch.Size([])),
+                    player1: Composite(
+                        reward: UnboundedContinuous(
+                            shape=torch.Size([1]),
+                            space=ContinuousBox(
+                                low=Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.float32, contiguous=True),
+                                high=Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.float32, contiguous=True)),
+                            dtype=torch.float32,
+                            domain=continuous),
+                        shape=torch.Size([])),
+                    shape=torch.Size([])),
+                full_done_spec: Composite(
+                    done: Categorical(
+                        shape=torch.Size([1]),
+                        space=DiscreteBox(n=2),
+                        dtype=torch.bool,
+                        domain=discrete),
+                    terminated: Categorical(
+                        shape=torch.Size([1]),
+                        space=DiscreteBox(n=2),
+                        dtype=torch.bool,
+                        domain=discrete),
+                    truncated: Categorical(
+                        shape=torch.Size([1]),
+                        space=DiscreteBox(n=2),
+                        dtype=torch.bool,
+                        domain=discrete),
+                    shape=torch.Size([])),
+                shape=torch.Size([])),
+            input_spec: Composite(
+                full_state_spec: Composite(
+                    board: Categorical(
+                        shape=torch.Size([3, 3]),
+                        space=DiscreteBox(n=2),
+                        dtype=torch.int32,
+                        domain=discrete),
+                    turn: Categorical(
+                        shape=torch.Size([1]),
+                        space=DiscreteBox(n=2),
+                        dtype=torch.int32,
+                        domain=discrete),
+                    mask: Categorical(
+                        shape=torch.Size([9]),
+                        space=DiscreteBox(n=2),
+                        dtype=torch.bool,
+                        domain=discrete), shape=torch.Size([])),
+                full_action_spec: Composite(
+                    action: Categorical(
+                        shape=torch.Size([1]),
+                        space=DiscreteBox(n=9),
+                        dtype=torch.int64,
+                        domain=discrete),
+                    shape=torch.Size([])),
+                shape=torch.Size([])),
+            shape=torch.Size([]))
+
+    To run a dummy rollout, execute the following command:
+
+    Examples:
+        >>> env = TicTacToeEnv()
+        >>> env.rollout(10)
+        TensorDict(
+            fields={
+                action: Tensor(shape=torch.Size([9, 1]), device=cpu, dtype=torch.int64, is_shared=False),
+                board: Tensor(shape=torch.Size([9, 3, 3]), device=cpu, dtype=torch.int32, is_shared=False),
+                done: Tensor(shape=torch.Size([9, 1]), device=cpu, dtype=torch.bool, is_shared=False),
+                mask: Tensor(shape=torch.Size([9, 9]), device=cpu, dtype=torch.bool, is_shared=False),
+                next: TensorDict(
+                    fields={
+                        board: Tensor(shape=torch.Size([9, 3, 3]), device=cpu, dtype=torch.int32, is_shared=False),
+                        done: Tensor(shape=torch.Size([9, 1]), device=cpu, dtype=torch.bool, is_shared=False),
+                        mask: Tensor(shape=torch.Size([9, 9]), device=cpu, dtype=torch.bool, is_shared=False),
+                        player0: TensorDict(
+                            fields={
+                                reward: Tensor(shape=torch.Size([9, 1]), device=cpu, dtype=torch.float32, is_shared=False)},
+                            batch_size=torch.Size([9]),
+                            device=None,
+                            is_shared=False),
+                        player1: TensorDict(
+                            fields={
+                                reward: Tensor(shape=torch.Size([9, 1]), device=cpu, dtype=torch.float32, is_shared=False)},
+                            batch_size=torch.Size([9]),
+                            device=None,
+                            is_shared=False),
+                        terminated: Tensor(shape=torch.Size([9, 1]), device=cpu, dtype=torch.bool, is_shared=False),
+                        truncated: Tensor(shape=torch.Size([9, 1]), device=cpu, dtype=torch.bool, is_shared=False),
+                        turn: Tensor(shape=torch.Size([9, 1]), device=cpu, dtype=torch.int32, is_shared=False)},
+                    batch_size=torch.Size([9]),
+                    device=None,
+                    is_shared=False),
+                terminated: Tensor(shape=torch.Size([9, 1]), device=cpu, dtype=torch.bool, is_shared=False),
+                truncated: Tensor(shape=torch.Size([9, 1]), device=cpu, dtype=torch.bool, is_shared=False),
+                turn: Tensor(shape=torch.Size([9, 1]), device=cpu, dtype=torch.int32, is_shared=False)},
+            batch_size=torch.Size([9]),
+            device=None,
+            is_shared=False)
+
+    """
+
+    # batch_locked is set to False since various batch sizes can be provided to the env
+    batch_locked: bool = False
+
+    def __init__(self, *, single_player: bool = False, device=None):
+        super().__init__(device=device)
+        self.single_player = single_player
+        self.action_spec: Unbounded = Categorical(
+            n=9,
+            shape=(),
+            device=device,
+        )
+
+        self.full_observation_spec: Composite = Composite(
+            board=Unbounded(shape=(3, 3), dtype=torch.int, device=device),
+            turn=Categorical(
+                2,
+                shape=(1,),
+                dtype=torch.int,
+                device=device,
+            ),
+            mask=Categorical(
+                2,
+                shape=(9,),
+                dtype=torch.bool,
+                device=device,
+            ),
+            device=device,
+        )
+        self.state_spec: Composite = self.observation_spec.clone()
+
+        self.reward_spec: Unbounded = Composite(
+            {
+                ("player0", "reward"): Unbounded(shape=(1,), device=device),
+                ("player1", "reward"): Unbounded(shape=(1,), device=device),
+            },
+            device=device,
+        )
+
+        self.full_done_spec: Categorical = Composite(
+            done=Categorical(2, shape=(1,), dtype=torch.bool, device=device),
+            device=device,
+        )
+        self.full_done_spec["terminated"] = self.full_done_spec["done"].clone()
+        self.full_done_spec["truncated"] = self.full_done_spec["done"].clone()
+
+    def _reset(self, reset_td: TensorDict) -> TensorDict:
+        shape = reset_td.shape if reset_td is not None else ()
+        state = self.state_spec.zero(shape)
+        state["board"] -= 1
+        state["mask"].fill_(True)
+        return state.update(self.full_done_spec.zero(shape))
+
+    def _step(self, state: TensorDict) -> TensorDict:
+        board = state["board"].clone()
+        turn = state["turn"].clone()
+        action = state["action"]
+        board.flatten(-2, -1).scatter_(index=action.unsqueeze(-1), dim=-1, value=1)
+        wins = self.win(board, action)
+
+        mask = board.flatten(-2, -1) == -1
+        done = wins | ~mask.any(-1, keepdim=True)
+        terminated = done.clone()
+
+        reward_0 = wins & (turn == 0)
+        reward_1 = wins & (turn == 1)
+
+        state = TensorDict(
+            {
+                "done": done,
+                "terminated": terminated,
+                ("player0", "reward"): reward_0.float(),
+                ("player1", "reward"): reward_1.float(),
+                "board": torch.where(board == -1, board, 1 - board),
+                "turn": 1 - turn,
+                "mask": mask,
+            },
+            batch_size=state.batch_size,
+        )
+        if self.single_player:
+            select = (~done & (turn == 0)).squeeze(-1)
+            if select.all():
+                state_select = state
+            elif select.any():
+                state_select = state[select]
+            else:
+                return state
+            state_select = self._step(self.rand_action(state_select))
+            if select.all():
+                return state_select
+            return torch.where(done, state, state_select)
+        return state
+
+    def _set_seed(self, seed: int | None) -> None:
+        ...
+
+    @staticmethod
+    def win(board: torch.Tensor, action: torch.Tensor):
+        row = action // 3  # type: ignore
+        col = action % 3  # type: ignore
+        if board[..., row, :].sum() == 3:
+            return True
+        if board[..., col].sum() == 3:
+            return True
+        if board.diagonal(0, -2, -1).sum() == 3:
+            return True
+        if board.flip(-1).diagonal(0, -2, -1).sum() == 3:
+            return True
+        return False
+
+    @staticmethod
+    def full(board: torch.Tensor) -> bool:
+        return torch.sym_int(board.abs().sum()) == 9
+
+    @staticmethod
+    def get_action_mask():
+        pass
+
+    def rand_action(self, tensordict: TensorDictBase | None = None):
+        mask = tensordict.get("mask")
+        action_spec = self.action_spec
+        if tensordict.ndim:
+            action_spec = action_spec.expand(tensordict.shape)
+        else:
+            action_spec = action_spec.clone()
+        action_spec.update_mask(mask)
+        tensordict.set(self.action_key, action_spec.rand())
+        return tensordict

torchrl/envs/env_creator.py

@@ -0,0 +1,263 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+from __future__ import annotations
+
+from collections import OrderedDict
+from collections.abc import Callable
+from multiprocessing.sharedctypes import Synchronized
+from multiprocessing.synchronize import Lock, RLock
+
+import torch
+from tensordict import TensorDictBase
+from torchrl._utils import logger as torchrl_logger
+from torchrl.data.utils import CloudpickleWrapper
+from torchrl.envs.common import EnvBase, EnvMetaData
+
+
+class EnvCreator:
+    """Environment creator class.
+
+    EnvCreator is a generic environment creator class that can substitute
+    lambda functions when creating environments in multiprocessing contexts.
+    If the environment created on a subprocess must share information with the
+    main process (e.g. for the VecNorm transform), EnvCreator will pass the
+    pointers to the tensordicts in shared memory to each process such that
+    all of them are synchronised.
+
+    Args:
+        create_env_fn (callable): a callable that returns an EnvBase
+            instance.
+        create_env_kwargs (dict, optional): the kwargs of the env creator.
+        share_memory (bool, optional): if False, the resulting tensordict
+            from the environment won't be placed in shared memory.
+        **kwargs: additional keyword arguments to be passed to the environment
+            during construction.
+
+    Examples:
+        >>> # We create the same environment on 2 processes using VecNorm
+        >>> # and check that the discounted count of observations matches on
+        >>> # both workers, even if one has not executed any step
+        >>> import time
+        >>> from torchrl.envs.libs.gym import GymEnv
+        >>> from torchrl.envs.transforms import VecNorm, TransformedEnv
+        >>> from torchrl.envs import EnvCreator
+        >>> from torch import multiprocessing as mp
+        >>> env_fn = lambda: TransformedEnv(GymEnv("Pendulum-v1"), VecNorm())
+        >>> env_creator = EnvCreator(env_fn)
+        >>>
+        >>> def test_env1(env_creator):
+        ...     env = env_creator()
+        ...     tensordict = env.reset()
+        ...     for _ in range(10):
+        ...         env.rand_step(tensordict)
+        ...         if tensordict.get(("next", "done")):
+        ...             tensordict = env.reset(tensordict)
+        ...     print("env 1: ", env.transform._td.get(("next", "observation_count")))
+        >>>
+        >>> def test_env2(env_creator):
+        ...     env = env_creator()
+        ...     time.sleep(5)
+        ...     print("env 2: ", env.transform._td.get(("next", "observation_count")))
+        >>>
+        >>> if __name__ == "__main__":
+        ...     ps = []
+        ...     p1 = mp.Process(target=test_env1, args=(env_creator,))
+        ...     p1.start()
+        ...     ps.append(p1)
+        ...     p2 = mp.Process(target=test_env2, args=(env_creator,))
+        ...     p2.start()
+        ...     ps.append(p1)
+        ...     for p in ps:
+        ...         p.join()
+        env 1:  tensor([11.9934])
+        env 2:  tensor([11.9934])
+    """
+
+    def __init__(
+        self,
+        create_env_fn: Callable[..., EnvBase],
+        create_env_kwargs: dict | None = None,
+        share_memory: bool = True,
+        **kwargs,
+    ) -> None:
+        if not isinstance(create_env_fn, (EnvCreator, CloudpickleWrapper)):
+            self.create_env_fn = CloudpickleWrapper(create_env_fn)
+        else:
+            self.create_env_fn = create_env_fn
+
+        self.create_env_kwargs = kwargs
+        if isinstance(create_env_kwargs, dict):
+            self.create_env_kwargs.update(create_env_kwargs)
+        self.initialized = False
+        self._meta_data = None
+        self._share_memory = share_memory
+        self.init_()
+
+    def make_variant(self, **kwargs) -> EnvCreator:
+        """Creates a variant of the EnvCreator, pointing to the same underlying metadata but with different keyword arguments during construction.
+
+        This can be useful with transforms that share a state, like :class:`~torchrl.envs.TrajCounter`.
+
+        Examples:
+            >>> from torchrl.envs import GymEnv
+            >>> env_creator_pendulum = EnvCreator(GymEnv, env_name="Pendulum-v1")
+            >>> env_creator_cartpole = env_creator_pendulum.make_variant(env_name="CartPole-v1")
+
+        """
+        # Copy self
+        out = type(self).__new__(type(self))
+        out.__dict__.update(self.__dict__)
+        out.create_env_kwargs.update(kwargs)
+        return out
+
+    def share_memory(self, state_dict: OrderedDict) -> None:
+        for key, item in list(state_dict.items()):
+            if isinstance(item, (TensorDictBase,)):
+                if not item.is_shared():
+                    item.share_memory_()
+                else:
+                    torchrl_logger.info(
+                        f"{self.env_type}: {item} is already shared"
+                    )  # , deleting key'val)
+                    del state_dict[key]
+            elif isinstance(item, OrderedDict):
+                self.share_memory(item)
+            elif isinstance(item, torch.Tensor):
+                del state_dict[key]
+
+    @property
+    def meta_data(self) -> EnvMetaData:
+        if self._meta_data is None:
+            raise RuntimeError(
+                "meta_data is None in EnvCreator. " "Make sure init_() has been called."
+            )
+        return self._meta_data
+
+    @meta_data.setter
+    def meta_data(self, value: EnvMetaData):
+        self._meta_data = value
+
+    @staticmethod
+    def _is_mp_value(val):
+        if isinstance(val, (Synchronized,)) and hasattr(val, "_obj"):
+            return True
+        # Also check for lock types which need to be shared across processes
+        if isinstance(val, (Lock, RLock)):
+            return True
+        return False
+
+    @classmethod
+    def _find_mp_values(cls, env_or_transform, values, prefix=()):
+        from torchrl.envs.transforms.transforms import Compose, TransformedEnv
+
+        if isinstance(env_or_transform, EnvBase) and isinstance(
+            env_or_transform, TransformedEnv
+        ):
+            cls._find_mp_values(
+                env_or_transform.transform,
+                values=values,
+                prefix=prefix + ("transform",),
+            )
+            cls._find_mp_values(
+                env_or_transform.base_env, values=values, prefix=prefix + ("base_env",)
+            )
+        elif isinstance(env_or_transform, Compose):
+            for i, t in enumerate(env_or_transform.transforms):
+                cls._find_mp_values(t, values=values, prefix=prefix + (i,))
+        for k, v in env_or_transform.__dict__.items():
+            if cls._is_mp_value(v):
+                values.append((prefix + (k,), v))
+        return values
+
+    def init_(self) -> EnvCreator:
+        shadow_env = self.create_env_fn(**self.create_env_kwargs)
+        tensordict = shadow_env.reset()
+        shadow_env.rand_step(tensordict)
+        self.env_type = type(shadow_env)
+        self._transform_state_dict = shadow_env.state_dict()
+        # Extract any mp.Value object from the env
+        self._mp_values = self._find_mp_values(shadow_env, values=[])
+
+        if self._share_memory:
+            self.share_memory(self._transform_state_dict)
+        self.initialized = True
+        self.meta_data = EnvMetaData.metadata_from_env(shadow_env)
+        shadow_env.close()
+        del shadow_env
+        return self
+
+    @classmethod
+    def _set_mp_value(cls, env, key, value):
+        if len(key) > 1:
+            if isinstance(key[0], int):
+                return cls._set_mp_value(env[key[0]], key[1:], value)
+            else:
+                return cls._set_mp_value(getattr(env, key[0]), key[1:], value)
+        else:
+            setattr(env, key[0], value)
+
+    def __call__(self, **kwargs) -> EnvBase:
+        if not self.initialized:
+            raise RuntimeError("EnvCreator must be initialized before being called.")
+        kwargs.update(self.create_env_kwargs)  # create_env_kwargs precedes
+        env = self.create_env_fn(**kwargs)
+        if self._mp_values:
+            for k, v in self._mp_values:
+                self._set_mp_value(env, k, v)
+        env.load_state_dict(self._transform_state_dict, strict=False)
+        return env
+
+    def state_dict(self) -> OrderedDict:
+        if self._transform_state_dict is None:
+            return OrderedDict()
+        return self._transform_state_dict
+
+    def load_state_dict(self, state_dict: OrderedDict) -> None:
+        if self._transform_state_dict is not None:
+            for key, item in state_dict.items():
+                item_to_update = self._transform_state_dict[key]
+                item_to_update.copy_(item)
+
+    def __repr__(self) -> str:
+        substr = ", ".join(
+            [f"{key}: {type(item)}" for key, item in self.create_env_kwargs]
+        )
+        return f"EnvCreator({self.create_env_fn}({substr}))"
+
+
+def env_creator(fun: Callable) -> EnvCreator:
+    """Helper function to call `EnvCreator`."""
+    return EnvCreator(fun)
+
+
+def get_env_metadata(env_or_creator: EnvBase | Callable, kwargs: dict | None = None):
+    """Retrieves a EnvMetaData object from an env."""
+    if isinstance(env_or_creator, (EnvBase,)):
+        return EnvMetaData.metadata_from_env(env_or_creator)
+    elif not isinstance(env_or_creator, EnvBase) and not isinstance(
+        env_or_creator, EnvCreator
+    ):
+        # then env is a creator
+        if kwargs is None:
+            kwargs = {}
+        env = env_or_creator(**kwargs)
+        return EnvMetaData.metadata_from_env(env)
+    elif isinstance(env_or_creator, EnvCreator):
+        if not (
+            kwargs == env_or_creator.create_env_kwargs
+            or kwargs is None
+            or len(kwargs) == 0
+        ):
+            raise RuntimeError(
+                "kwargs mismatch between EnvCreator and the kwargs provided to get_env_metadata:"
+                f"got EnvCreator.create_env_kwargs={env_or_creator.create_env_kwargs} and "
+                f"kwargs = {kwargs}"
+            )
+        return env_or_creator.meta_data.clone()
+    else:
+        raise NotImplementedError(
+            f"env of type {type(env_or_creator)} is not supported by get_env_metadata."
+        )