fix(inference): use spawn context to avoid CUDA fork error in inference worker

Replace fork-based multiprocessing with spawn context in LerobotPolicy to
prevent RuntimeError when CUDA is already initialized in the parent process.
Extract observation_space and env_metadata from env before spawning to avoid
pickling unpicklable ROS2 handles. Add spawn verification test to test script.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

Author: ZhangYi1999

crisp_gym/policy/lerobot_policy.py

@@ -2,7 +2,7 @@
 
 import json
 import logging
-from multiprocessing import Pipe, Process
+import multiprocessing
 from multiprocessing.connection import Connection
 from pathlib import Path
 from typing import Any, Callable, Tuple
@@ -52,16 +52,23 @@ def __init__(
             env (ManipulatorBaseEnv): The environment in which the policy will be applied.
             overrides (dict | None): Optional overrides for the policy configuration.
         """
-        self.parent_conn, self.child_conn = Pipe()
         self.env = env
         self.overrides = overrides if overrides is not None else {}
 
-        self.inf_proc = Process(
+        ctx = multiprocessing.get_context("spawn")
+        self.parent_conn, self.child_conn = ctx.Pipe()
+
+        # Extract env data before spawning (env may not be picklable — ROS2 handles)
+        observation_space = env.observation_space
+        env_metadata = env.get_metadata()
+
+        self.inf_proc = ctx.Process(
             target=inference_worker,
             kwargs={
                 "conn": self.child_conn,
                 "pretrained_path": pretrained_path,
-                "env": env,
+                "observation_space": observation_space,
+                "env_metadata": env_metadata,
                 "overrides": self.overrides,
             },
             daemon=True,
@@ -114,15 +121,17 @@ def shutdown(self):
 def inference_worker(
     conn: Connection,
     pretrained_path: str,
-    env: ManipulatorBaseEnv,
+    observation_space,
+    env_metadata: dict,
     overrides: dict | None = None,
 ):  # noqa: ANN001
     """Policy inference process: loads policy on GPU, receives observations via conn, returns actions, and exits on None.
 
     Args:
         conn (Connection): The connection to the parent process for sending and receiving data.
         pretrained_path (str): Path to the pretrained policy model.
-        env (ManipulatorBaseEnv): The environment in which the policy will be applied.
+        observation_space: The environment's observation space (pre-extracted for spawn compatibility).
+        env_metadata (dict): The environment metadata (pre-extracted for spawn compatibility).
         overrides (dict | None): Optional overrides for the policy configuration.
     """
     setup_logging()
@@ -145,7 +154,7 @@ def inference_worker(
 
         train_config = TrainPipelineConfig.from_pretrained(pretrained_path)
 
-        _check_dataset_metadata(train_config, env, logger)
+        _check_dataset_metadata(train_config, env_metadata, logger)
 
         logger.info("[Inference] Loaded training config.")
 
@@ -176,7 +185,7 @@ def inference_worker(
         if USE_LEROBOT_PROCESSORS:
             preprocessor, postprocessor = make_pre_post_processors(policy_cfg=policy.config, pretrained_path=pretrained_path)
 
-        warmup_obs_raw = env.observation_space.sample()
+        warmup_obs_raw = observation_space.sample()
         warmup_obs_raw["observation.state"] = concatenate_state_features(warmup_obs_raw)
         warmup_obs = numpy_obs_to_torch(warmup_obs_raw)
         if USE_LEROBOT_PROCESSORS:
@@ -238,15 +247,15 @@ def inference_worker(
 
 def _check_dataset_metadata(
     train_config: TrainPipelineConfig,
-    env: ManipulatorBaseEnv,
+    env_metadata: dict,
     logger: logging.Logger,
     keys_to_skip: list[str] | None = None,
 ):
     """Check if the dataset metadata matches the environment configuration.
 
     Args:
         train_config (TrainPipelineConfig): The training pipeline configuration.
-        env (ManipulatorBaseEnv): The environment to compare against.
+        env_metadata (dict): The environment metadata dict to compare against.
         logger (logging.Logger): Logger for logging information.
         keys_to_skip (list[str] | None): List of metadata keys to skip during comparison.
     """
@@ -272,7 +281,6 @@ def _warn_if_missing(key: str):
             logger.info(
                 "[Inference] Found crisp_meta.json in dataset, comparing environment and policy configs..."
             )
-            env_metadata = env.get_metadata()
             with open(path_to_metadata, "r") as f:
                 dataset_metadata = json.load(f)
             for key, value in dataset_metadata.items():

crisp_gym/scripts/test_lerobot_policy.py

@@ -8,6 +8,7 @@
 import argparse
 import json
 import logging
+import multiprocessing
 import os
 import time
 
@@ -19,6 +20,35 @@
 logger = logging.getLogger(__name__)
 
 
+def _spawn_worker(conn, policy_path: str, device: str):
+    """Top-level function for spawn test: loads policy on the given device and reports success/failure.
+
+    Must be a top-level function (not a lambda/closure) to be picklable for spawn.
+    """
+    try:
+        import json as _json
+        from lerobot.policies.factory import get_policy_class
+
+        if os.path.isdir(policy_path):
+            config_path = os.path.join(policy_path, "config.json")
+            with open(config_path) as f:
+                cfg = _json.load(f)
+        else:
+            from huggingface_hub import hf_hub_download
+            with open(hf_hub_download(policy_path, "config.json")) as f:
+                cfg = _json.load(f)
+
+        policy_type = cfg.get("type") or cfg.get("policy_type")
+        policy_cls = get_policy_class(policy_type)
+        policy = policy_cls.from_pretrained(policy_path)
+        policy.to(device).eval()
+        conn.send({"ok": True, "device": device})
+    except Exception as e:
+        conn.send({"ok": False, "error": str(e)})
+    finally:
+        conn.close()
+
+
 def make_dummy_obs(policy, device: str) -> dict[str, torch.Tensor]:
     """Build a random observation batch (size 1) matching policy.config.input_features.
 
@@ -73,6 +103,8 @@ def main():
     parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
     parser.add_argument("--n-warmup", type=int, default=5, help="Warmup iterations.")
     parser.add_argument("--n-bench", type=int, default=10, help="Benchmark iterations.")
+    parser.add_argument("--test-spawn", action=argparse.BooleanOptionalAction, default=True,
+                        help="After inference test, spawn a subprocess to verify no CUDA fork error.")
     args = parser.parse_args()
 
     print(f"\n{'='*60}")
@@ -155,6 +187,26 @@ def main():
     print(f"  min / max    : {min(times)*1000:.2f} / {max(times)*1000:.2f} ms")
     print("\nPASS: policy loaded and inference completed successfully.")
 
+    # ── Step 5: Spawn subprocess test (CUDA fork safety) ──────────
+    if args.test_spawn:
+        print(f"\n[5/5] Spawn subprocess test (CUDA already initialized in main process)...")
+        ctx = multiprocessing.get_context("spawn")
+        parent_conn, child_conn = ctx.Pipe(duplex=False)
+        proc = ctx.Process(
+            target=_spawn_worker,
+            args=(child_conn, args.policy_path, args.device),
+            daemon=True,
+        )
+        proc.start()
+        child_conn.close()  # close child end in parent
+        result = parent_conn.recv()
+        proc.join(timeout=300)
+        if result.get("ok"):
+            print(f"PASS: spawn worker loaded policy on {result['device']} without fork error")
+        else:
+            print(f"FAIL: spawn worker reported error: {result.get('error')}")
+            raise SystemExit(1)
+
 
 if __name__ == "__main__":
     logging.basicConfig(level=logging.INFO)