[BugFix,Test] test chess rendering

ghstack-source-id: 59b37e6
Pull Request resolved: #2721

Author: Vincent Moens

.github/unittest/linux_libs/scripts_chess/environment.yml

@@ -18,3 +18,5 @@ dependencies:
     - scipy
     - hydra-core
     - chess
+    - transformers
+    - cairosvg

.github/unittest/linux_libs/scripts_chess/install.sh

@@ -28,15 +28,15 @@ git submodule sync && git submodule update --init --recursive
 printf "Installing PyTorch with cu121"
 if [[ "$TORCH_VERSION" == "nightly" ]]; then
   if [ "${CU_VERSION:-}" == cpu ] ; then
-      pip3 install --pre torch --index-url https://download.pytorch.org/whl/nightly/cpu -U
+      pip3 install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cpu -U
   else
-      pip3 install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu121 -U
+      pip3 install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cu121 -U
   fi
 elif [[ "$TORCH_VERSION" == "stable" ]]; then
     if [ "${CU_VERSION:-}" == cpu ] ; then
-      pip3 install torch --index-url https://download.pytorch.org/whl/cpu
+      pip3 install torch torchvision --index-url https://download.pytorch.org/whl/cpu
   else
-      pip3 install torch --index-url https://download.pytorch.org/whl/cu121
+      pip3 install torch torchvision --index-url https://download.pytorch.org/whl/cu121
   fi
 else
   printf "Failed to install pytorch"

test/test_env.py

@@ -172,7 +172,8 @@
     mp_ctx = "fork"
 
 _has_chess = importlib.util.find_spec("chess") is not None
-
+_has_tv = importlib.util.find_spec("torchvision") is not None
+_has_cairosvg = importlib.util.find_spec("cairosvg") is not None
 ## TO BE FIXED: DiscreteActionProjection queries a randint on each worker, which leads to divergent results between
 ## the serial and parallel batched envs
 # def _make_atari_env(atari_env):
@@ -3471,6 +3472,15 @@ def test_env(self, stateful, include_pgn, include_fen, include_hash, include_san
                 if include_san:
                     assert "san_hash" in env.observation_spec.keys()
 
+    @pytest.mark.skipif(not _has_tv, reason="torchvision not found.")
+    @pytest.mark.skipif(not _has_cairosvg, reason="cairosvg not found.")
+    @pytest.mark.parametrize("stateful", [False, True])
+    def test_chess_rendering(self, stateful):
+        env = ChessEnv(stateful=stateful, include_fen=True, pixels=True)
+        env.check_env_specs()
+        r = env.rollout(3)
+        assert "pixels" in r
+
     def test_pgn_bijectivity(self):
         np.random.seed(0)
         pgn = ChessEnv._PGN_RESTART

torchrl/envs/custom/chess.py

@@ -316,7 +316,9 @@ def __init__(
                 raise ImportError(
                     "Please install torchvision to use this environment with pixel rendering."
                 )
-            self.full_observation_spec["pixels"] = Unbounded(shape=())
+            self.full_observation_spec["pixels"] = Unbounded(
+                shape=(3, 390, 390), dtype=torch.uint8
+            )
 
         self.full_action_spec = Composite(
             action=Categorical(n=len(self.san_moves), shape=(), dtype=torch.int64)