lazy compile

Author: vmoens

torchrl/collectors/collectors.py

@@ -1227,27 +1227,53 @@ def _setup_policy_and_weights(self, policy: TensorDictModule | Callable) -> None
         else:
             self.policy = self._wrapped_policy = policy
 
-        # Extract policy weights
-        if isinstance(self._wrapped_policy, nn.Module):
+        # Extract policy weights from the uncompiled policy
+        # Access _wrapped_policy_uncompiled directly to avoid triggering compilation
+        if isinstance(self._wrapped_policy_uncompiled, nn.Module):
             self.policy_weights = TensorDict.from_module(
-                self._wrapped_policy, as_module=True
+                self._wrapped_policy_uncompiled, as_module=True
             ).data
         else:
             self.policy_weights = TensorDict()
 
-        # Apply compilation/cudagraph
-        if self.compiled_policy:
-            self._wrapped_policy = compile_with_warmup(
-                self._wrapped_policy, **self.compiled_policy_kwargs
+        # If policy doesn't have meta params, compile immediately
+        # Otherwise, defer until first use (after weights are loaded)
+        if not has_meta_params and (self.compiled_policy or self.cudagraphed_policy):
+            self._wrapped_policy_compiled = self._compile_wrapped_policy(
+                self._wrapped_policy_uncompiled
             )
+
+    def _compile_wrapped_policy(self, policy):
+        """Apply compilation and/or cudagraph to a policy."""
+        if self.compiled_policy:
+            policy = compile_with_warmup(policy, **self.compiled_policy_kwargs)
         if self.cudagraphed_policy:
-            self._wrapped_policy = CudaGraphModule(
-                self._wrapped_policy,
+            policy = CudaGraphModule(
+                policy,
                 in_keys=[],
                 out_keys=[],
                 device=self.policy_device,
                 **self.cudagraphed_policy_kwargs,
             )
+        return policy
+
+    @property
+    def _wrapped_policy(self):
+        """Returns the compiled policy, compiling it lazily if needed."""
+        if (policy := self._wrapped_policy_compiled) is None:
+            if self.compiled_policy or self.cudagraphed_policy:
+                policy = self._wrapped_policy_compiled = self._compile_wrapped_policy(
+                    self._wrapped_policy_uncompiled
+                )
+            else:
+                policy = self._wrapped_policy_compiled = self._wrapped_policy_uncompiled
+        return policy
+
+    @_wrapped_policy.setter
+    def _wrapped_policy(self, value):
+        """Allow setting the wrapped policy during initialization."""
+        self._wrapped_policy_uncompiled = value
+        self._wrapped_policy_compiled = None
 
     def _apply_env_device(self) -> None:
         """Apply device to environment if specified."""
@@ -1425,22 +1451,57 @@ def _maybe_make_final_rollout(self, make_rollout: bool):
                 # erase all devices
                 self._final_rollout.clear_device_()
 
+        # Check if policy has meta-device parameters (not yet initialized)
+        has_meta_params = False
+        if hasattr(self, "_wrapped_policy_uncompiled") and isinstance(
+            self._wrapped_policy_uncompiled, nn.Module
+        ):
+            for p in self._wrapped_policy_uncompiled.parameters():
+                if p.device.type == "meta":
+                    has_meta_params = True
+                    break
+
         # If the policy has a valid spec, we use it
         self._policy_output_keys = set()
         if (
             make_rollout
-            and hasattr(self._wrapped_policy, "spec")
-            and self._wrapped_policy.spec is not None
-            and all(v is not None for v in self._wrapped_policy.spec.values(True, True))
+            and hasattr(
+                self._wrapped_policy_uncompiled
+                if has_meta_params
+                else self._wrapped_policy,
+                "spec",
+            )
+            and (
+                self._wrapped_policy_uncompiled
+                if has_meta_params
+                else self._wrapped_policy
+            ).spec
+            is not None
+            and all(
+                v is not None
+                for v in (
+                    self._wrapped_policy_uncompiled
+                    if has_meta_params
+                    else self._wrapped_policy
+                ).spec.values(True, True)
+            )
         ):
             if any(
                 key not in self._final_rollout.keys(isinstance(key, tuple))
-                for key in self._wrapped_policy.spec.keys(True, True)
+                for key in (
+                    self._wrapped_policy_uncompiled
+                    if has_meta_params
+                    else self._wrapped_policy
+                ).spec.keys(True, True)
             ):
                 # if policy spec is non-empty, all the values are not None and the keys
                 # match the out_keys we assume the user has given all relevant information
                 # the policy could have more keys than the env:
-                policy_spec = self._wrapped_policy.spec
+                policy_spec = (
+                    self._wrapped_policy_uncompiled
+                    if has_meta_params
+                    else self._wrapped_policy
+                ).spec
                 if policy_spec.ndim < self._final_rollout.ndim:
                     policy_spec = policy_spec.expand(self._final_rollout.shape)
                 for key, spec in policy_spec.items(True, True):
@@ -1450,10 +1511,32 @@ def _maybe_make_final_rollout(self, make_rollout: bool):
                     self._final_rollout.set(key, spec.zero())
         elif (
             not make_rollout
-            and hasattr(self._wrapped_policy, "out_keys")
-            and self._wrapped_policy.out_keys
+            and hasattr(
+                self._wrapped_policy_uncompiled
+                if has_meta_params
+                else self._wrapped_policy,
+                "out_keys",
+            )
+            and (
+                self._wrapped_policy_uncompiled
+                if has_meta_params
+                else self._wrapped_policy
+            ).out_keys
         ):
-            self._policy_output_keys = list(self._wrapped_policy.out_keys)
+            self._policy_output_keys = list(
+                (
+                    self._wrapped_policy_uncompiled
+                    if has_meta_params
+                    else self._wrapped_policy
+                ).out_keys
+            )
+        elif has_meta_params:
+            # Policy has meta params and no spec/out_keys - defer initialization
+            # Mark that we need to initialize later when weights are loaded
+            self._policy_output_keys = set()
+            if make_rollout:
+                # We'll populate keys on first actual rollout after weights are loaded
+                self._final_rollout_needs_init = True
         else:
             if make_rollout:
                 # otherwise, we perform a small number of steps with the policy to