加入一个提前停止的机制 还有减少环境步中间步骤传递至cpu

.venv/lib/python3.10/site-packages/stable_worldmodel/solver/solver.py

@@ -1,6 +1,7 @@
 from typing import Any, Protocol, runtime_checkable
 
 import gymnasium as gym
+import numpy as np
 import torch
 
 
@@ -61,7 +62,12 @@ class Solver(Protocol):
         """Planning horizon length in timesteps."""
         ...
 
-    def solve(self, info_dict: dict, init_action: torch.Tensor | None = None) -> dict:
+    def solve(
+        self,
+        info_dict: dict,
+        init_action: torch.Tensor | None = None,
+        active_mask: torch.Tensor | np.ndarray | None = None,
+    ) -> dict:
         """Solve the planning optimization problem to find optimal actions.
 
         Args: