[Major, BugFix, Test] Refactor Transforms tests (#878)

Author: vmoens

.circleci/unittest/linux_libs/scripts_gym/batch_scripts.sh

@@ -14,45 +14,87 @@ conda activate ./env
 $DIR/install.sh
 
 # Extracted from run_test.sh to run once.
-yum makecache && yum install libglvnd-devel mesa-libGL mesa-libGL-devel mesa-libEGL mesa-libEGL-devel glfw mesa-libOSMesa-devel glew glew-devel egl-utils freeglut xorg-x11-server-Xvfb -y
+#yum makecache && yum install libglvnd-devel mesa-libGL mesa-libGL-devel mesa-libEGL mesa-libEGL-devel glfw mesa-libOSMesa-devel glew glew-devel egl-utils freeglut xorg-x11-server-Xvfb -y
+yum makecache && yum install libglvnd-devel glew xorg-x11-server-Xvfb zlib-devel egl-utils mesa-libEGL -y
+
 
 # This version is installed initially (see environment.yml)
 for GYM_VERSION in '0.13'
 do
+  # Create a copy of the conda env and work with this
+  conda deactivate
+  conda create --prefix ./cloned_env --clone ./env -y
+
   echo "Testing gym version: ${GYM_VERSION}"
   pip3 install 'gym[atari]'==$GYM_VERSION
   $DIR/run_test.sh
+
+  # delete the conda copy
+  conda deactivate
+  conda env remove -n ./cloned_env
 done
 
 # gym[atari]==0.19 is broken, so we install only gym without dependencies.
 for GYM_VERSION in '0.19'
 do
+  # Create a copy of the conda env and work with this
+  conda deactivate
+  conda create --prefix ./cloned_env --clone ./env -y
+
   echo "Testing gym version: ${GYM_VERSION}"
   pip3 install gym==$GYM_VERSION
   $DIR/run_test.sh
+
+  # delete the conda copy
+  conda deactivate
+  conda env remove -n ./cloned_env
 done
 
 # gym[atari]==0.20 installs ale-py==0.8, but this version is not compatible with gym<0.26, so we downgrade it.
 for GYM_VERSION in '0.20'
 do
+  # Create a copy of the conda env and work with this
+  conda deactivate
+  conda create --prefix ./cloned_env --clone ./env -y
+
   echo "Testing gym version: ${GYM_VERSION}"
   pip3 install 'gym[atari]'==$GYM_VERSION
   pip3 install ale-py==0.7
   $DIR/run_test.sh
+
+  # delete the conda copy
+  conda deactivate
+  conda env remove -n ./cloned_env
 done
 
 for GYM_VERSION in '0.25'
 do
+  # Create a copy of the conda env and work with this
+  conda deactivate
+  conda create --prefix ./cloned_env --clone ./env -y
+
   echo "Testing gym version: ${GYM_VERSION}"
   pip3 install 'gym[atari]'==$GYM_VERSION
   $DIR/run_test.sh
+
+  # delete the conda copy
+  conda deactivate
+  conda env remove -n ./cloned_env
 done
 
 # For this version "gym[accept-rom-license]" is required.
 for GYM_VERSION in '0.26'
 do
+  # Create a copy of the conda env and work with this
+  conda deactivate
+  conda create --prefix ./cloned_env --clone ./env -y
+
   echo "Testing gym version: ${GYM_VERSION}"
   pip3 install 'gym[accept-rom-license]'==$GYM_VERSION
   pip3 install 'gym[atari]'==$GYM_VERSION
   $DIR/run_test.sh
+
+  # delete the conda copy
+  conda deactivate
+  conda env remove -n ./cloned_env
 done

.circleci/unittest/linux_libs/scripts_gym/run_test.sh

@@ -5,7 +5,8 @@ set -e
 eval "$(./conda/bin/conda shell.bash hook)"
 conda activate ./env
 
-yum makecache && yum install libglvnd-devel mesa-libGL mesa-libGL-devel mesa-libEGL mesa-libEGL-devel glfw mesa-libOSMesa-devel glew glew-devel egl-utils freeglut xorg-x11-server-Xvfb -y
+#yum makecache && yum install libglvnd-devel mesa-libGL mesa-libGL-devel mesa-libEGL mesa-libEGL-devel glfw mesa-libOSMesa-devel glew glew-devel egl-utils freeglut xorg-x11-server-Xvfb -y
+#yum makecache && yum install glew -y
 
 export PYTORCH_TEST_WITH_SLOW='1'
 python -m torch.utils.collect_env
@@ -22,5 +23,6 @@ export MKL_THREADING_LAYER=GNU
 
 coverage run -m pytest test/smoke_test.py -v --durations 20
 coverage run -m pytest test/smoke_test_deps.py -v --durations 20 -k 'test_gym or test_dm_control_pixels or test_dm_control'
+MUJOCO_GL=egl python test/test_libs.py -k test_collect
 MUJOCO_GL=egl coverage run -m pytest --instafail -v --durations 20 -k 'test_libs'
 coverage xml -i

.circleci/unittest/linux_libs/scripts_gym/setup_env.sh

@@ -94,7 +94,8 @@ wget http://www.atarimania.com/roms/Roms.rar
 ./rar/unrar e Roms.rar ./Roms -y
 python -m atari_py.import_roms Roms
 
-yum makecache && yum install libglvnd-devel mesa-libGL mesa-libGL-devel mesa-libEGL glfw mesa-libOSMesa-devel glew egl-utils freeglut -y
+#yum makecache && yum install libglvnd-devel mesa-libGL mesa-libGL-devel mesa-libEGL glfw mesa-libOSMesa-devel glew egl-utils freeglut -y
+yum makecache && yum install libglvnd-devel glew zlib-devel -y
 
 # install mujoco-py locally
 cd ${root_dir}/.mujoco/mujoco-py

examples/dreamer/dreamer_utils.py

@@ -89,7 +89,7 @@ def make_env_transforms(
         env.append_transform(Resize(cfg.image_size, cfg.image_size))
         if cfg.grayscale:
             env.append_transform(GrayScale())
-        env.append_transform(FlattenObservation(0, -3))
+        env.append_transform(FlattenObservation(0, -3, allow_positive_dim=True))
         env.append_transform(CatFrames(N=cfg.catframes, in_keys=["pixels"], dim=-3))
         if stats is None:
             obs_stats = {

test/_utils_internal.py

@@ -120,6 +120,16 @@ def dtype_fixture():
     torch.set_default_dtype(dtype)
 
 
+@contextlib.contextmanager
+def set_global_var(module, var_name, value):
+    old_value = getattr(module, var_name)
+    setattr(module, var_name, value)
+    try:
+        yield
+    finally:
+        setattr(module, var_name, old_value)
+
+
 def _make_envs(
     env_name,
     frame_skip,
@@ -263,13 +273,3 @@ def t_out():
             )
 
     return t_out
-
-
-@contextlib.contextmanager
-def set_global_var(module, var_name, value):
-    old_value = getattr(module, var_name)
-    setattr(module, var_name, value)
-    try:
-        yield
-    finally:
-        setattr(module, var_name, old_value)

test/mocking_classes.py

@@ -360,7 +360,7 @@ def __new__(
                 if categorical_action_encoding
                 else OneHotDiscreteTensorSpec
             )
-            action_spec = action_spec_cls(*batch_size, 7)
+            action_spec = action_spec_cls(n=7, shape=(*batch_size, 7))
         if reward_spec is None:
             reward_spec = UnboundedContinuousTensorSpec(shape=(1,))
 

test/test_collector.py

@@ -149,7 +149,7 @@ def _is_consistent_device_type(
     _os_is_windows and _python_is_3_10,
     reason="Windows Access Violation in torch.multiprocessing / BrokenPipeError in multiprocessing.connection",
 )
-@pytest.mark.parametrize("num_env", [1, 3])
+@pytest.mark.parametrize("num_env", [1, 2])
 @pytest.mark.parametrize("device", ["cuda", "cpu", None])
 @pytest.mark.parametrize("policy_device", ["cuda", "cpu", None])
 @pytest.mark.parametrize("passing_device", ["cuda", "cpu", None])
@@ -233,7 +233,7 @@ def env_fn(seed):
     ccollector.shutdown()
 
 
-@pytest.mark.parametrize("num_env", [1, 3])
+@pytest.mark.parametrize("num_env", [1, 2])
 @pytest.mark.parametrize("env_name", ["conv", "vec"])
 def test_concurrent_collector_consistency(num_env, env_name, seed=40):
     if num_env == 1:
@@ -309,7 +309,7 @@ def make_env():
         return GymEnv(PONG_VERSIONED, frame_skip=4)
 
     env = SerialEnv(2, make_env)
-    # env = SerialEnv(3, lambda: GymEnv("CartPole-v1", frame_skip=4))
+    # env = SerialEnv(2, lambda: GymEnv("CartPole-v1", frame_skip=4))
     env.set_seed(0)
     collector = SyncDataCollector(
         env, total_frames=10000, frames_per_batch=10000, split_trajs=False
@@ -331,7 +331,7 @@ def make_env():
     assert _data["reward"].sum(-2).min() == -21
 
 
-@pytest.mark.parametrize("num_env", [1, 3])
+@pytest.mark.parametrize("num_env", [1, 2])
 @pytest.mark.parametrize("env_name", ["vec"])
 def test_collector_done_persist(num_env, env_name, seed=5):
     if num_env == 1:
@@ -381,7 +381,7 @@ def make_env(seed):
 
 
 @pytest.mark.parametrize("frames_per_batch", [200, 10])
-@pytest.mark.parametrize("num_env", [1, 3])
+@pytest.mark.parametrize("num_env", [1, 2])
 @pytest.mark.parametrize("env_name", ["vec"])
 def test_split_trajs(num_env, env_name, frames_per_batch, seed=5):
     if num_env == 1:
@@ -475,7 +475,7 @@ def make_env(seed):
 #         ccollector.shutdown()
 
 
-@pytest.mark.parametrize("num_env", [1, 3])
+@pytest.mark.parametrize("num_env", [1, 2])
 @pytest.mark.parametrize("env_name", ["vec", "conv"])
 def test_collector_batch_size(num_env, env_name, seed=100):
     if num_env == 3 and _os_is_windows:
@@ -498,7 +498,7 @@ def env_fn():
 
     torch.manual_seed(0)
     np.random.seed(0)
-    num_workers = 4
+    num_workers = 2
     frames_per_batch = 20
     ccollector = MultiaSyncDataCollector(
         create_env_fn=[env_fn for _ in range(num_workers)],
@@ -534,7 +534,7 @@ def env_fn():
     ccollector.shutdown()
 
 
-@pytest.mark.parametrize("num_env", [1, 3])
+@pytest.mark.parametrize("num_env", [1, 2])
 @pytest.mark.parametrize("env_name", ["vec", "conv"])
 def test_concurrent_collector_seed(num_env, env_name, seed=100):
     if num_env == 1:
@@ -581,7 +581,7 @@ def env_fn():
     ccollector.shutdown()
 
 
-@pytest.mark.parametrize("num_env", [1, 3])
+@pytest.mark.parametrize("num_env", [1, 2])
 @pytest.mark.parametrize("env_name", ["conv", "vec"])
 def test_collector_consistency(num_env, env_name, seed=100):
     if num_env == 1:
@@ -644,7 +644,7 @@ def env_fn(seed):
     collector.shutdown()
 
 
-@pytest.mark.parametrize("num_env", [1, 3])
+@pytest.mark.parametrize("num_env", [1, 2])
 @pytest.mark.parametrize("collector_class", [SyncDataCollector, aSyncDataCollector])
 @pytest.mark.parametrize("env_name", ["conv", "vec"])
 def test_traj_len_consistency(num_env, env_name, collector_class, seed=100):
@@ -1100,7 +1100,7 @@ def env_fn(seed):
     ],
 )
 class TestAutoWrap:
-    num_envs = 3
+    num_envs = 2
 
     @pytest.fixture
     def env_maker(self):

test/test_env.py

@@ -40,7 +40,7 @@
 from torchrl.envs.gym_like import default_info_dict_reader
 from torchrl.envs.libs.dm_control import _has_dmc, DMControlEnv
 from torchrl.envs.libs.gym import _has_gym, GymEnv, GymWrapper
-from torchrl.envs.transforms import Compose, TransformedEnv
+from torchrl.envs.transforms import Compose, StepCounter, TransformedEnv
 from torchrl.envs.utils import step_mdp
 from torchrl.envs.vec_env import ParallelEnv, SerialEnv
 from torchrl.modules import Actor, ActorCriticOperator, MLP, SafeModule, ValueOperator
@@ -195,6 +195,38 @@ def test_rollout_predictability(device):
     ).all()
 
 
+@pytest.mark.skipif(not _has_gym, reason="no gym")
+@pytest.mark.parametrize(
+    "env_name",
+    [
+        PENDULUM_VERSIONED,
+    ],
+)
+@pytest.mark.parametrize(
+    "frame_skip",
+    [
+        1,
+    ],
+)
+@pytest.mark.parametrize("parallel", [False, True])
+def test_rollout_reset(env_name, frame_skip, parallel, seed=0):
+    envs = []
+    for horizon in [20, 30, 40]:
+        envs.append(
+            lambda horizon=horizon: TransformedEnv(
+                GymEnv(env_name, frame_skip=frame_skip), StepCounter(horizon)
+            )
+        )
+    if parallel:
+        env = ParallelEnv(3, envs)
+    else:
+        env = SerialEnv(3, envs)
+    env.set_seed(100)
+    out = env.rollout(100, break_when_any_done=False)
+    assert out.shape == torch.Size([3, 100])
+    assert (out["done"].squeeze().sum(-1) == torch.tensor([5, 3, 2])).all()
+
+
 class TestModelBasedEnvBase:
     @pytest.mark.parametrize("device", get_available_devices())
     def test_mb_rollout(self, device, seed=0):
@@ -585,9 +617,9 @@ def test_parallel_env_custom_method(self, parallel):
         # define env
 
         if parallel:
-            env = ParallelEnv(3, lambda: DiscreteActionVecMockEnv())
+            env = ParallelEnv(2, lambda: DiscreteActionVecMockEnv())
         else:
-            env = SerialEnv(3, lambda: DiscreteActionVecMockEnv())
+            env = SerialEnv(2, lambda: DiscreteActionVecMockEnv())
 
         # we must start the environment first
         env.reset()

test/test_exploration.py

@@ -265,7 +265,7 @@ def test_gsde(
         device=device,
     )
     if gSDE:
-        gSDENoise().reset(td)
+        gSDENoise(shape=[batch]).reset(td)
         assert "_eps_gSDE" in td.keys()
         assert td.get("_eps_gSDE").device == device
     actor(td)