Upload ai/environments/vec_env_adapter.py with huggingface_hub

ai/environments/vec_env_adapter.py

@@ -0,0 +1,191 @@
+import os
+
+import numpy as np
+from gymnasium import spaces
+from stable_baselines3.common.vec_env import VecEnv
+
+# RUST Engine Toggle
+USE_RUST_ENGINE = os.getenv("USE_RUST_ENGINE", "0") == "1"
+
+if USE_RUST_ENGINE:
+    print(" [VecEnvAdapter] RUST Engine ENABLED (USE_RUST_ENGINE=1)")
+    from ai.vec_env_rust import RustVectorEnv
+
+    # Wrapper to inject MCTS_SIMS from env
+    class VectorEnvAdapter(RustVectorEnv):
+        def __init__(self, num_envs, action_space=None, opp_mode=0, force_start_order=-1):
+            mcts_sims = int(os.getenv("MCTS_SIMS", "50"))
+            super().__init__(num_envs, action_space, opp_mode, force_start_order, mcts_sims)
+else:
+    # GPU Environment Toggle
+    USE_GPU_ENV = os.getenv("USE_GPU_ENV", "0") == "1" or os.getenv("GPU_ENV", "0") == "1"
+
+    if USE_GPU_ENV:
+        try:
+            from ai.vector_env_gpu import HAS_CUDA, VectorEnvGPU
+
+            if HAS_CUDA:
+                print(" [VecEnvAdapter] GPU Environment ENABLED (USE_GPU_ENV=1)")
+            else:
+                print(" [VecEnvAdapter] Warning: USE_GPU_ENV=1 but CUDA not available. Falling back to CPU.")
+                USE_GPU_ENV = False
+        except ImportError as e:
+            print(f" [VecEnvAdapter] Warning: Failed to import GPU env: {e}. Falling back to CPU.")
+            USE_GPU_ENV = False
+
+    if not USE_GPU_ENV:
+        from ai.environments.vector_env import VectorGameState
+
+    class VectorEnvAdapter(VecEnv):
+        """
+        Wraps the Numba-accelerated VectorGameState to be compatible with Stable-Baselines3.
+
+        When USE_GPU_ENV=1 is set, uses VectorEnvGPU for GPU-resident environments
+        with zero-copy observation transfer to PyTorch.
+        """
+
+        metadata = {"render_modes": ["rgb_array"]}
+
+        def __init__(self, num_envs, action_space=None, opp_mode=0, force_start_order=-1):
+            self.num_envs = num_envs
+            self.use_gpu = USE_GPU_ENV
+
+            # For Legacy Adapter: Read MCTS_SIMS env var or default
+            mcts_sims = int(os.getenv("MCTS_SIMS", "50"))
+
+            if self.use_gpu:
+                # GPU Env doesn't support MCTS yet, pass legacy args
+                self.game_state = VectorEnvGPU(num_envs, opp_mode=opp_mode, force_start_order=force_start_order)
+            else:
+                self.game_state = VectorGameState(num_envs, opp_mode=opp_mode, force_start_order=force_start_order)
+
+            # Use Dynamic Dimension from Engine (IMAX 8k, Standard 2k, or Compressed 512)
+            obs_dim = self.game_state.obs_dim
+            self.observation_space = spaces.Box(low=0, high=1, shape=(obs_dim,), dtype=np.float32)
+            if action_space is None:
+                # Check if game_state has defined action_space_dim (default 2000)
+                if hasattr(self.game_state, "action_space_dim"):
+                    action_dim = self.game_state.action_space_dim
+                else:
+                    # Fallback: The Engine always produces 2000-dim masks (Action IDs 0-1999)
+                    action_dim = 2000
+
+                action_space = spaces.Discrete(action_dim)
+
+            # Manually initialize VecEnv fields to bypass render_modes crash
+            self.action_space = action_space
+            self.actions = None
+            self.render_mode = None
+
+            # Track previous scores for delta-based rewards
+            self.prev_scores = np.zeros(num_envs, dtype=np.int32)
+            self.prev_turns = np.zeros(num_envs, dtype=np.int32)
+            # Pre-allocate empty infos list (reused when no envs done)
+            self._empty_infos = [{} for _ in range(num_envs)]
+
+        def reset(self):
+            """
+            Reset all environments.
+            """
+            self.game_state.reset()
+            self.prev_scores.fill(0)  # Reset score tracking
+            self.prev_turns.fill(0)  # Reset turn tracking
+
+            obs = self.game_state.get_observations()
+            # Convert CuPy to NumPy if using GPU (SB3 expects numpy)
+            if self.use_gpu:
+                try:
+                    import cupy as cp
+
+                    if isinstance(obs, cp.ndarray):
+                        obs = cp.asnumpy(obs)
+                except:
+                    pass
+            return obs
+
+        def step_async(self, actions):
+            """
+            Tell the generic VecEnv wrapper to hold these actions.
+            """
+            self.actions = actions
+
+        def step_wait(self):
+            """
+            Execute the actions on the Numba engine.
+            """
+            # Ensure actions are int32 for Numba (avoid copy if already correct type)
+            if self.actions.dtype != np.int32:
+                actions_int32 = self.actions.astype(np.int32)
+            else:
+                actions_int32 = self.actions
+
+            # Step the engine
+            obs, rewards, dones, infos = self.game_state.step(actions_int32)
+
+            # Convert CuPy arrays to NumPy if using GPU (SB3 expects numpy)
+            if self.use_gpu:
+                try:
+                    import cupy as cp
+
+                    if isinstance(obs, cp.ndarray):
+                        obs = cp.asnumpy(obs)
+                    if isinstance(rewards, cp.ndarray):
+                        rewards = cp.asnumpy(rewards)
+                    if isinstance(dones, cp.ndarray):
+                        dones = cp.asnumpy(dones)
+                except:
+                    pass
+
+            return obs, rewards, dones, infos
+
+        def close(self):
+            pass
+
+        def get_attr(self, attr_name, indices=None):
+            """
+            Return attribute from vectorized environments.
+            """
+            if attr_name == "action_masks":
+                # Return function reference or result? SB3 usually looks for method
+                pass
+            return [None] * self.num_envs
+
+        def set_attr(self, attr_name, value, indices=None):
+            pass
+
+        def env_method(self, method_name, *method_args, **method_kwargs):
+            """
+            Call instance methods of vectorized environments.
+            """
+            if method_name == "action_masks":
+                # Return list of masks for all envs
+                masks = self.game_state.get_action_masks()
+                if self.use_gpu:
+                    try:
+                        import cupy as cp
+
+                        if isinstance(masks, cp.ndarray):
+                            masks = cp.asnumpy(masks)
+                    except:
+                        pass
+                return [masks[i] for i in range(self.num_envs)]
+
+            return [None] * self.num_envs
+
+        def env_is_wrapped(self, wrapper_class, indices=None):
+            return [False] * self.num_envs
+
+        def action_masks(self):
+            """
+            Required for MaskablePPO. Returns (num_envs, action_space.n) boolean array.
+            """
+            masks = self.game_state.get_action_masks()
+            if self.use_gpu:
+                try:
+                    import cupy as cp
+
+                    if isinstance(masks, cp.ndarray):
+                        masks = cp.asnumpy(masks)
+                except:
+                    pass
+            return masks