(feat): add sokoban sandbox env and scripts

Author: jingyushen

examples/agentic_demo/agent_val_sokoban_sandbox.yaml

@@ -0,0 +1,178 @@
+defaults:
+  - ../config/traj_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"
+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}_sokoban
+#  notes: "agentic_pipeline"
+#  tags:
+#    - agentic
+#    - roll
+#    - baseline
+
+track_with: tensorboard
+tracker_kwargs:
+  log_dir: /data/tensorboard/roll_exp/agentic_sandbox
+
+checkpoint_config:
+  type: file_system
+  output_dir: /data/cpfs_0/rl_examples/models/${exp_name}
+
+num_gpus_per_node: 8
+
+max_steps: 1024
+save_steps: 10000
+logging_steps: 1
+eval_steps: 10
+resume_from_checkpoint: false
+
+rollout_batch_size: 32
+val_batch_size: 32
+sequence_length: 4096
+
+advantage_clip: 0.2
+ppo_epochs: 1
+adv_estimator: "grpo"
+#pg_clip: 0.1
+#dual_clip_loss: True
+init_kl_coef: 0.0
+whiten_advantages: true
+entropy_loss_coef: 0
+max_grad_norm: 1.0
+
+pretrain: Qwen/Qwen2.5-0.5B-Instruct
+reward_pretrain: Qwen/Qwen2.5-0.5B-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
+    lr_scheduler_type: cosine
+  data_args:
+    template: qwen2_5
+  strategy_args:
+#    strategy_name: deepspeed_train
+#    strategy_config: ${deepspeed_zero3}
+    strategy_name: megatron_train
+    strategy_config:
+      tensor_model_parallel_size: 1
+      pipeline_model_parallel_size: 1
+      expert_model_parallel_size: 1
+      use_distributed_optimizer: true
+      recompute_granularity: full
+  device_mapping: list(range(1,2))
+  infer_batch_size: 1
+
+actor_infer:
+  model_args:
+    disable_gradient_checkpointing: true
+    dtype: bf16
+  generating_args:
+    max_new_tokens: 32 # 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(1,2))
+
+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(1,2))
+  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.15 # sokoban env penalty_for_step=-0.1
+  max_env_num_per_worker: 16
+  num_env_groups: 2
+  # under the same group, the env config and env seed are ensured to be equal
+  group_size: 8
+  tags: [SokobanSandbox]
+  num_groups_partition: [2] # 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:
+  max_env_num_per_worker: 1
+  num_env_groups: 2
+  group_size: 1 # should be set to 1 because val temperature is set to 0 and same prompt leads to same output
+  tags: [SimpleSokoban, LargerSokoban]
+  num_groups_partition: [1, 1] # 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
+
+# Here, you can override variables defined in the imported envs. max_tokens_per_step: 128 in custom_env.SimpleSokoban, here replaced by 64
+max_tokens_per_step: 64
+
+custom_envs:
+  SimpleSokoban:
+    ${custom_env.SimpleSokoban}
+  LargerSokoban:
+    ${custom_env.LargerSokoban}
+  SokobanDifferentGridVocab:
+    ${custom_env.SokobanDifferentGridVocab}
+  FrozenLake:
+    ${custom_env.FrozenLake}
+  FrozenLakeThink:
+    ${custom_env.FrozenLakeThink}
+  SokobanSandbox:    
+    env_type: sokoban_sandbox
+    env_config: 
+      base_url: 'http://localhost:8080' # change to your own service address if needed
+    max_steps: ${max_actions_per_traj}
+    max_tokens_per_step: ${max_tokens_per_step}
+    env_manager_cls: ${env_manager_cls}
+    use_thread_lock: true
+    agent_system_template: "You are an expert Sokoban-playing agent. Your only task is to output the next action like <answer>Right</answer>. Do not output any other text, reasoning, or explanations."
+    agent_template: |
+      Current State:
+      {observation}
+      Turn {turn_idx}:
+      {suffix}
+      Respond with one action from the available list.
+      Format: <answer>ACTION</answer>
+      Choose your action.

examples/agentic_demo/run_agentic_pipeline_sokoban_sandbox_multi_nodes.sh

@@ -0,0 +1,7 @@
+#!/bin/bash
+set +x
+
+CONFIG_PATH=$(basename $(dirname $0))
+export PYTHONPATH="$PWD:$PYTHONPATH"
+python examples/start_agentic_pipeline.py --config_path $CONFIG_PATH  --config_name agent_val_sokoban_sandbox
+

examples/agentic_demo/run_agentic_pipeline_sokoban_sandbox_single_node.sh

@@ -0,0 +1,21 @@
+#!/bin/bash
+set +x
+
+
+echo "Stopping existing services..."
+pkill -9 -f "admin --env local" 2>/dev/null || true
+lsof -ti:8080 | xargs kill -9 2>/dev/null || true
+sleep 2
+
+echo "Starting ROCK service..."
+cd /workspace/ROCK && nohup admin --env local > /tmp/rock_service.log 2>&1 &
+
+echo "Waiting for ROCK service to be ready..."
+sleep 5
+
+cd $PWD
+
+CONFIG_PATH=$(basename $(dirname $0))
+export PYTHONPATH="$PWD:$PYTHONPATH"
+python examples/start_agentic_pipeline.py --config_path $CONFIG_PATH  --config_name agent_val_sokoban_sandbox
+

roll/pipeline/agentic/env/__init__.py

@@ -11,6 +11,7 @@
 gem.register("roll_math", entry_point="roll.pipeline.agentic.env.gem.math_env:MathEnv")
 gem.register("roll_code", entry_point="roll.pipeline.agentic.env.gem.code_env:CodeEnv")
 gem.register("roll_qa", entry_point="roll.pipeline.agentic.env.gem.qa_env:QaEnv")
+gem.register("sokoban_sandbox", entry_point="roll.pipeline.agentic.env.sandbox:SokobanSandboxEnv")
 
 
 try:

roll/pipeline/agentic/env/sandbox/__init__.py

@@ -0,0 +1,3 @@
+from .sokoban_sandbox_env import SokobanSandboxEnv
+
+__all__ = ["SokobanSandboxEnv"]