Init test

CLAUDE.md

@@ -0,0 +1,108 @@
+# CLAUDE.md
+
+This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository.
+
+## What this repo is
+
+RAGEN is a multi-turn RL training framework for LLM reasoning agents, built on top of [veRL](https://github.com/volcengine/verl). It implements **StarPO** (State-Thinking-Actions-Reward Policy Optimization): agents receive text observations, reason in `<think>` blocks, emit actions in `<answer>` blocks, receive environment rewards, and are trained with PPO/GRPO.
+
+`verl/` is a git submodule providing the underlying PPO/GRPO infrastructure, FSDP workers, and Ray orchestration. RAGEN adds the multi-turn rollout loop, environment registry, rollout filtering, and collapse diagnostics on top.
+
+## Key commands
+
+```bash
+# Installation (creates conda env + installs deps)
+bash scripts/setup_ragen.sh              # base
+bash scripts/setup_ragen.sh --with-search  # add SearchQA env (~87 GB index)
+
+# Training — config-name selects a task overlay on base.yaml
+python train.py --config-name _2_sokoban
+python train.py --config-name _2_sokoban trainer.experiment_name=my_run
+
+# Config dry-run (verify Hydra resolves without launching)
+python train.py --config-name _2_sokoban --cfg job
+
+# Tests (no test runner config; run individually)
+python -m pytest tests/
+
+# Zero-shot eval (vast.ai)
+python scripts/eval_zeroshot_pomdp.py --model Qwen/Qwen2.5-7B-Instruct
+```
+
+## Architecture
+
+Training entry point is `train.py` → `RayAgentTrainer` (`ragen/trainer/agent_trainer.py`). Each training step:
+
+1. **Rollout** — `LLMAgentProxy.rollout()` (`ragen/llm_agent/agent_proxy.py`) drives the multi-turn loop:
+   - `EnvStateManager` (`es_manager.py`) resets/steps environments in parallel groups
+   - `ContextManager` (`ctx_manager.py`) converts `(env_obs, history)` → LLM prompt and parses `<answer>` tags back into environment actions
+   - vLLM generates responses; loop repeats up to `agent_proxy.max_turn` turns
+
+2. **Filter** — `RolloutFilter` (`ragen/trainer/rollout_filter.py`) drops low-signal trajectories. Default strategy: `top_p` by reward variance (keeps top fraction with highest variance). Key config: `rollout_filter_strategy`, `rollout_filter_value`.
+
+3. **Advantage** — GAE (PPO) or GRPO depending on `algorithm.adv_estimator`.
+
+4. **Update** — veRL FSDP workers apply policy gradient.
+
+**Collapse detection** (`ragen/trainer/collapse_metrics.py`) runs every `collapse_detection.compute_freq` steps, computing mutual information I(X;Z) and conditional entropy H(Z|X) to distinguish template collapse from entropy collapse.
+
+## Config system
+
+Hydra-based. All configs live in `config/`. The hierarchy is:
+
+```
+config/base.yaml              ← global defaults (model, rollout, PPO params, env groups)
+config/envs.yaml              ← all environment definitions under custom_envs.*
+config/_2_sokoban.yaml        ← task overlay: sets experiment_name, inherits base
+```
+
+Task configs (`_1_bandit.yaml` … `_11_lights_out.yaml`) only override what differs from `base.yaml`. The active environment is set via:
+```yaml
+es_manager.train.env_configs.tags: ["CoordSokoban"]   # name from envs.yaml custom_envs.*
+```
+
+Override any field at the command line: `python train.py --config-name _2_sokoban model_path=Qwen/Qwen2.5-3B-Instruct`.
+
+Important batch-size constraint that `train.py` validates:
+```
+env_groups × group_size × rollout_filter_ratio ≥ ppo_mini_batch_size
+```
+
+## Environment system
+
+Environments are registered in `ragen/env/__init__.py` (`REGISTERED_ENVS` dict). Each env implements `BaseEnv` (`ragen/env/base.py`): `reset(seed) → obs_str` and `step(action) → (obs_str, reward, done, info)`.
+
+Defined in `config/envs.yaml` under `custom_envs.<Name>`:
+- `env_type`: key into `REGISTERED_ENVS`
+- `env_instruction`: the system prompt shown to the LLM
+- `max_actions_per_traj`: hard step limit
+- `env_config`: passed as kwargs to the env's config dataclass
+
+**Adding a new environment**: (1) create `ragen/env/<name>/` with a config dataclass and env class extending `BaseDiscreteActionEnv` or `BaseLanguageBasedEnv`; (2) register in `ragen/env/__init__.py`; (3) add an entry to `config/envs.yaml`.
+
+## POMDP Sokoban (this branch)
+
+Partial-observation Sokoban where the agent sees only a 3×3 window. Key additions:
+
+- `SokobanEnvConfig.partial_obs: bool` and `partial_obs_window: int` — in `ragen/env/sokoban/config.py`
+- `SokobanEnv._render_partial_obs()` — renders the local patch + `Position: (r, c)` footer — in `ragen/env/sokoban/env.py`
+- `POMDPSokoban` and `POMDPSokobanHard` env variants in `config/envs.yaml`
+- Training config: `config/_2_pomdp_sokoban.yaml`
+- vast.ai scripts: `scripts/setup_vast.sh`, `scripts/run_pomdp_sanity.sh`, `scripts/eval_zeroshot_pomdp.py`
+
+The agent is soft-scaffolded (via `env_instruction`) to maintain a progressive map in its `<think>` block. Memory compression is emergent from RL reward; no format penalty is applied.
+
+## Context window modes
+
+Controlled by `agent_proxy.context_window_mode`:
+- `full` — entire conversation history passed each turn
+- `limited_multi_turn` — last `max_context_window` turns only
+- `single_turn` — current observation only (no history)
+
+These are the axes of the memory baseline experiments.
+
+## Cluster setup notes
+
+**Quest (Northwestern HPC)**: use `set -eo pipefail` (not `-euo`); `source ~/.bashrc` before `conda activate ragen`; conda env at `/home/eiu4164/.conda/envs/ragen/`; HF cache at `/projects/p32139/hf_cache`.
+
+**vast.ai**: Python 3.12 env required (verl uses `X | Y` union syntax); install with `python -m pip install -e verl/ --ignore-requires-python`; `pkg_resources` missing on fresh Python 3.12 → `pip install setuptools`.

bandit.slurm

@@ -0,0 +1,67 @@
+#!/bin/bash
+#SBATCH --job-name=bandit
+#SBATCH --account=e52938
+#SBATCH --partition=gengpu
+#SBATCH --nodes=1
+#SBATCH --ntasks=1
+#SBATCH --gres=gpu:h100:1
+#SBATCH --cpus-per-task=16
+#SBATCH --mem=64G
+#SBATCH --time=6:00:00
+#SBATCH --output=/home/eiu4164/projects/RAGEN/logs/%x_%j.out
+#SBATCH --error=/home/eiu4164/projects/RAGEN/logs/%x_%j.err
+#SBATCH --mail-type=BEGIN,END,FAIL
+#SBATCH --mail-user=wenlanji2026@u.northwestern.edu
+
+set -eo pipefail
+
+mkdir -p /home/eiu4164/projects/RAGEN/logs
+cd /home/eiu4164/projects/RAGEN
+
+module load python-miniconda3/4.10.3
+source ~/.bashrc
+conda activate ragen
+
+export PYTHONNOUSERSITE=1
+export HF_HOME=/projects/p32139/hf_cache
+export OMP_NUM_THREADS=16
+
+export JOB_TMP="/tmp/j${SLURM_JOB_ID}"
+export TMPDIR="${JOB_TMP}/t"
+export RAY_TMPDIR="${JOB_TMP}/r"
+mkdir -p "${TMPDIR}" "${RAY_TMPDIR}"
+
+CKPT_ROOT="${RUN_STORAGE_ROOT:-/projects/p32139/gesture_data/ragen_runs}"
+export WANDB_DIR="${CKPT_ROOT}/wandb"
+mkdir -p "${CKPT_ROOT}" "${WANDB_DIR}"
+
+PY=$(which python)
+echo "Python:       ${PY}"
+echo "SLURM_JOB_ID: ${SLURM_JOB_ID}"
+echo "CUDA_VISIBLE_DEVICES: ${CUDA_VISIBLE_DEVICES:-unset}"
+
+EXP_NAME="bandit_grpo_0p5b"
+
+"${PY}" train.py \
+    --config-name _1_bandit \
+    model_path=Qwen/Qwen2.5-0.5B-Instruct \
+    trainer.n_gpus_per_node=1 \
+    micro_batch_size_per_gpu=8 \
+    ppo_mini_batch_size=64 \
+    es_manager.train.group_size=128 \
+    actor_rollout_ref.rollout.rollout_filter_value=0.10 \
+    actor_rollout_ref.actor.entropy_coeff=0.01 \
+    actor_rollout_ref.rollout.max_model_len=2048 \
+    actor_rollout_ref.rollout.response_length=500 \
+    actor_rollout_ref.rollout.gpu_memory_utilization=0.80 \
+    trainer.total_training_steps=200 \
+    trainer.save_freq=200 \
+    trainer.test_freq=20 \
+    trainer.resume_mode=disable \
+    algorithm.adv_estimator=grpo \
+    system.CUDA_VISIBLE_DEVICES="0" \
+    agent_proxy.context_window_mode=full \
+    agent_proxy.max_context_window=-1 \
+    trainer.experiment_name="${EXP_NAME}" \
+    "trainer.default_local_dir=${CKPT_ROOT}/checkpoints/\${trainer.project_name}/\${trainer.experiment_name}" \
+    > "/home/eiu4164/projects/RAGEN/logs/${EXP_NAME}.${SLURM_JOB_ID}.log" 2>&1

config/_11_lights_out.yaml

@@ -0,0 +1,25 @@
+defaults:
+  - base
+
+trainer:
+  experiment_name: lights-out
+
+actor_rollout_ref:
+  rollout:
+    response_length: 256
+    val_kwargs:
+      do_sample: True
+      temperature: 0.5
+
+agent_proxy:
+  max_turn: 8
+  max_actions_per_turn: 1
+  max_context_window: 5
+
+es_manager:
+  train:
+    env_configs:
+      tags: ["LightsOut"]
+  val:
+    env_configs:
+      tags: ["LightsOut"]

config/_2_pomdp_sokoban.yaml

@@ -0,0 +1,24 @@
+defaults:
+  - base
+
+hydra:
+  searchpath:
+    - pkg://verl.trainer.config
+
+trainer:
+  experiment_name: pomdp-sokoban-main
+
+agent_proxy:
+  max_turn: 30          # 30 steps/episode; full context peaks at ~7k tokens on 5×5
+  max_actions_per_turn: 1
+
+es_manager:
+  train:
+    env_configs:
+      tags: ["POMDPSokoban"]
+      n_groups: [8]
+  val:
+    env_groups: 32      # 512 × 30 steps is too expensive; 32 is enough
+    env_configs:
+      tags: ["POMDPSokoban"]
+      n_groups: [32]

config/base.yaml

@@ -14,7 +14,7 @@ log_prob_micro_batch_size_per_gpu: ${micro_batch_size_per_gpu}
 ppo_mini_batch_size: 32
 model_path: Qwen/Qwen2.5-3B-Instruct
 enable_response_mask: True # Enabling response mask could improve stability of rollout/old_log_prob, as P(st|history) are no longer calculated in loss here. See https://docs.google.com/document/d/1bg7obeiKTExuHHBl5uOiSpec5uLDZ2Tgvxy6li5pHX4/edit?usp=sharing for more details.
-grpo_advantage_length_weight: False # if you do not enable this and critic/advantage_estimator is GRPO, and the critic/advantages/mean is too low, then you can try enabling this to encourage reasoning and forbid collapse
+grpo_advantage_length_weight: True # if you do not enable this and critic/advantage_estimator is GRPO, and the critic/advantages/mean is too low, then you can try enabling this to encourage reasoning and forbid collapse
 
 lora:
   rank: 0
@@ -134,6 +134,7 @@ trainer:
 agent_proxy:
   context_window_mode: "full" # "full" | "limited_multi_turn" | "single_turn"
   max_context_window: -1 # k value: -1 for full history, 1 for no history (like without_history)
+  enable_summary_memory: false # if true, summarize turns outside the context window instead of dropping them
   batch_adjust_mode: copy # "copy" to duplicate samples, "delete" to remove samples when batch size is not divisible
   max_turn: 5
   action_sep: "||"
@@ -153,7 +154,12 @@ collapse_detection:
   num_samples: 64  # N or all
 
 es_manager:
-  format_penalty: -0.1
+  format_penalty: -0.5
+  step_penalty: 0.0
+  think_quality_penalty: 0.0
+  think_len_penalty: 0.0
+  think_len_limit: -1
+  think_coord_fmt_penalty: 0.0
   train:
     env_groups: 8
     # under the same group, the env config and env seed are ensured to be equal