accelerate env

agent_system/multi_turn_rollout/rollout_loop.py

@@ -328,11 +328,17 @@ class TrajectoryCollector:
         episode_lengths = np.zeros(batch_size, dtype=np.float32)
         episode_rewards = np.zeros(batch_size, dtype=np.float32)
         tool_callings = np.zeros(batch_size, dtype=np.float32)
+        import time as _time
+        _total_preprocess_time = 0.0
+        _total_infer_time = 0.0
+        _total_env_time = 0.0
         # Trajectory collection loop
         for _step in range(self.config.env.max_steps):
             active_masks = np.logical_not(is_done)
 
+            _t0 = _time.time()
             batch = self.preprocess_batch(gen_batch=gen_batch, obs=obs)
+            _total_preprocess_time += _time.time() - _t0
 
             batch_keys_to_pop = ["input_ids", "attention_mask", "position_ids"]
             non_tensor_batch_keys_to_pop = ["raw_prompt_ids"]
@@ -351,7 +357,9 @@ class TrajectoryCollector:
 
             # pad to be divisible by dp_size
             batch_input_padded, pad_size = pad_dataproto_to_divisor(batch_input, actor_rollout_wg.world_size)
+            _t0 = _time.time()
             batch_output_padded = actor_rollout_wg.generate_sequences(batch_input_padded)
+            _total_infer_time += _time.time() - _t0
             # # unpad
             batch_output = unpad_dataproto(batch_output_padded, pad_size=pad_size)
 
@@ -359,10 +367,12 @@ class TrajectoryCollector:
             batch.non_tensor_batch['traj_uid'] = traj_uid
 
             batch = batch.union(batch_output)
-            
+
             text_actions = self.tokenizer.batch_decode(batch.batch['responses'], skip_special_tokens=True)
-            
+
+            _t0 = _time.time()
             next_obs, rewards, dones, infos = envs.step(text_actions)
+            _total_env_time += _time.time() - _t0
 
             
             if len(rewards.shape) == 2:
@@ -411,8 +421,9 @@ class TrajectoryCollector:
                     episode_lengths=episode_lengths,
                     )
         
