(feat): refactor webshop env.

Author: PanAndy

examples/qwen2.5-7B-agentic_megatron/agentic_val_webshop.yaml

@@ -10,7 +10,7 @@ hydra:
     dir: .
   output_subdir: null
 
-exp_name: "agentic_pipeline"
+exp_name: "agentic_pipeline_webshop"
 seed: 42
 logging_dir: ./output/logs
 output_dir: ./output
@@ -29,9 +29,21 @@ system_envs:
 #    - roll
 #    - baseline
 
+#track_with: swanlab
+#tracker_kwargs:
+#  login_kwargs:
+#    api_key: your_api_key
+#  project: roll-agentic
+#  logdir: debug
+#  experiment_name: ${exp_name}
+#  tags:
+#    - roll
+#    - agentic
+#    - debug
+
 track_with: tensorboard
 tracker_kwargs:
-  log_dir: /data/oss_bucket_0/yali/llm/tensorboard/roll_exp/agentic_sokoban
+  log_dir: /data/oss_bucket_0/yali/llm/tensorboard/roll_exp/agentic_webshop
 
 num_gpus_per_node: 8
 
@@ -48,8 +60,8 @@ sequence_length: 8192
 reward_clip: 20
 advantage_clip: 0.2 # 0.1-0.3
 ppo_epochs: 1
-adv_estimator: "reinforce"
-pg_clip: 0.1
+adv_estimator: "grpo"
+#pg_clip: 0.1
 max_grad_norm: 1.0
 #dual_clip_loss: True
 init_kl_coef: 0.0
@@ -69,7 +81,7 @@ actor_train:
     learning_rate: 1.0e-6
     weight_decay: 0
     per_device_train_batch_size: 1
-    gradient_accumulation_steps: 4
+    gradient_accumulation_steps: 8
     warmup_steps: 10
   data_args:
     template: qwen2_5
@@ -123,51 +135,42 @@ reference:
   device_mapping: list(range(0,8))
   infer_batch_size: 1
 
-enable_response_mask: True
-action_sep: "||"
-use_turn_scores: False # important to GAE when applying token-level rewards to token-level advantages. If False, will take the sum of scores as the reward for the last turn.
-enable_think: False # False -> no think RL
-max_actions_per_traj: 20
 reward_normalization:
-  grouping: tags # 可以tags(env_type)/traj_group_id(group)/batch(rollout_batch)... group_by计算reward/adv
-  method: identity # asym_clip / identity / mean_std
-
-custom_envs:
-  WebShopEnv:
-    env_type: webshop
-    max_actions_per_traj:  ${max_actions_per_traj}
-    max_steps_per_traj: ${max_actions_per_traj}
-    env_instruction: |
-      You are web shopping.
-      I will give you instructions about what to do.
-      You have to follow the instructions.
-      Every round I will give you an observation and a list of available actions, you have to respond an action based on the state and instruction.
-      You can use search action if search is available.
-      You can click one of the buttons in clickables.
-      An action should be of the following structure:
-      search[keywords]
-      click[value]
-      If the action is not valid, perform nothing.
-      Keywords in search are up to you, but the value in click must be a value in the list of available actions.
-      Remember that your keywords in search should be carefully designed.
-      Your response should use the following format:
-
-      Thought: I think ...
-      Action: click[something]
-    max_tokens: 1024
-    env_config:
-      observation_mode: text
+  grouping: traj_group_id # 可以tags(env_type)/traj_group_id(group)/batch(rollout_batch)... group_by计算reward/adv
+  method: mean_std # asym_clip / identity / mean_std
 
 train_env_manager:
   format_penalty: -0.05
-  num_env_groups: 64
-  group_size: 1
+  num_env_groups: 8
+  group_size: 8
+  max_env_num_per_worker: 4
   tags: [WebShopEnv]
-  num_groups_partition: [64] # If not set, all env names divide nums equally. Under the same group, the env config and env seed (prompt) are equal in each generation
+  num_groups_partition: [8] # If not set, all env names divide nums equally. Under the same group, the env config and env seed (prompt) are equal in each generation
 
 val_env_manager:
   num_env_groups: 64
   group_size: 1 # should be set to 1 because val temperature is set to 0 and same prompt leads to same output
+  max_env_num_per_worker: 4
   tags: [WebShopEnv]
   num_groups_partition: [64] # TODO: If not set, all env names divide nums equally. Under the same group, the env config and env seed (prompt) are equal in each generation
 
