Update

[ghstack-poisoned]

Author: vmoens

test/test_weightsync.py

@@ -272,114 +272,6 @@ def test_no_weight_sync_scheme(self):
         transport.send_weights("policy", weights)
 
 
-class TestMultiModelUpdates:
-    def test_multi_model_state_dict_updates(self):
-        env = GymEnv("CartPole-v1")
-
-        policy = TensorDictModule(
-            nn.Linear(
-                env.observation_spec["observation"].shape[-1], env.action_spec.shape[-1]
-            ),
-            in_keys=["observation"],
-            out_keys=["action"],
-        )
-
-        value = TensorDictModule(
-            nn.Linear(env.observation_spec["observation"].shape[-1], 1),
-            in_keys=["observation"],
-            out_keys=["value"],
-        )
-
-        weight_sync_schemes = {
-            "policy": MultiProcessWeightSyncScheme(strategy="state_dict"),
-            "value": MultiProcessWeightSyncScheme(strategy="state_dict"),
-        }
-
-        collector = SyncDataCollector(
-            create_env_fn=lambda: GymEnv("CartPole-v1"),
-            policy=policy,
-            frames_per_batch=64,
-            total_frames=128,
-            weight_sync_schemes=weight_sync_schemes,
-        )
-
-        policy_weights = policy.state_dict()
-        value_weights = value.state_dict()
-
-        with torch.no_grad():
-            for key in policy_weights:
-                policy_weights[key].fill_(1.0)
-            for key in value_weights:
-                value_weights[key].fill_(2.0)
-
-        collector.update_policy_weights_(
-            weights_dict={
-                "policy": policy_weights,
-                "value": value_weights,
-            }
-        )
-
-        for data in collector:
-            assert data.numel() > 0
-            break
-
-        collector.shutdown()
-        env.close()
-
-    def test_multi_model_tensordict_updates(self):
-        env = GymEnv("CartPole-v1")
-
-        policy = TensorDictModule(
-            nn.Linear(
-                env.observation_spec["observation"].shape[-1], env.action_spec.shape[-1]
-            ),
-            in_keys=["observation"],
-            out_keys=["action"],
-        )
-
-        value = TensorDictModule(
-            nn.Linear(env.observation_spec["observation"].shape[-1], 1),
-            in_keys=["observation"],
-            out_keys=["value"],
-        )
-
-        weight_sync_schemes = {
-            "policy": MultiProcessWeightSyncScheme(strategy="tensordict"),
-            "value": MultiProcessWeightSyncScheme(strategy="tensordict"),
-        }
-
-        collector = SyncDataCollector(
-            create_env_fn=lambda: GymEnv("CartPole-v1"),
-            policy=policy,
-            frames_per_batch=64,
-            total_frames=128,
-            weight_sync_schemes=weight_sync_schemes,
-        )
-
-        policy_weights = TensorDict.from_module(policy)
-        value_weights = TensorDict.from_module(value)
-
-        with torch.no_grad():
-            policy_weights["module"]["weight"].fill_(1.0)
-            policy_weights["module"]["bias"].fill_(1.0)
-            value_weights["module"]["weight"].fill_(2.0)
-            value_weights["module"]["bias"].fill_(2.0)
-
-        collector.update_policy_weights_(
-            weights_dict={
-                "policy": policy_weights,
-                "value": value_weights,
-            }
-        )
-
-        for data in collector:
-            assert data.numel() > 0
-            break
-
-        collector.shutdown()
-        env.close()
-
-
 class TestHelpers:
     def test_resolve_model_simple(self):
         class Context: