[RLlib] Add torch.compile config options to old API stack.

rllib/policy/torch_policy_v2.py

@@ -9,6 +9,7 @@
 
 import gymnasium as gym
 import numpy as np
+from packaging import version
 import tree  # pip install dm_tree
 
 import ray
@@ -41,7 +42,10 @@
 from ray.rllib.utils.numpy import convert_to_numpy
 from ray.rllib.utils.spaces.space_utils import normalize_action
 from ray.rllib.utils.threading import with_lock
-from ray.rllib.utils.torch_utils import convert_to_torch_tensor
+from ray.rllib.utils.torch_utils import (
+    convert_to_torch_tensor,
+    TORCH_COMPILE_REQUIRED_VERSION,
+)
 from ray.rllib.utils.typing import (
     AlgorithmConfigDict,
     GradInfoDict,
@@ -520,6 +524,23 @@ def _init_model_and_dist_class(self):
                 framework=self.framework,
             )
 
+        # Compile the model, if requested by the user.
+        if self.config.get("torch_compile_learner"):
+            if (
+                torch is not None
+                and version.parse(torch.__version__) < TORCH_COMPILE_REQUIRED_VERSION
+            ):
+                raise ValueError("`torch.compile` is not supported for torch < 2.0.0!")
+
+            lw = "learner" if self.config.get("worker_index") else "worker"
+            model = torch.compile(
+                model,
+                backend=self.config.get(
+                    f"torch_compile_{lw}_dynamo_backend", "inductor"
+                ),
+                dynamic=False,
+                mode=self.config.get(f"torch_compile_{lw}_dynamo_mode"),
+            )
         return model, dist_class
 
     @override(Policy)