Tianshou/test/base/env.py

import random
import time
from copy import deepcopy
from typing import Any, Literal

import gymnasium as gym
import networkx as nx
import numpy as np
from gymnasium.spaces import Box, Dict, Discrete, MultiDiscrete, Space, Tuple


class MoveToRightEnv(gym.Env):
    """A task for "going right". The task is to go right ``size`` steps.

    The observation is the current index, and the action is to go left or right.
    Action 0 is to go left, and action 1 is to go right.
    Taking action 0 at index 0 will keep the index at 0.
    Arriving at index ``size`` means the task is done.
    In the current implementation, stepping after the task is done is possible, which will
    lead the index to be larger than ``size``.

    Index 0 is the starting point. If reset is called with default options, the index will
    be reset to 0.
    """

    def __init__(
        self,
        size: int,
        sleep: float = 0.0,
        dict_state: bool = False,
        recurse_state: bool = False,
        ma_rew: int = 0,
        multidiscrete_action: bool = False,
        random_sleep: bool = False,
        array_state: bool = False,
    ) -> None:
        assert (
            dict_state + recurse_state + array_state <= 1
        ), "dict_state / recurse_state / array_state can be only one true"
        self.size = size
        self.sleep = sleep
        self.random_sleep = random_sleep
        self.dict_state = dict_state
        self.recurse_state = recurse_state
        self.array_state = array_state
        self.ma_rew = ma_rew
        self._md_action = multidiscrete_action
        # how many steps this env has stepped
        self.steps = 0
        if dict_state:
            self.observation_space = Dict(
                {
                    "index": Box(shape=(1,), low=0, high=size - 1),
                    "rand": Box(shape=(1,), low=0, high=1, dtype=np.float64),
                },
            )
        elif recurse_state:
            self.observation_space = Dict(
                {
                    "index": Box(shape=(1,), low=0, high=size - 1),
                    "dict": Dict(
                        {
                            "tuple": Tuple(
                                (
                                    Discrete(2),
                                    Box(shape=(2,), low=0, high=1, dtype=np.float64),
                                ),
                            ),
                            "rand": Box(shape=(1, 2), low=0, high=1, dtype=np.float64),
                        },
                    ),
                },
            )
        elif array_state:
            self.observation_space = Box(shape=(4, 84, 84), low=0, high=255)
        else:
            self.observation_space = Box(shape=(1,), low=0, high=size - 1)
        if multidiscrete_action:
            self.action_space = MultiDiscrete([2, 2])
        else:
            self.action_space = Discrete(2)
        self.terminated = False
        self.index = 0

    def reset(
        self,
        seed: int | None = None,
        # TODO: passing a dict here doesn't make any sense
        options: dict[str, Any] | None = None,
    ) -> tuple[dict[str, Any] | np.ndarray, dict]:
        """:param seed:
        :param options: the start index is provided in options["state"]
        :return:
        """
        if options is None:
            options = {"state": 0}
        super().reset(seed=seed)
        self.terminated = False
        self.do_sleep()
        self.index = options["state"]
        return self._get_state(), {"key": 1, "env": self}

    def _get_reward(self) -> list[int] | int:
        """Generate a non-scalar reward if ma_rew is True."""
        end_flag = int(self.terminated)
        if self.ma_rew > 0:
            return [end_flag] * self.ma_rew
        return end_flag

    def _get_state(self) -> dict[str, Any] | np.ndarray:
        """Generate state(observation) of MyTestEnv."""
        if self.dict_state:
            return {
                "index": np.array([self.index], dtype=np.float32),
                "rand": self.np_random.random(1),
            }
        if self.recurse_state:
            return {
                "index": np.array([self.index], dtype=np.float32),
                "dict": {
                    "tuple": (np.array([1], dtype=int), self.np_random.random(2)),
                    "rand": self.np_random.random((1, 2)),
                },
            }
        if self.array_state:
            img = np.zeros([4, 84, 84], int)
            img[3, np.arange(84), np.arange(84)] = self.index
            img[2, np.arange(84)] = self.index
            img[1, :, np.arange(84)] = self.index
            img[0] = self.index
            return img
        return np.array([self.index], dtype=np.float32)

    def do_sleep(self) -> None:
        if self.sleep > 0:
            sleep_time = random.random() if self.random_sleep else 1
            sleep_time *= self.sleep
            time.sleep(sleep_time)

    def step(self, action: np.ndarray | int):  # type: ignore[no-untyped-def]  # cf. issue #1080
        self.steps += 1
        if self._md_action and isinstance(action, np.ndarray):
            action = action[0]
        if self.terminated:
            raise ValueError("step after done !!!")
        self.do_sleep()
        if self.index == self.size:
            self.terminated = True
            return self._get_state(), self._get_reward(), self.terminated, False, {}
        if action == 0:
            self.index = max(self.index - 1, 0)
            return (
                self._get_state(),
                self._get_reward(),
                self.terminated,
                False,
                {"key": 1, "env": self} if self.dict_state else {},
            )
        if action == 1:
            self.index += 1
            self.terminated = self.index == self.size
            return (
                self._get_state(),
                self._get_reward(),
                self.terminated,
                False,
                {"key": 1, "env": self},
            )
        return None