+max_tokens_per_step: 128
+max_actions_per_traj: 20
+action_pattern: <answer>(.*?)</answer>
+think_action_pattern: <think>(.*?)</think>\s*<answer>(.*?)</answer>
+user_prompt_no_think_format: <answer> [your answer] </answer>
+user_prompt_think_format: <think> [Your thoughts] </think> <answer> [your answer] </answer>
+added_text_no_think: <answer>
+added_text_think: <think>
+
+custom_envs:
+  WebShopEnv:
+    env_type: webshop
+    max_tokens_per_step: ${max_tokens_per_step}
+    user_prompt_format: ${user_prompt_no_think_format}
+    added_text: ${added_text_no_think}
+    env_manager_cls: roll.pipeline.agentic.env_manager.traj_env_manager.TrajEnvManager
+    use_thread_lock: true
+    env_config:
+      observation_mode: text
+      max_steps: ${max_actions_per_traj}

examples/qwen2.5-7B-agentic_megatron/agentic_val_webshop_async.yaml

@@ -0,0 +1,177 @@
+defaults:
+  - ../config/envs@_here_
+  - ../config/deepspeed_zero@_here_
+  - ../config/deepspeed_zero2@_here_
+  - ../config/deepspeed_zero3@_here_
+  - ../config/deepspeed_zero3_cpuoffload@_here_
+
+hydra:
+  run:
+    dir: .
+  output_subdir: null
+
+exp_name: "agentic_pipeline_webshop_async"
+seed: 42
+logging_dir: ./output/logs
+output_dir: ./output
+render_save_dir: ./output/render
+system_envs:
+  USE_MODELSCOPE: '1'
+
+#track_with: wandb
+#tracker_kwargs:
+#  api_key:
+#  project: roll-agentic
+#  name: ${exp_name}_webshop
+#  notes: "agentic_pipeline"
+#  tags:
+#    - agentic
+#    - roll
+#    - baseline
+
+#track_with: swanlab
+#tracker_kwargs:
+#  login_kwargs:
+#    api_key: your_api_key
+#  project: roll-agentic
+#  logdir: debug
+#  experiment_name: ${exp_name}
+#  tags:
+#    - roll
+#    - agentic
+#    - debug
+
+track_with: tensorboard
+tracker_kwargs:
+  log_dir: /data/oss_bucket_0/yali/llm/tensorboard/roll_exp/agentic_webshop
+
+num_gpus_per_node: 8
+
+max_steps: 1024
+save_steps: 10000
+logging_steps: 1
+eval_steps: 10
+resume_from_checkpoint: false
+
+async_generation_ratio: 1
+
+rollout_batch_size: 64
+val_batch_size: 64
+sequence_length: 8192
+
+reward_clip: 20
+advantage_clip: 0.2 # 0.1-0.3
+ppo_epochs: 1
+adv_estimator: "grpo"
+#pg_clip: 0.1
+max_grad_norm: 1.0
+#dual_clip_loss: True
+init_kl_coef: 0.0
+whiten_advantages: true
+entropy_loss_coef: 0
+
+pretrain: Qwen/Qwen2.5-7B-Instruct
+reward_pretrain: Qwen/Qwen2.5-7B-Instruct
+
+actor_train:
+  model_args:
+    attn_implementation: fa2
+    disable_gradient_checkpointing: false
+    dtype: bf16
+    model_type: ~
+  training_args:
+    learning_rate: 1.0e-6
+    weight_decay: 0
+    per_device_train_batch_size: 1
+    gradient_accumulation_steps: 16
+    warmup_steps: 10
+  data_args:
+    template: qwen2_5
+  strategy_args:
+    strategy_name: megatron_train
+    strategy_config:
+      tensor_model_parallel_size: 1
+      context_parallel_size: 1
+      pipeline_model_parallel_size: 1
+      expert_model_parallel_size: 1
+      use_distributed_optimizer: true
+      recompute_granularity: full
+      max_grad_norm: ${max_grad_norm}
+  device_mapping: list(range(0,4))
+  infer_batch_size: 1
+
+actor_infer:
+  model_args:
+    disable_gradient_checkpointing: true
+    dtype: bf16
+  generating_args:
+    max_new_tokens: 1024 # single-turn response length
+    top_p: 0.99
+    top_k: 100
+    num_beams: 1
+    temperature: 0.99
+    num_return_sequences: 1
+  data_args:
+    template: qwen2_5
+  strategy_args:
+    strategy_name: vllm
+    strategy_config:
+      gpu_memory_utilization: 0.8
+      block_size: 16
+      load_format: auto
+  device_mapping: list(range(4,8))
+  infer_batch_size: 1
+
+reference:
+  model_args:
+    attn_implementation: fa2
+    disable_gradient_checkpointing: true
+    dtype: bf16
+    model_type: ~
+  data_args:
+    template: qwen2_5
+  strategy_args:
+    strategy_name: hf_infer
+    strategy_config: ~
+  device_mapping: list(range(0,4))
+  infer_batch_size: 1
+
+reward_normalization:
+  grouping: traj_group_id # 可以tags(env_type)/traj_group_id(group)/batch(rollout_batch)... group_by计算reward/adv
+  method: mean_std # asym_clip / identity / mean_std
+
+train_env_manager:
+  format_penalty: -0.05
+  num_env_groups: 8
+  group_size: 8
+  max_env_num_per_worker: 4
+  tags: [WebShopEnv]
+  num_groups_partition: [8] # If not set, all env names divide nums equally. Under the same group, the env config and env seed (prompt) are equal in each generation
+
+val_env_manager:
+  num_env_groups: 64
+  group_size: 1 # should be set to 1 because val temperature is set to 0 and same prompt leads to same output
+  max_env_num_per_worker: 4
+  tags: [WebShopEnv]
+  num_groups_partition: [64] # TODO: If not set, all env names divide nums equally. Under the same group, the env config and env seed (prompt) are equal in each generation
+
+max_tokens_per_step: 128
+max_actions_per_traj: 20
+action_pattern: <answer>(.*?)</answer>
+think_action_pattern: <think>(.*?)</think>\s*<answer>(.*?)</answer>
+user_prompt_no_think_format: <answer> [your answer] </answer>
+user_prompt_think_format: <think> [Your thoughts] </think> <answer> [your answer] </answer>
+added_text_no_think: <answer>
+added_text_think: <think>
+
+custom_envs:
+  WebShopEnv:
+    env_type: webshop
+    max_tokens_per_step: ${max_tokens_per_step}
+    user_prompt_format: ${user_prompt_no_think_format}
+    added_text: ${added_text_no_think}
+    env_manager_cls: roll.pipeline.agentic.env_manager.traj_env_manager.TrajEnvManager
+    use_thread_lock: true
+    env_config:
+      observation_mode: text
+      max_steps: ${max_actions_per_traj}

requirements_common.txt

@@ -23,6 +23,7 @@ codetiming
 more_itertools
 
 wandb
+swanlab
 
 math-verify
 openai

roll/agentic/env/parse_action_utils.py

@@ -15,16 +15,20 @@ def default_parser_action_func(text, action_pattern, action_lookup, special_toke
         think_content, action_content = "", match.group(1)
     else:
         think_content, action_content = match.group(1), match.group(2)
-    action_content_strip = action_content
-    think_content_strip = think_content
-    for special_token in special_token_list:
-        action_content_strip = action_content.replace(special_token, "").strip()
-        think_content_strip = think_content.replace(special_token, "").strip()
+    action_content_strip = action_content.strip()
+    think_content_strip = think_content.strip()
 
-    rev_action_lookup = {v.lower(): k for k, v in action_lookup.items()}
+    if special_token_list is not None:
+        for special_token in special_token_list:
+            action_content_strip = action_content.replace(special_token, "").strip()
+            think_content_strip = think_content.replace(special_token, "").strip()
 
-    if action_content_strip.lower() in rev_action_lookup:
-        action = rev_action_lookup[action_content_strip.lower()]
+    action = action_content_strip
+    if action_lookup is not None:
+        action = None
+        rev_action_lookup = {v.lower(): k for k, v in action_lookup.items()}
+        if action_content_strip.lower() in rev_action_lookup:
+            action = rev_action_lookup[action_content_strip.lower()]
 
     action_info = {
         "action": action,

roll/agentic/env/webshop/config.py

@@ -5,11 +5,13 @@
 from webshop_minimal import init_basedir
 from webshop_minimal.utils import DEFAULT_FILE_PATH
 
+from roll.agentic.env.base import BaseEnvConfig
+
 init_basedir()  # init DEFAULT_FILE_PATH, hardcoded dataset to small
 
 
 @dataclass
-class WebShopEnvConfig:
+class WebShopEnvConfig(BaseEnvConfig):
     """Configuration for WebAgentText environment"""
 
     # dataset: str = field(default="small", metadata={"description": "Small or full dataset"})
@@ -34,3 +36,14 @@ class WebShopEnvConfig:
     show_attrs: bool = field(
         default=False, metadata={"description": "SimServer arg: Whether to show additional attributes"}
     )
+
+    max_steps: int = 10
+    env_instruction: str = ("You are web shopping. I will give you instructions about what to do. "
+                            "You have to follow the instructions. Every round I will give you an observation and "
+                            "a list of available actions, you have to respond an action based on the state and instruction. "
+                            "You can use search action if search is available. You can click one of the buttons in clickables. "
+                            "An action should be of the following structure: search[keywords] click[value] If the action is not valid, perform nothing. "
+                            "Keywords in search are up to you, but the value in click must be a value in the list of available actions. "
+                            "Remember that your keywords in search should be carefully designed. "
+                            "Your response should use the following format Thought: I think ... Action: click[something]")
+    action_pattern: str = r"<answer>(.*?)</answer>"