amend

pytorch · vmoens · Apr 25, 2024 · Apr 24, 2024 · Apr 24, 2024 · Apr 24, 2024
commit b3f06c09ce813f938f79a651e44420c8a07625b1
diff --git a/test/_utils_internal.py b/test/_utils_internal.py
@@ -560,10 +560,6 @@ def __init__(
         *,
         lstm_backend: str | None = None,
     ) -> None:
-        warnings.warn(
-            "LSTMNet is being deprecated in favour of torchrl.modules.LSTMModule, and will be removed in v0.4.0.",
-            category=DeprecationWarning,
-        )
         super().__init__()
         lstm_kwargs.update({"batch_first": True})
         self.mlp = MLP(device=device, **mlp_kwargs)

diff --git a/test/mocking_classes.py b/test/mocking_classes.py
@@ -925,7 +925,6 @@ def __init__(
         super().__init__(
             world_model,
             device=device,
-            dtype=dtype,
             batch_size=batch_size,
         )
         self.observation_spec = CompositeSpec(

diff --git a/test/test_helpers.py b/test/test_helpers.py
@@ -50,8 +50,6 @@
     DiscreteModelConfig,
     DreamerConfig,
     make_dqn_actor,
-    make_redq_model,
-    REDQModelConfig,
 )
 
 TORCH_VERSION = version.parse(torch.__version__)
@@ -162,123 +160,6 @@ def test_dqn_maker(
         proof_environment.close()
 
 
-@pytest.mark.skipif(not _has_functorch, reason="functorch not installed")
-@pytest.mark.skipif(not _has_hydra, reason="No hydra library found")
-@pytest.mark.skipif(not _has_gym, reason="No gym library found")
-@pytest.mark.parametrize("device", get_default_devices())
-@pytest.mark.parametrize("from_pixels", [(), ("from_pixels=True", "catframes=4")])
-@pytest.mark.parametrize("gsde", [(), ("gSDE=True",)])
-@pytest.mark.parametrize("exploration", [ExplorationType.MODE, ExplorationType.RANDOM])
-def test_redq_make(device, from_pixels, gsde, exploration):
-    if not gsde and exploration != ExplorationType.RANDOM:
-        pytest.skip("no need to test this setting")
-    flags = list(from_pixels + gsde)
-    if gsde and from_pixels:
-        pytest.skip("gsde and from_pixels are incompatible")
-
-    config_fields = [
-        (config_field.name, config_field.type, config_field)
-        for config_cls in (
-            EnvConfig,
-            REDQModelConfig,
-        )
-        for config_field in dataclasses.fields(config_cls)
-    ]
-
-    Config = dataclasses.make_dataclass(cls_name="Config", fields=config_fields)
-    cs = ConfigStore.instance()
-    cs.store(name="config", node=Config)
-    with initialize(version_base="1.1", config_path=None):
-        cfg = compose(config_name="config", overrides=flags)
-
-        env_maker = (
-            ContinuousActionConvMockEnvNumpy
-            if from_pixels
-            else ContinuousActionVecMockEnv
-        )
-        env_maker = transformed_env_constructor(
-            cfg,
-            use_env_creator=False,
-            custom_env_maker=env_maker,
-            stats={"loc": 0.0, "scale": 1.0},
-        )
-        proof_environment = env_maker()
-
-        model = make_redq_model(
-            proof_environment,
-            device=device,
-            cfg=cfg,
-        )
-        actor, qvalue = model
-        td = proof_environment.reset().to(device)
-        with set_exploration_type(exploration):
-            actor(td)
-        expected_keys = [
-            "done",
-            "terminated",
-            "action",
-            "sample_log_prob",
-            "loc",
-            "scale",
-            "step_count",
-            "is_init",
-        ]
-        if len(gsde):
-            expected_keys += ["_eps_gSDE"]
-        if from_pixels:
-            expected_keys += [
-                "hidden",
-                "pixels",
-                "pixels_orig",
-            ]
-        else:
-            expected_keys += ["observation_vector", "observation_orig"]
-
-        try:
-            assert set(td.keys()) == set(expected_keys)
-        except AssertionError:
-            proof_environment.close()
-            raise
-
-        if cfg.gSDE:
-            tsf_loc = actor.module[0].module[-1].module.transform(td.get("loc"))
-            if exploration == ExplorationType.RANDOM:
-                with pytest.raises(AssertionError):
-                    torch.testing.assert_close(td.get("action"), tsf_loc)
-            else:
-                torch.testing.assert_close(td.get("action"), tsf_loc)
-
-        qvalue(td)
-        expected_keys = [
-            "done",
-            "terminated",
-            "action",
-            "sample_log_prob",
-            "state_action_value",
-            "loc",
-            "scale",
-            "step_count",
-            "is_init",
-        ]
-        if len(gsde):
-            expected_keys += ["_eps_gSDE"]
-        if from_pixels:
-            expected_keys += [
-                "hidden",
-                "pixels",
-                "pixels_orig",
-            ]
-        else:
-            expected_keys += ["observation_vector", "observation_orig"]
-        try:
-            assert set(td.keys()) == set(expected_keys)
-        except AssertionError:
-            proof_environment.close()
-            raise
-        proof_environment.close()
-        del proof_environment
-
-
 @pytest.mark.parametrize("initial_seed", range(5))
 def test_seed_generator(initial_seed):
     num_seeds = 100

diff --git a/test/test_modules.py b/test/test_modules.py
@@ -23,7 +23,6 @@
     GRUCell,
     LSTM,
     LSTMCell,
-    LSTMNet,
     MultiAgentConvNet,
     MultiAgentMLP,
     OnlineDTActor,
@@ -350,140 +349,6 @@ def test_noisy(layer_class, device, seed=0):
         torch.testing.assert_close(y1, y2)
 
 
-@pytest.mark.parametrize("device", get_default_devices())
-@pytest.mark.parametrize("out_features", [3, 4])
-@pytest.mark.parametrize("hidden_size", [8, 9])
-@pytest.mark.parametrize("num_layers", [1, 2])
-@pytest.mark.parametrize("has_precond_hidden", [True, False])
-def test_lstm_net(
-    device,
-    out_features,
-    hidden_size,
-    num_layers,
-    has_precond_hidden,
-    double_prec_fixture,
-):
-    torch.manual_seed(0)
-    batch = 5
-    time_steps = 6
-    in_features = 7
-    net = LSTMNet(
-        out_features,
-        {
-            "input_size": hidden_size,
-            "hidden_size": hidden_size,
-            "num_layers": num_layers,
-        },
-        {"out_features": hidden_size},
-        device=device,
-    )
-    # test single step vs multi-step
-    x = torch.randn(batch, time_steps, in_features, device=device)
-    x_unbind = x.unbind(1)
-    tds_loop = []
-    if has_precond_hidden:
-        hidden0_out0, hidden1_out0 = torch.randn(
-            2, batch, time_steps, num_layers, hidden_size, device=device
-        )
-        hidden0_out0[:, 1:] = 0.0
-        hidden1_out0[:, 1:] = 0.0
-        hidden0_out = hidden0_out0[:, 0]
-        hidden1_out = hidden1_out0[:, 0]
-    else:
-        hidden0_out, hidden1_out = None, None
-        hidden0_out0, hidden1_out0 = None, None
-
-    for _x in x_unbind:
-        y, hidden0_in, hidden1_in, hidden0_out, hidden1_out = net(
-            _x, hidden0_out, hidden1_out
-        )
-        td = TensorDict(
-            {
-                "y": y,
-                "hidden0_in": hidden0_in,
-                "hidden1_in": hidden1_in,
-                "hidden0_out": hidden0_out,
-                "hidden1_out": hidden1_out,
-            },
-            [batch],
-        )
-        tds_loop.append(td)
-    tds_loop = torch.stack(tds_loop, 1)
-
-    y, hidden0_in, hidden1_in, hidden0_out, hidden1_out = net(
-        x, hidden0_out0, hidden1_out0
-    )
-    tds_vec = TensorDict(
-        {
-            "y": y,
-            "hidden0_in": hidden0_in,
-            "hidden1_in": hidden1_in,
-            "hidden0_out": hidden0_out,
-            "hidden1_out": hidden1_out,
-        },
-        [batch, time_steps],
-    )
-    torch.testing.assert_close(tds_vec["y"], tds_loop["y"])
-    torch.testing.assert_close(
-        tds_vec["hidden0_out"][:, -1], tds_loop["hidden0_out"][:, -1]
-    )
-    torch.testing.assert_close(
-        tds_vec["hidden1_out"][:, -1], tds_loop["hidden1_out"][:, -1]
-    )
-
-
-@pytest.mark.parametrize("device", get_default_devices())
-@pytest.mark.parametrize("out_features", [3, 5])
-@pytest.mark.parametrize("hidden_size", [3, 5])
-def test_lstm_net_nobatch(device, out_features, hidden_size):
-    time_steps = 6
-    in_features = 4
-    net = LSTMNet(
-        out_features,
-        {"input_size": hidden_size, "hidden_size": hidden_size},
-        {"out_features": hidden_size},
-        device=device,
-    )
-    # test single step vs multi-step
-    x = torch.randn(time_steps, in_features, device=device)
-    x_unbind = x.unbind(0)
-    tds_loop = []
-    hidden0_in, hidden1_in, hidden0_out, hidden1_out = [
-        None,
-    ] * 4
-    for _x in x_unbind:
-        y, hidden0_in, hidden1_in, hidden0_out, hidden1_out = net(
-            _x, hidden0_out, hidden1_out
-        )
-        td = TensorDict(
-            {
-                "y": y,
-                "hidden0_in": hidden0_in,
-                "hidden1_in": hidden1_in,
-                "hidden0_out": hidden0_out,
-                "hidden1_out": hidden1_out,
-            },
-            [],
-        )
-        tds_loop.append(td)
-    tds_loop = torch.stack(tds_loop, 0)
-
-    y, hidden0_in, hidden1_in, hidden0_out, hidden1_out = net(x.unsqueeze(0))
-    tds_vec = TensorDict(
-        {
-            "y": y,
-            "hidden0_in": hidden0_in,
-            "hidden1_in": hidden1_in,
-            "hidden0_out": hidden0_out,
-            "hidden1_out": hidden1_out,
-        },
-        [1, time_steps],
-    ).squeeze(0)
-    torch.testing.assert_close(tds_vec["y"], tds_loop["y"])
-    torch.testing.assert_close(tds_vec["hidden0_out"][-1], tds_loop["hidden0_out"][-1])
-    torch.testing.assert_close(tds_vec["hidden1_out"][-1], tds_loop["hidden1_out"][-1])
-
-
 @pytest.mark.parametrize("device", get_default_devices())
 @pytest.mark.parametrize("batch_size", [3, 5])
 class TestPlanner:

diff --git a/torchrl/envs/model_based/common.py b/torchrl/envs/model_based/common.py
@@ -117,13 +117,11 @@ def __init__(
         params: Optional[List[torch.Tensor]] = None,
         buffers: Optional[List[torch.Tensor]] = None,
         device: DEVICE_TYPING = "cpu",
-        dtype: Optional[Union[torch.dtype, np.dtype]] = None,
         batch_size: Optional[torch.Size] = None,
         run_type_checks: bool = False,
     ):
         super(ModelBasedEnvBase, self).__init__(
             device=device,
-            dtype=dtype,
             batch_size=batch_size,
             run_type_checks=run_type_checks,
         )

diff --git a/torchrl/envs/model_based/dreamer.py b/torchrl/envs/model_based/dreamer.py
@@ -27,11 +27,10 @@ def __init__(
         belief_shape: Tuple[int, ...],
         obs_decoder: TensorDictModule = None,
         device: DEVICE_TYPING = "cpu",
-        dtype: Optional[Union[torch.dtype, np.dtype]] = None,
         batch_size: Optional[torch.Size] = None,
     ):
         super(DreamerEnv, self).__init__(
-            world_model, device=device, dtype=dtype, batch_size=batch_size
+            world_model, device=device, batch_size=batch_size
         )
         self.obs_decoder = obs_decoder
         self.prior_shape = prior_shape

diff --git a/torchrl/trainers/helpers/__init__.py b/torchrl/trainers/helpers/__init__.py
@@ -16,7 +16,7 @@
     transformed_env_constructor,
 )
 from .logger import LoggerConfig
-from .losses import make_dqn_loss, make_redq_loss, make_target_updater
-from .models import make_dqn_actor, make_dreamer, make_redq_model
+from .losses import make_dqn_loss, make_target_updater
+from .models import make_dqn_actor, make_dreamer
 from .replay_buffer import make_replay_buffer
 from .trainers import make_trainer