Feature/crafter

README.md

@@ -54,6 +54,7 @@ The environments supported by sheeprl are:
 | MineRL             | `pip install -e .[minerl]`   | [how_to/minerl](./howto/learn_in_minerl.md)     | :heavy_check_mark: |
 | MineDojo           | `pip install -e .[minedojo]` | [how_to/minedojo](./howto/learn_in_minedojo.md) | :heavy_check_mark: |
 | DIAMBRA            | `pip install -e .[diambra]`  | [how_to/diambra](./howto/learn_in_diambra.md)   | :heavy_check_mark: |
+| Crafter            | `pip install -e .[crafter]`  | https://github.com/danijar/crafter              | :heavy_check_mark: |
 
 
 ## Why

pyproject.toml

@@ -67,6 +67,7 @@ atari = [
 minedojo = ["minedojo==0.1", "importlib_resources==5.12.0"]
 minerl = ["minerl==0.4.4"]
 diambra = ["wheel==0.38.4", "setuptools<=66.0.0", "gym==0.21.0", "diambra==0.0.16", "diambra-arena==2.1.2"]
+crafter = ["crafter==1.8.1"]
 
 [tool.ruff]
 line-length = 120

sheeprl/algos/dreamer_v1/dreamer_v1.py

@@ -413,8 +413,8 @@ def main(fabric: Fabric, cfg: DictConfig):
         raise RuntimeError(f"Unexpected observation type, should be of type Dict, got: {observation_space}")
     if cfg.cnn_keys.encoder == [] and cfg.mlp_keys.encoder == []:
         raise RuntimeError(
-            "You should specify at least one CNN keys or MLP keys from the cli: `cnn_keys.encoder=[rgb]` "
-            "or `mlp_keys.encoder=[state]` "
+            "You should specify at least one CNN keys or MLP keys from the cli: "
+            "`cnn_keys.encoder=[rgb]` or `mlp_keys.encoder=[state]`"
         )
     if (
         len(set(cfg.cnn_keys.encoder).intersection(set(cfg.cnn_keys.decoder))) == 0

sheeprl/algos/dreamer_v2/dreamer_v2.py

@@ -432,8 +432,8 @@ def main(fabric: Fabric, cfg: DictConfig):
         raise RuntimeError(f"Unexpected observation type, should be of type Dict, got: {observation_space}")
     if cfg.cnn_keys.encoder == [] and cfg.mlp_keys.encoder == []:
         raise RuntimeError(
-            "You should specify at least one CNN keys or MLP keys from the cli: `cnn_keys.encoder=[rgb]` "
-            "or `mlp_keys.encoder=[state]` "
+            "You should specify at least one CNN keys or MLP keys from the cli: "
+            "`cnn_keys.encoder=[rgb]` or `mlp_keys.encoder=[state]`"
         )
     if (
         len(set(cfg.cnn_keys.encoder).intersection(set(cfg.cnn_keys.decoder))) == 0

sheeprl/algos/dreamer_v3/dreamer_v3.py

@@ -389,8 +389,8 @@ def main(fabric: Fabric, cfg: DictConfig):
         raise RuntimeError(f"Unexpected observation type, should be of type Dict, got: {observation_space}")
     if cfg.cnn_keys.encoder == [] and cfg.mlp_keys.encoder == []:
         raise RuntimeError(
-            "You should specify at least one CNN keys or MLP keys from the cli: `cnn_keys.encoder=[rgb]` "
-            "or `mlp_keys.encoder=[state]` "
+            "You should specify at least one CNN keys or MLP keys from the cli: "
+            "`cnn_keys.encoder=[rgb]` or `mlp_keys.encoder=[state]`"
         )
     if (
         len(set(cfg.cnn_keys.encoder).intersection(set(cfg.cnn_keys.decoder))) == 0

sheeprl/algos/p2e_dv1/p2e_dv1.py

@@ -417,8 +417,8 @@ def main(fabric: Fabric, cfg: DictConfig):
         raise RuntimeError(f"Unexpected observation type, should be of type Dict, got: {observation_space}")
     if cfg.cnn_keys.encoder == [] and cfg.mlp_keys.encoder == []:
         raise RuntimeError(
-            "You should specify at least one CNN keys or MLP keys from the cli: `cnn_keys.encoder=[rgb]` "
-            "or `mlp_keys.encoder=[state]` "
+            "You should specify at least one CNN keys or MLP keys from the cli: "
+            "`cnn_keys.encoder=[rgb]` or `mlp_keys.encoder=[state]`"
         )
     if (
         len(set(cfg.cnn_keys.encoder).intersection(set(cfg.cnn_keys.decoder))) == 0

sheeprl/algos/p2e_dv2/p2e_dv2.py

@@ -519,8 +519,8 @@ def main(fabric: Fabric, cfg: DictConfig):
         raise RuntimeError(f"Unexpected observation type, should be of type Dict, got: {observation_space}")
     if cfg.cnn_keys.encoder == [] and cfg.mlp_keys.encoder == []:
         raise RuntimeError(
-            "You should specify at least one CNN keys or MLP keys from the cli: `cnn_keys.encoder=[rgb]` "
-            "or `mlp_keys.encoder=[state]` "
+            "You should specify at least one CNN keys or MLP keys from the cli: "
+            "`cnn_keys.encoder=[rgb]` or `mlp_keys.encoder=[state]`"
         )
     if (
         len(set(cfg.cnn_keys.encoder).intersection(set(cfg.cnn_keys.decoder))) == 0

sheeprl/algos/ppo/ppo.py

@@ -164,8 +164,8 @@ def main(fabric: Fabric, cfg: DictConfig):
         raise RuntimeError(f"Unexpected observation type, should be of type Dict, got: {observation_space}")
     if cfg.cnn_keys.encoder + cfg.mlp_keys.encoder == []:
         raise RuntimeError(
-            "You should specify at least one CNN keys or MLP keys from the cli: `cnn_keys.encoder=[rgb]` "
-            "or `mlp_keys.encoder=[state]` "
+            "You should specify at least one CNN keys or MLP keys from the cli: "
+            "`cnn_keys.encoder=[rgb]` or `mlp_keys.encoder=[state]`"
         )
     fabric.print("Encoder CNN keys:", cfg.cnn_keys.encoder)
     fabric.print("Encoder MLP keys:", cfg.mlp_keys.encoder)

sheeprl/algos/ppo/ppo_decoupled.py

@@ -75,8 +75,8 @@ def player(
         raise RuntimeError(f"Unexpected observation type, should be of type Dict, got: {observation_space}")
     if cfg.cnn_keys.encoder + cfg.mlp_keys.encoder == []:
         raise RuntimeError(
-            "You should specify at least one CNN keys or MLP keys from the cli: `cnn_keys.encoder=[rgb]` "
-            "or `mlp_keys.encoder=[state]` "
+            "You should specify at least one CNN keys or MLP keys from the cli: "
+            "`cnn_keys.encoder=[rgb]` or `mlp_keys.encoder=[state]`"
         )
     fabric.print("Encoder CNN keys:", cfg.cnn_keys.encoder)
     fabric.print("Encoder MLP keys:", cfg.mlp_keys.encoder)

sheeprl/algos/sac_ae/sac_ae.py

@@ -187,8 +187,8 @@ def main(fabric: Fabric, cfg: DictConfig):
         raise RuntimeError(f"Unexpected observation type, should be of type Dict, got: {observation_space}")
     if cfg.cnn_keys.encoder == [] and cfg.mlp_keys.encoder == []:
         raise RuntimeError(
-            "You should specify at least one CNN keys or MLP keys from the cli: `cnn_keys.encoder=[rgb]` "
-            "or `mlp_keys.encoder=[state]` "
+            "You should specify at least one CNN keys or MLP keys from the cli: "
+            "`cnn_keys.encoder=[rgb]` or `mlp_keys.encoder=[state]`"
         )
     if (
         len(set(cfg.cnn_keys.encoder).intersection(set(cfg.cnn_keys.decoder))) == 0

sheeprl/configs/env/crafter.yaml

@@ -0,0 +1,16 @@
+defaults:
+  - default
+  - _self_
+
+# Override from `default` config
+id: reward
+action_repeat: 1
+capture_video: False
+reward_as_observation: True
+
+# Wrapper to be instantiated
+wrapper:
+  _target_: sheeprl.envs.crafter.CrafterWrapper
+  id: ${env.id}
+  screen_size: ${env.screen_size}
+  seed: ${seed}

sheeprl/configs/exp/dreamer_v3_XL_crafter.yaml

@@ -0,0 +1,50 @@
+# @package _global_
+
+defaults:
+  - dreamer_v3
+  - override /env: crafter
+  - _self_
+
+# Experiment
+seed: 5
+total_steps: 1000000
+
+# Environment
+env:
+  num_envs: 1
+  id: reward
+
+# Checkpoint
+checkpoint:
+  every: 100000
+
+# Buffer
+buffer:
+  checkpoint: True
+
+# The CNN and MLP keys of the decoder are the same as those of the encoder by default
+cnn_keys:
+  encoder:
+    - rgb
+  decoder:
+    - rgb
+mlp_keys:
+  encoder:
+    - reward
+  decoder: []
+
+# Algorithm
+algo:
+  train_every: 2
+  learning_starts: 1024
+  dense_units: 1024
+  mlp_layers: 5
+  world_model:
+    encoder:
+      cnn_channels_multiplier: 96
+    recurrent_model:
+      recurrent_state_size: 4096
+    transition_model:
+      hidden_size: 1024
+    representation_model:
+      hidden_size: 1024

sheeprl/envs/crafter.py

@@ -0,0 +1,60 @@
+from sheeprl.utils.imports import _IS_CRAFTER_AVAILABLE
+
+if not _IS_CRAFTER_AVAILABLE:
+    raise ModuleNotFoundError(_IS_CRAFTER_AVAILABLE)
+
+from typing import Any, Dict, List, Optional, SupportsFloat, Tuple, Union
+
+import crafter
+import numpy as np
+from gymnasium import core, spaces
+from gymnasium.core import RenderFrame
+
+
+class CrafterWrapper(core.Env):
+    def __init__(self, id: str, screen_size: Union[int, Tuple[int, int]] = 64, seed: Optional[int] = None) -> None:
+        assert id in {"reward", "nonreward"}
+        if isinstance(screen_size, int):
+            screen_size = (screen_size,) * 2
+
+        self._env = crafter.Env(size=screen_size, seed=seed, reward=(id == "reward"))
+        self.observation_space = spaces.Dict(
+            {
+                "rgb": spaces.Box(
+                    self._env.observation_space.low,
+                    self._env.observation_space.high,
+                    self._env.observation_space.shape,
+                    self._env.observation_space.dtype,
+                )
+            }
+        )
+        self.action_space = spaces.Discrete(self._env.action_space.n)
+        self.reward_range = self._env.reward_range or (-np.inf, np.inf)
+        self.observation_space.seed(seed)
+        self.action_space.seed(seed)
+
+        # render
+        self._render_mode: str = "rgb_array"
+
+    @property
+    def render_mode(self) -> str:
+        return self._render_mode
+
+    def _convert_obs(self, obs: np.ndarray) -> Dict[str, np.ndarray]:
+        return {"rgb": obs}
+
+    def step(self, action: Any) -> Tuple[Any, SupportsFloat, bool, bool, Dict[str, Any]]:
+        obs, reward, done, info = self._env.step(action)
+        return self._convert_obs(obs), reward, done, False, info
+
+    def reset(
+        self, *, seed: Optional[int] = None, options: Optional[Dict[str, Any]] = None
+    ) -> Tuple[Any, Dict[str, Any]]:
+        obs = self._env.reset()
+        return self._convert_obs(obs), {}
+
+    def render(self) -> Optional[Union[RenderFrame, List[RenderFrame]]]:
+        return self._env.render()
+
+    def close(self) -> None:
+        return super().close()

sheeprl/envs/wrappers.py

@@ -201,7 +201,9 @@ class RewardAsObservationWrapper(gym.Wrapper):
     def __init__(self, env: Env) -> None:
         super().__init__(env)
         self._env = env
-        reward_range = self._env.reward_range if hasattr(self._env, "reward_range") else (-np.inf, np.inf)
+        reward_range = (
+            self._env.reward_range or (-np.inf, np.inf) if hasattr(self._env, "reward_range") else (-np.inf, np.inf)
+        )
         # The reward is assumed to be a scalar
         if isinstance(self._env.observation_space, gym.spaces.Dict):
             self.observation_space = gym.spaces.Dict(

sheeprl/utils/imports.py

@@ -4,6 +4,7 @@
 
 _IS_ATARI_AVAILABLE = RequirementCache("gymnasium[atari]")
 _IS_ATARI_ROMS_AVAILABLE = RequirementCache("gymnasium[accept-rom-license]")
+_IS_CRAFTER_AVAILABLE = RequirementCache("crafter")
 _IS_DIAMBRA_AVAILABLE = RequirementCache("diambra")
 _IS_DIAMBRA_ARENA_AVAILABLE = RequirementCache("diambra-arena")
 _IS_DMC_AVAILABLE = RequirementCache("dm_control")