class NXEnv(gym.Env):
    def __init__(self, size: int, obs_type: str, feat_dim: int = 32) -> None:
        self.size = size
        self.feat_dim = feat_dim
        self.graph = nx.Graph()
        self.graph.add_nodes_from(list(range(size)))
        assert obs_type in ["array", "object"]
        self.obs_type = obs_type

    def _encode_obs(self) -> np.ndarray | nx.Graph:
        if self.obs_type == "array":
            return np.stack([v["data"] for v in self.graph._node.values()])
        return deepcopy(self.graph)

    def reset(
        self,
        seed: int | None = None,
        options: dict[str, Any] | None = None,
    ) -> tuple[np.ndarray | nx.Graph, dict]:
        super().reset(seed=seed)
        graph_state = np.random.rand(self.size, self.feat_dim)
        for i in range(self.size):
            self.graph.nodes[i]["data"] = graph_state[i]
        return self._encode_obs(), {}

    def step(
        self,
        action: Space,
    ) -> tuple[np.ndarray | nx.Graph, float, Literal[False], Literal[False], dict]:
        next_graph_state = np.random.rand(self.size, self.feat_dim)
        for i in range(self.size):
            self.graph.nodes[i]["data"] = next_graph_state[i]
        return self._encode_obs(), 1.0, False, False, {}


class MyGoalEnv(MoveToRightEnv):
    def __init__(self, *args: Any, **kwargs: Any) -> None:
        assert (
            kwargs.get("dict_state", 0) + kwargs.get("recurse_state", 0) == 0
        ), "dict_state / recurse_state not supported"
        super().__init__(*args, **kwargs)
        obs, _ = super().reset(options={"state": 0})
        obs, _, _, _, _ = super().step(1)
        self._goal = obs * self.size
        super_obsv = self.observation_space
        self.observation_space = gym.spaces.Dict(
            {
                "observation": super_obsv,
                "achieved_goal": super_obsv,
                "desired_goal": super_obsv,
            },
        )

    def reset(self, *args: Any, **kwargs: Any) -> tuple[dict[str, Any], dict]:
        obs, info = super().reset(*args, **kwargs)
        new_obs = {"observation": obs, "achieved_goal": obs, "desired_goal": self._goal}
        return new_obs, info

    def step(self, *args: Any, **kwargs: Any) -> tuple[dict[str, Any], float, bool, bool, dict]:
        obs_next, rew, terminated, truncated, info = super().step(*args, **kwargs)
        new_obs_next = {
            "observation": obs_next,
            "achieved_goal": obs_next,
            "desired_goal": self._goal,
        }
        return new_obs_next, rew, terminated, truncated, info

    def compute_reward_fn(
        self,
        achieved_goal: np.ndarray,
        desired_goal: np.ndarray,
        info: dict,
    ) -> np.ndarray:
        axis: tuple[int, ...] = (-3, -2, -1) if self.array_state else (-1,)
        return (achieved_goal == desired_goal).all(axis=axis)