-        return total_batch_list, episode_rewards, episode_lengths, success, traj_uid, tool_callings
-    
+        rollout_timing = {"inference_s": _total_infer_time, "env_s": _total_env_time, "preprocess_s": _total_preprocess_time}
+        return total_batch_list, episode_rewards, episode_lengths, success, traj_uid, tool_callings, rollout_timing
+
     def dynamic_multi_turn_loop(
             self,
             gen_batch: DataProto, 
@@ -451,7 +462,7 @@ class TrajectoryCollector:
                 print(f"valid num={len(total_batch_list)} < target num={self.config.data.train_batch_size * self.config.env.rollout.n}. Keep generating... ({try_count}/{max_try_count})")
             try_count += 1
 
-            batch_list, episode_rewards, episode_lengths, success, traj_uid, tool_callings = self.vanilla_multi_turn_loop(
+            batch_list, episode_rewards, episode_lengths, success, traj_uid, tool_callings, _ = self.vanilla_multi_turn_loop(
                 gen_batch=gen_batch,
                 actor_rollout_wg=actor_rollout_wg,
                 envs=envs,
@@ -479,7 +490,7 @@ class TrajectoryCollector:
         total_traj_uid = np.concatenate(total_traj_uid, axis=0)
         total_tool_callings = np.concatenate(total_tool_callings, axis=0)
 
-        return total_batch_list, total_episode_rewards, total_episode_lengths, total_success, total_traj_uid, total_tool_callings
+        return total_batch_list, total_episode_rewards, total_episode_lengths, total_success, total_traj_uid, total_tool_callings, {}
 
     def multi_turn_loop(
             self,
@@ -506,15 +517,15 @@ class TrajectoryCollector:
         # Initial observations from the environment
         if self.config.algorithm.filter_groups.enable and is_train:
             # Dynamic Sampling (for DAPO and Dynamic GiGPO)
-            total_batch_list, total_episode_rewards, total_episode_lengths, total_success, total_traj_uid, totoal_tool_callings = \
+            total_batch_list, total_episode_rewards, total_episode_lengths, total_success, total_traj_uid, totoal_tool_callings, rollout_timing = \
                 self.dynamic_multi_turn_loop(
                 gen_batch=gen_batch,
                 actor_rollout_wg=actor_rollout_wg,
                 envs=envs,
             )
         else:
-            # Vanilla Sampling   
-            total_batch_list, total_episode_rewards, total_episode_lengths, total_success, total_traj_uid, totoal_tool_callings = \
+            # Vanilla Sampling
+            total_batch_list, total_episode_rewards, total_episode_lengths, total_success, total_traj_uid, totoal_tool_callings, rollout_timing = \
                 self.vanilla_multi_turn_loop(
                 gen_batch=gen_batch,
                 actor_rollout_wg=actor_rollout_wg,
@@ -524,7 +535,7 @@ class TrajectoryCollector:
         assert len(total_batch_list) == len(total_episode_lengths)
         assert len(total_batch_list) == len(total_traj_uid)
         assert len(total_batch_list) == len(totoal_tool_callings)
-        
+
 
         # Create trajectory data
         gen_batch_output: DataProto = self.gather_rollout_data(
@@ -535,5 +546,9 @@ class TrajectoryCollector:
             traj_uid=total_traj_uid,
             tool_callings=totoal_tool_callings,
         )
-        
+
+        if gen_batch_output.meta_info is None:
+            gen_batch_output.meta_info = {}
+        gen_batch_output.meta_info['rollout_timing'] = rollout_timing
+
         return gen_batch_output

poisonclaw/envs/browsergym_env.py

@@ -6,6 +6,9 @@ enabling GRPO/GiGPO/PPO training on web environments:
   - VisualWebArena:    requires VWA servers (Docker or remote)
   - WebArena:          requires WebArena servers
 
+Each BrowserGym env runs in its own Ray Actor process, so all envs
+step/reset in parallel (no GIL or Playwright thread-affinity issues).
+
 Config fields (under env.*):
   env_name        routing key, must contain "browsergym" (e.g. "browsergym-miniwob")
   gym_id          single BrowserGym task ID  (mutually exclusive with task_list)
@@ -16,6 +19,7 @@ Config fields (under env.*):
   seed            base random seed  (default: 42)
   viewport_width  screenshot width  (default: 1280)
   viewport_height screenshot height (default: 720)
+  pre_observation_delay  seconds to wait before obs extraction (default: 0.5)
 """
 
 from __future__ import annotations
@@ -26,6 +30,7 @@ from collections import defaultdict
 from typing import Optional
 
 import numpy as np
+import ray
 
 from agent_system.environments.base import EnvironmentManagerBase
 
@@ -33,20 +38,74 @@ logger = logging.getLogger(__name__)
 
 # ── Action regex patterns (coordinate-based, matching VLM output format) ─────
 _RE_ACTION_TAG = re.compile(r"<action>(.*?)</action>", re.DOTALL | re.IGNORECASE)
-_RE_CLICK      = re.compile(r"click\(\s*(\d+(?:\.\d+)?)\s*,\s*(\d+(?:\.\d+)?)\s*\)")
-_RE_TYPE       = re.compile(r"type\((.+?)\)", re.DOTALL)
-_RE_PRESS      = re.compile(r"press\((.+?)\)")
-_RE_NAVIGATE   = re.compile(r"(?:navigate|goto)\((.+?)\)", re.IGNORECASE)
-_RE_SCROLL     = re.compile(r"scroll\(\s*(\d+)\s*,\s*(\d+)\s*,\s*(.+?)\s*\)")
+_RE_CLICK      = re.compile(r"click\s*\(\s*(\d+(?:\.\d+)?)\s*,\s*(\d+(?:\.\d+)?)\s*\)", re.IGNORECASE)
+_RE_TYPE       = re.compile(r"type\s*\(\s*(.+?)\s*\)", re.DOTALL)
+_RE_PRESS      = re.compile(r"press\s*\(\s*(.+?)\s*\)", re.IGNORECASE)
+_RE_NAVIGATE   = re.compile(r"(?:navigate|goto)\s*\(\s*(.+?)\s*\)", re.IGNORECASE)
+_RE_SCROLL     = re.compile(r"scroll\s*\(\s*(\d+)\s*,\s*(\d+)\s*,\s*(.+?)\s*\)")
+
+# Playwright key name mapping (VLM may output lowercase)
+_KEY_MAP = {
+    "enter": "Enter", "tab": "Tab", "escape": "Escape", "esc": "Escape",
+    "backspace": "Backspace", "delete": "Delete", "space": " ",
+    "arrowup": "ArrowUp", "arrowdown": "ArrowDown",
+    "arrowleft": "ArrowLeft", "arrowright": "ArrowRight",
+}
+
+
+# ── Ray Actor: one BrowserGym env per process ────────────────────────────────
+
+class BrowserGymWorker:
+    """Ray Actor wrapping a single BrowserGym gymnasium env.
+
+    Runs in its own process — no GIL or Playwright thread issues.
+    """
+
+    def __init__(self, task_id: str, action_mapping, pre_obs_delay: float = 0.5):
+        import gymnasium as gym
+        self._import_browsergym_namespaces()
+        self.env = gym.make(
+            task_id,
+            action_mapping=action_mapping,
+            pre_observation_delay=pre_obs_delay,
+        )
+
+    def step(self, action: str):
+        obs, reward, terminated, truncated, info = self.env.step(action)
+        return obs, float(reward), terminated, truncated, info
+
+    def reset(self, seed: int):
+        obs, info = self.env.reset(seed=seed)
+        return obs, info
+
+    def close(self):
+        try:
+            self.env.close()
+        except Exception:
+            pass
+
+    @staticmethod
+    def _import_browsergym_namespaces():
+        import importlib
+        for module in [
+            "browsergym.miniwob",
+            "browsergym.visualwebarena",
+            "browsergym.webarena",
+            "browsergym.workarena",
+            "browsergym.weblinx",
+        ]:
+            try:
+                importlib.import_module(module)
+            except ImportError:
+                pass
 
 
 class BrowserGymEnvManager(EnvironmentManagerBase):
     """Wraps BrowserGym gym environments for verl-agent training.
 
-    Each parallel slot gets a BrowserGym env instance.  On every episode
-    boundary the env is ``reset()``-ed in place (same task type, new seed).
-    For multi-task training supply ``config.env.task_list``; tasks are
-    assigned round-robin to the parallel slots.
+    Each parallel slot is a Ray Actor running a BrowserGym env in its own
+    process.  ``step()`` and ``reset()`` dispatch to all actors in parallel
+    via ``ray.get([actor.step.remote(...) for ...])``.
     """
 
     def __init__(self, config, split: str = "train") -> None:
@@ -71,26 +130,25 @@ class BrowserGymEnvManager(EnvironmentManagerBase):
             split, self.num_envs, self.task_ids,
         )
 
-        # Import BrowserGym namespace packages to trigger task registration
-        self._import_browsergym_namespaces()
-
-        # Build coordinate-based action mapping (VLM outputs pixel coords,
-        # not element bids).  Include 'coord' for mouse_click/keyboard_*,
-        # 'nav' for goto/go_back, and 'chat' for send_msg_to_user.
+        # Build coordinate-based action mapping
         from browsergym.core.action.highlevel import HighLevelActionSet
         self._action_set = HighLevelActionSet(subsets=["coord", "nav"])
 
-        # Create one gym.Env per slot (round-robin over task_ids)
-        import gymnasium as gym
-        self._gym_envs: list[gym.Env] = [
-            gym.make(
-                self.task_ids[i % len(self.task_ids)],
+        # Create Ray Actor workers (one browser per actor, fully parallel)
+        pre_obs_delay = float(getattr(config.env, "pre_observation_delay", 0.5))
+        resources = {"num_cpus": config.env.resources_per_worker.get("num_cpus", 0.5)}
+        WorkerActor = ray.remote(**resources)(BrowserGymWorker)
+
+        self._workers = [
+            WorkerActor.remote(
+                task_id=self.task_ids[i % len(self.task_ids)],
                 action_mapping=self._action_set.to_python_code,
+                pre_obs_delay=pre_obs_delay,
             )
             for i in range(self.num_envs)
         ]
 
-        # Per-env runtime state
+        # Per-env runtime state (kept on manager side for obs building)
         self._last_obs: list[Optional[dict]]  = [None] * self.num_envs
         self._steps:    list[int]             = [0]    * self.num_envs
         self._done:     list[bool]            = [True] * self.num_envs
@@ -106,28 +164,87 @@ class BrowserGymEnvManager(EnvironmentManagerBase):
     # ── EnvironmentManagerBase interface ─────────────────────────────────────
 
     def reset(self, kwargs=None) -> tuple[dict, list[dict]]:
+        futures = [
+            self._workers[i].reset.remote(self._seeds[i])
+            for i in range(self.num_envs)
+        ]
+        results = ray.get(futures)
+
         obs_list, info_list = [], []
-        for i in range(self.num_envs):
-            obs, info = self._reset_env(i)
+        for i, (obs, info) in enumerate(results):
+            self._last_obs[i] = obs
+            self._steps[i]    = 0
+            self._done[i]     = False
+            self._history[i]  = []
+            self._goals[i]    = self._extract_goal(obs)
+            self._seeds[i]   += self.num_envs
+            info.setdefault("won", False)
+            info["is_action_valid"] = np.array(True)
             obs_list.append(obs)
             info_list.append(info)
+
         return self._pack_obs(obs_list), info_list
 
     def step(
         self, text_actions: list[str]
     ) -> tuple[dict, np.ndarray, np.ndarray, list[dict]]:
-        obs_list, info_list = [], []
         rewards = np.zeros(self.num_envs, dtype=np.float32)
         dones   = np.zeros(self.num_envs, dtype=bool)
 
+        # Dispatch step or reset to each worker in parallel
+        futures = []
+        action_map = {}  # track which envs are stepping (vs resetting)
         for i, action_text in enumerate(text_actions):
             if self._done[i]:
-                obs, info = self._reset_env(i)
-                dones[i]  = False
+                seed = self._seeds[i]
+                self._seeds[i] += self.num_envs
+                futures.append(self._workers[i].reset.remote(seed))
+                action_map[i] = "reset"
+            else:
+                bg_action, is_valid = self._parse_action(action_text)
+                futures.append(self._workers[i].step.remote(bg_action))
+                action_map[i] = ("step", action_text, is_valid)
+
+        results = ray.get(futures)
+
+        obs_list, info_list = [], []
+        for i, result in enumerate(results):
+            if action_map[i] == "reset":
+                obs, info = result
+                self._last_obs[i] = obs
+                self._steps[i]    = 0
+                self._done[i]     = False
+                self._history[i]  = []
+                self._goals[i]    = self._extract_goal(obs)
+                info.setdefault("won", False)
+                info["is_action_valid"] = np.array(True)
+                dones[i] = False
             else:
-                obs, reward, done, info = self._step_env(i, action_text)
+                _, action_text, is_valid = action_map[i]
+                obs, reward, terminated, truncated, info = result
+                self._last_obs[i] = obs
+                self._steps[i]   += 1
+                done = terminated or truncated or (self._steps[i] >= self.max_steps)
+                self._done[i] = done
+                self._history[i].append(action_text)
                 rewards[i] = reward
                 dones[i]   = done
+                info["won"]              = bool(terminated and reward > 0)
+                info["is_action_valid"]  = np.array(is_valid)
+                info["last_action_error"] = obs.get("last_action_error", "")
+
+                # Debug: log first env's actions for the first few steps
+                if i == 0 and self._steps[i] <= 3:
+                    import sys
+                    err = obs.get("last_action_error", "")
+                    bg_action = action_text  # approximate for logging
+                    print(
+                        f"[DEBUG env0 step{self._steps[i]}] "
+                        f"vlm={action_text[:80]!r} "
+                        f"valid={is_valid} r={reward} term={terminated} err={err!r}",
+                        file=sys.stderr, flush=True,
+                    )
+
             obs_list.append(obs)
             info_list.append(info)
 
@@ -138,9 +255,13 @@ class BrowserGymEnvManager(EnvironmentManagerBase):
         return self._make_text_obs(obs_list)
 
     def close(self) -> None:
-        for env in self._gym_envs:
+        for worker in self._workers:
+            try:
+                ray.get(worker.close.remote())
+            except Exception:
+                pass
             try:
-                env.close()
+                ray.kill(worker)
             except Exception:
                 pass
 
@@ -155,71 +276,10 @@ class BrowserGymEnvManager(EnvironmentManagerBase):
         assert len(success["success_rate"]) == batch_size
         return {k: np.array(v) for k, v in success.items()}
 
-    # ── Internal helpers ──────────────────────────────────────────────────────
-
-    @staticmethod
-    def _import_browsergym_namespaces() -> None:
-        """Import BrowserGym sub-packages to register their tasks in gymnasium."""
-        _known = {
-            "miniwob":        "browsergym.miniwob",
-            "visualwebarena": "browsergym.visualwebarena",
-            "webarena":       "browsergym.webarena",
-            "workarena":      "browsergym.workarena",
-            "weblinx":        "browsergym.weblinx",
-        }
-        import importlib
-        for key, module in _known.items():
-            try:
-                importlib.import_module(module)
-            except ImportError:
-                pass  # optional package not installed
-
-    def _reset_env(self, idx: int) -> tuple[dict, dict]:
-        seed = self._seeds[idx]
-        self._seeds[idx] += self.num_envs  # advance seed for next episode
-
-        obs, info = self._gym_envs[idx].reset(seed=seed)
-        self._last_obs[idx] = obs
-        self._steps[idx]    = 0
-        self._done[idx]     = False
-        self._history[idx]  = []
-        self._goals[idx]    = self._extract_goal(obs)
-
-        info.setdefault("won", False)
-        info["is_action_valid"] = np.array(True)
-        return obs, info
-
-    def _step_env(
-        self, idx: int, action_text: str
-    ) -> tuple[dict, float, bool, dict]:
-        bg_action, is_valid = self._parse_action(action_text)
-
-        obs, reward, terminated, truncated, info = self._gym_envs[idx].step(bg_action)
-
-        self._last_obs[idx] = obs
-        self._steps[idx]   += 1
-        done = terminated or truncated or (self._steps[idx] >= self.max_steps)
-        self._done[idx] = done
-
-        self._history[idx].append(action_text)
-
-        info["won"]              = bool(terminated and reward > 0)
-        info["is_action_valid"]  = np.array(is_valid)
-        info["last_action_error"] = obs.get("last_action_error", "")
-        return obs, float(reward), done, info
-
     # ── Action parsing ────────────────────────────────────────────────────────
 
     def _parse_action(self, text: str) -> tuple[str, bool]:
-        """Convert VLM text output → BrowserGym coordinate-based action string.
-
-        BrowserGym uses its own action API (not raw Playwright calls):
-          - mouse_click(x, y)          for coordinate clicks
-          - keyboard_type(text)        for typing
-          - keyboard_press(key_comb)   for key presses
-          - goto(url)                  for navigation
-          - scroll(dx, dy)             for scrolling
-        """
+        """Convert VLM text output → BrowserGym coordinate-based action string."""
         # Unwrap optional <action> tags
         m = _RE_ACTION_TAG.search(text)
         text = m.group(1).strip() if m else text.strip()
@@ -240,6 +300,7 @@ class BrowserGymEnvManager(EnvironmentManagerBase):
         m = _RE_PRESS.search(text)
         if m:
             key = m.group(1).strip().strip("\"'")
+            key = _KEY_MAP.get(key.lower(), key)
             return f'keyboard_press("{key}")', True
 
         # navigate(url) / goto(url)
@@ -288,9 +349,12 @@ class BrowserGymEnvManager(EnvironmentManagerBase):
             if err:
                 parts.append(f"Last action error: {err}")
             parts.append(
-                "Next action (pick one format):\n"
-                "  click(x, y)  |  type(text)  |  press(key)"
-                "  |  navigate(url)  |  scroll(x, y, up/down)"
+                "Respond with exactly ONE action using the formats below. "
+                "Replace the placeholders with actual values.\n"
+                "  click(x, y) — click at pixel coordinates, e.g. click(120, 55)\n"
+                "  type(text) — type a string, e.g. type(hello world)\n"
+                "  press(key) — press a key, e.g. press(Enter)\n"
+                "Your response must start with the action, nothing else."
             )
             texts.append("\n\n".join(parts))
         return texts

scripts/run_browsergym_miniwob.sh

@@ -22,6 +22,9 @@ export MINIWOB_URL="http://localhost:${MINIWOB_PORT}/miniwob/"
 CONDA_ENV_LIB="$(python3 -c 'import sys, os; print(os.path.join(sys.prefix, "lib"))')"
 LOCAL_LIBS="/home/jovyan/project/verl-agent/local-libs/extracted/usr/lib/x86_64-linux-gnu"
 export LD_LIBRARY_PATH="${CONDA_ENV_LIB}:${LOCAL_LIBS}:${LD_LIBRARY_PATH:-}"
+export WANDB__SERVICE_WAIT=120
+WANDB_API_KEY=$(python3 -c "import wandb; print(wandb.api.api_key)" 2>/dev/null)
+export WANDB_API_KEY
 echo "[run_browsergym_miniwob] MINIWOB_URL=$MINIWOB_URL  (pid=$HTTP_PID)"
 
 # Cleanup HTTP server on exit
@@ -29,9 +32,9 @@ cleanup() { kill "$HTTP_PID" 2>/dev/null || true; }
 trap cleanup EXIT
 
 # ── Tunable knobs ─────────────────────────────────────────────────────────────
-train_data_size=8       # parallel train envs  (= train_batch_size)
-val_data_size=4         # parallel val envs
-group_size=4            # GRPO group size (rollout.n)
+train_data_size=4       # parallel train envs  (= train_batch_size)
+val_data_size=32         # parallel val envs
+group_size=8            # GRPO group size (rollout.n)
 HF_MODEL_ID="Qwen/Qwen2.5-VL-3B-Instruct"
 HF_CACHE_SNAPSHOT="$HOME/.cache/huggingface/hub/models--Qwen--Qwen2.5-VL-3B-Instruct/snapshots/66285546d2b821cf421d4f5eb2576359d3770cd3"
 LOCAL_MODEL_PATH="/tmp/Qwen2.5-VL-3B-Instruct"
@@ -92,14 +95,16 @@ fi
 # ── Training ──────────────────────────────────────────────────────────────────
 python3 -m verl.trainer.main_ppo \
     algorithm.adv_estimator=grpo \
-    algorithm.use_kl_in_reward=False \
-    algorithm.gamma=0.99 \
+    algorithm.use_kl_in_reward=True \
+    algorithm.kl_ctrl.type=fixed \
+    algorithm.kl_ctrl.kl_coef=0.001 \
+    algorithm.gamma=0.95 \
     \
     data.train_files="$HOME/data/verl-agent/visual/train.parquet" \
     data.val_files="$HOME/data/verl-agent/visual/test.parquet" \
     data.train_batch_size="$train_data_size" \
     data.val_batch_size="$val_data_size" \
-    data.max_prompt_length=2048 \
+    data.max_prompt_length=1024 \
     data.max_response_length=256 \
     data.filter_overlong_prompts=True \
     data.truncation=left \
@@ -115,7 +120,7 @@ python3 -m verl.trainer.main_ppo \
     actor_rollout_ref.model.enable_gradient_checkpointing=True \
     actor_rollout_ref.model.use_remove_padding=False \
     actor_rollout_ref.actor.strategy=fsdp \
-    actor_rollout_ref.actor.optim.lr=1e-5 \
+    actor_rollout_ref.actor.optim.lr=2e-5 \
     actor_rollout_ref.actor.ppo_mini_batch_size=8 \
     actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=4 \
     actor_rollout_ref.actor.use_kl_loss=False \
@@ -126,37 +131,39 @@ python3 -m verl.trainer.main_ppo \
     actor_rollout_ref.rollout.name="$ENGINE" \
     actor_rollout_ref.rollout.n=1 \
     actor_rollout_ref.rollout.tensor_model_parallel_size=1 \
-    actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \
-    actor_rollout_ref.rollout.enable_chunked_prefill=False \
+    actor_rollout_ref.rollout.gpu_memory_utilization=0.75 \
+    actor_rollout_ref.rollout.enable_chunked_prefill=True \
     actor_rollout_ref.rollout.enforce_eager=False \
     actor_rollout_ref.rollout.free_cache_engine=False \
     actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=4 \
     actor_rollout_ref.rollout.val_kwargs.temperature=0.0 \
     actor_rollout_ref.rollout.val_kwargs.do_sample=False \
-    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=4 \
+    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=8 \
     actor_rollout_ref.ref.fsdp_config.param_offload=False \
     "+ray_init.runtime_env.env_vars.LD_LIBRARY_PATH=${CONDA_ENV_LIB}:${LOCAL_LIBS}" \
+    "+ray_init.runtime_env.env_vars.WANDB_API_KEY=${WANDB_API_KEY}" \
     \
     \
     env.env_name=browsergym-miniwob \
     env.seed=42 \
-    env.max_steps=10 \
+    env.max_steps=7 \
     env.rollout.n="$group_size" \
-    env.resources_per_worker.num_cpus=0.1 \
+    env.resources_per_worker.num_cpus=0.5 \
     ++env.history_length=3 \
-    ++env.task_list="[browsergym/miniwob.click-button,browsergym/miniwob.click-dialog,browsergym/miniwob.click-link,browsergym/miniwob.click-checkboxes,browsergym/miniwob.enter-text]" \
+    ++env.task_list="[browsergym/miniwob.click-checkboxes,browsergym/miniwob.click-tab-2,browsergym/miniwob.email-inbox,browsergym/miniwob.search-engine,browsergym/miniwob.login-user,browsergym/miniwob.social-media,browsergym/miniwob.click-collapsible-2,browsergym/miniwob.book-flight]" \
+    ++env.pre_observation_delay=0.1 \
     ++env.viewport_width=332 \
     ++env.viewport_height=214 \
     \
     trainer.critic_warmup=0 \
-    trainer.logger="[console]" \
+    trainer.logger="[console,wandb]" \
     trainer.project_name="$project_name" \
     trainer.experiment_name="$experiment_name" \
     trainer.n_gpus_per_node=1 \
     trainer.nnodes=1 \
     trainer.save_freq=100 \
     trainer.test_freq=-1 \
-    trainer.total_epochs=200 \
-    trainer.val_before_train=True \
+    trainer.total_epochs=50 \
+    trainer.val_before_train=False \
     +ray_init.include_dashboard=False \
     "$@"

verl/trainer/ppo/ray_trainer.py

@@ -410,7 +410,8 @@ def _print_timing_breakdown(timing_raw: Dict[str, float], global_step: int):
     other = total - accounted
     if other > 0.5:
         parts.append(f"Other: {other:.1f}s ({other / total * 100:.0f}%)")
-    print(f"\n[Step {global_step}] Total: {total:.1f}s | {' | '.join(parts)}\n")
+    import sys
+    print(f"\n[Step {global_step}] Total: {total:.1f}s | {' | '.join(parts)}\n", file=sys.stderr, flush=True)
 
 
 class RayPPOTrainer:
@@ -1062,7 +1063,12 @@ class RayPPOTrainer:
         if self.val_reward_fn is not None and self.config.trainer.get("val_before_train", True):
             val_metrics = self._validate()
             assert val_metrics, f"{val_metrics=}"
-            pprint(f"Initial validation metrics: {val_metrics}")
+            import sys
+            print("\n" + "=" * 60, file=sys.stderr, flush=True)
+            print("INITIAL VALIDATION METRICS:", file=sys.stderr, flush=True)
+            for k, v in val_metrics.items():
+                print(f"  {k}: {v:.4f}", file=sys.stderr, flush=True)
+            print("=" * 60 + "\n", file=sys.stderr, flush=True)
             logger.log(data=val_metrics, step=self.global_steps)
             if self.config.trainer.get("val_only", False):
                 return
@@ -1115,6 +1121,14 @@ class RayPPOTrainer:
                                                                 envs=self.envs,
                                                                 is_train=True,
                                                                 )
+                    # Extract rollout sub-timing into timing_raw
+                    _rt = getattr(gen_batch_output, 'meta_info', None) or {}
+                    _rt = _rt.get('rollout_timing', {})
+                    if _rt:
+                        timing_raw['gen_preprocess'] = _rt.get('preprocess_s', 0.0)
+                        timing_raw['gen_inference'] = _rt.get('inference_s', 0.0)
+                        timing_raw['gen_env'] = _rt.get('env_s', 0.0)
+
                     if self.config.algorithm.adv_estimator == AdvantageEstimator.REMAX:
                         with _timer("gen_max", timing_raw):
                             gen_baseline_batch = deepcopy(gen_batch)

verl/utils/logger/aggregate_logger.py

@@ -40,7 +40,8 @@ class LocalLogger:
 
     def log(self, data, step):
         if self.print_to_console:
-            print(concat_dict_to_str(data, step=step), flush=True)
+            import sys
+            print(concat_dict_to_str(data, step=step), file=sys.stderr, flush=True)
 
 
 class DecoratorLoggerBase:

verl/utils/tracking.py

@@ -51,7 +51,12 @@ class Tracking:
         if "tracking" in default_backend or "wandb" in default_backend:
             import wandb
 
-            wandb.init(project=project_name, name=experiment_name, config=config)
+            wandb.init(
+                project=project_name,
+                name=experiment_name,
+                config=config,
+                settings=wandb.Settings(start_method="thread"),
+            )
             self.logger["wandb"] = wandb
 
         if "mlflow" in default_backend: