Release CoD v2

.gitignore

@@ -101,3 +101,12 @@ wandb/
 
 # checkpoints
 checkpoints/
+
+# local experiment artifacts
+.codex_snapshots/
+tmp_vllm_stub/
+examples/research_cod/data/
+examples/data/optimal_control/
+examples/research_cod/pde_discovery/results/
+
+playground.ipynb

examples/learn_to_ask/data_prepare/message_splitter.py

@@ -10,27 +10,22 @@ def split_single_message_list(messages):
 
     Returns:
         list: List of rounds, where each round contains messages and remaining chat
+
+    Note: consecutive same-role messages are preserved as-is (RealMedConv is IM-style).
     """
     rounds = []
     round_number = 1
     i = 0
 
     while i < len(messages):
-        # Collect messages for this round
-        round_messages = []
-
-        # Add messages until we reach a user message
+        # Advance i past one [non-user]* [user]* block (one "turn")
         while i < len(messages) and messages[i].get("role") != "user":
-            round_messages.append(messages[i])
             i += 1
-
-        # Add user message(s) - if there are consecutive user messages,
-        # include all of them in this round
         while i < len(messages) and messages[i].get("role") == "user":
-            round_messages.append(messages[i])
             i += 1
 
-        # The remaining messages (if any) form the remaining_chat
+        # Use full prefix as round messages (paper Sec 3.1: C_{t-1} = (u_0, ..., u_{t-1}))
+        round_messages = messages[:i]
         remaining_messages = messages[i:]
         round_entry = {"round_number": round_number, "messages": round_messages}
 

examples/research_cod/README.md

@@ -10,8 +10,6 @@ CoD groups related tasks into a pack and generates a long rollout trajectory int
 The whole pack is trained end-to-end with RL, with fine-grained credit assignment rewarding context updates that make future tasks easier.
 A trained model's reward rises across pack positions, which is the signature of the elicited CoD meta-capability.
 
-> **The full implementation lives on the [`research/cod`](https://github.com/agentscope-ai/Trinity-RFT/tree/research/cod/examples/research_cod) branch.** Check out that branch to view or run the full code.
-
 <p align="center">
   <img src="assets/cod_overview.png" alt="CoD overview" width="760">
   <br><sub><em>Figure 1: a visualization of CoD-Deploy and CoD-Train (compared with standard task-by-task RL). Environments A and B are used for training; M is a new environment for deployment or evaluation. Each block is one rollout episode, for solving a task x_i (which may itself be a long-horizon multi-turn task) or for updating the agent's context z_i about the current environment.</em></sub>
@@ -40,7 +38,12 @@ Key metric: `reward_iterative_hint_e2e_taskset_{ts}_pos_{pos}`, the mean reward
 
 <p align="center">
   <img src="assets/cod_effect.png" alt="reward rising across pack positions" width="820">
-  <br><sub><em>Figure 2: the CoD effect. Reward rises across pack positions during training and at OOD evaluation, both in-domain (harder FrozenLake) and cross-domain (Alchemy, Terminal).</em></sub>
+  <br><sub><em>Figure 2: the CoD effect with Qwen3-8B. Reward rises across pack positions during training and at OOD evaluation, both in-domain (harder FrozenLake) and cross-domain (Alchemy, Terminal).</em></sub>
+</p>
+
+<p align="center">
+  <img src="assets/cod_qwen27b.png" alt="CoD-trained Qwen3.6-27B compared with Qwen3.8 models across three environments" width="1000">
+  <br><sub><em>CoD-trained Qwen3.6-27B compared with the base model, Qwen3.8-27B, and Qwen3.8-Max on PDE Discovery, Optimal Control, and Grid Navigation.</em></sub>
 </p>
 
 ---
@@ -55,16 +58,41 @@ In each environment, the tasks in a pack share something reusable: sometimes a h
 | Alchemy-Random | `cod_random_alchemy_workflow` | A hidden crafting recipe: which elements combine into which new element |
 | Terminal | `cod_terminal_workflow` | How commands and paths work and their pitfalls, and roughly where files live |
 | Learn2Ask | `cod_learn2ask_workflow` | When to keep asking vs. when to stop and give a diagnosis |
+| Optimal Control | `cod_optimalcontrol_workflow` | Learn the system's hidden dynamics from feedback and use them to reach new target states |
+| PDE Discovery | `trinity.common.workflows.connect_the_dots.pde_discovery.workflow.CoDPDEDiscoveryWorkflow` | Discover the unknown reaction term from sampled data and refine it across tasks |
+| Grid Navigation | `cod_grid_navigation_workflow` | Explore a shared cost map and use accumulated observations to choose lower-cost routes |
+
+Implementations live in [`trinity/common/workflows/connect_the_dots/`](../../trinity/common/workflows/connect_the_dots/).
 
-The workflow implementations live under [`trinity/common/workflows/connect_the_dots/`](https://github.com/agentscope-ai/Trinity-RFT/tree/research/cod/trinity/common/workflows/connect_the_dots) on the `research/cod` branch.
+---
+
+## Layout
+
+```
+examples/research_cod/
+├── get_*_data.py        # environment data generators
+├── exp_plan_final/      # main study
+│   ├── train/           # training configs
+│   └── bench/           # eval configs
+└── exp_plan_learn2ask/   # data prep for learn2ask
+
+trinity/common/workflows/connect_the_dots/   # CoD workflow implementation
+├── cod_workflow.py     # packing / iterative hint / task reward
+├── base_workflow.py    # AsyncCoDMultiStepWorkflow base class
+└── <env>/              # each environment has its own subdir
+    ├── workflow.py     # task rendering / scoring
+    └── prompts/        # system / user prompts
+
+trinity/algorithm/advantage_fn/cod_advantage.py   # reward-to-go credit assignment (CoDAdvantageFn)
+```
 
 ---
 
-## Run it
+## Quickstart
 
 ```bash
-git clone -b research/cod https://github.com/agentscope-ai/Trinity-RFT.git
-cd Trinity-RFT
+git clone <REPO_URL>
+cd <REPO_DIR>
 conda create -n trinity python=3.12 && conda activate trinity
 pip install -e ".[vllm,flash_attn]"
 pip install gymnasium jinja2 pandas
@@ -78,27 +106,127 @@ python examples/research_cod/get_frozen_lake_data.py --local_dir examples/resear
 python examples/research_cod/get_frozen_lake_data.py --local_dir examples/research_cod/data/frozen_lake_6767 \
     --train_size 50000 --test_size 4000 --map_min_size 6 --map_max_size 7 --tile_min_prob 0.6 --tile_max_prob 0.7
 # Alchemy-Random
-python examples/research_cod/get_alchemy_data.py --local_dir examples/research_cod/data/alchemy_random --train_size 50000 --test_size 4000 --seed 42
+python examples/research_cod/get_alchemy_data.py  --local_dir examples/research_cod/data/alchemy_random --train_size 50000 --test_size 4000 --seed 42
 # Terminal
 python examples/research_cod/get_terminal_data.py --local_dir examples/research_cod/data/terminal --train_size 50000 --test_size 4000 --seed 42 --composite_ratio 0.5
+# PDE Discovery
+python examples/research_cod/get_pde_discovery_data.py \
+    --local_dir examples/research_cod/data/pde_discovery_runtime_seed \
+    --train_size 50000 --test_size 32 --seed 42
+# Optimal Control
+python examples/research_cod/get_optimal_control_data.py \
+    --local_dir examples/research_cod/data/optimal_control \
+    --train_size 50000 --test_size 4000 --difficulty hard --train_seed 42 --test_seed 2024
+# Grid Navigation
+python examples/research_cod/get_grid_navigation_data.py \
+    --local_dir examples/research_cod/data/grid_navigation \
+    --train_size 50000 --test_size 4000 --seed 42
 ```
 
-**2. Train.**
+**2. Train CoD models.**
+Set `TRINITY_MODEL_PATH` to the local model directory and adjust `cluster` (node count / GPUs per node) in the YAML for your hardware. In the W&B project selected by the config, watch `rollout/reward_iterative_hint_e2e_taskset_0_pos_{pos}/mean` to compare rewards across pack positions.
 ```bash
 # FrozenLake-Obscure
 trinity run --config examples/research_cod/exp_plan_final/train/frozen_lake_obscure.yaml
 # Mixed (joint training on FrozenLake-Obscure + Alchemy-Random)
 trinity run --config examples/research_cod/exp_plan_final/train/mixed_flobs_alchran.yaml
+
+# Qwen3.6-27B
+pip install -e ".[qwen3_5]"
+export TRINITY_MODEL_PATH=/path/to/Qwen3.6-27B
+
+# PDE Discovery
+trinity run --config examples/research_cod/exp_plan_final/train/pde_discovery_cod_600steps_hard.yaml
+# Optimal Control
+trinity run --config examples/research_cod/exp_plan_final/train/optimal_control_improve.yaml
+# Grid Navigation
+trinity run --config examples/research_cod/exp_plan_final/train/grid_navigation_27b.yaml
 ```
 
-**3. Evaluate.** (per-checkpoint OOD generalization: harder in-domain + unseen cross-domain)
+Mixed OPD uses one teacher per domain, with tokenizers compatible with the student:
+
+```bash
+# Convert teacher checkpoints to Hugging Face format
+PDE_CKPT=/path/to/pde-run/global_step_75
+CONTROL_CKPT=/path/to/control-run/global_step_200
+GRID_CKPT=/path/to/grid-run/global_step_100
+for ckpt in "$PDE_CKPT" "$CONTROL_CKPT" "$GRID_CKPT"; do
+    trinity convert --checkpoint-dir "$ckpt" --base-model-dir "$TRINITY_MODEL_PATH"
+done
+
+# Mixed OPD: PDE + Optimal Control + Grid Navigation
+export TRINITY_PDE_TEACHER_MODEL_PATH="$PDE_CKPT/actor/huggingface"
+export TRINITY_OPTIMAL_CONTROL_TEACHER_MODEL_PATH="$CONTROL_CKPT/actor/huggingface"
+export TRINITY_GRID_NAVIGATION_TEACHER_MODEL_PATH="$GRID_CKPT/actor/huggingface"
+trinity run --config examples/research_cod/exp_plan_final/train/pde_control_grid_opd.yaml
+```
+
+**3. Evaluate.**
+Evaluate each saved checkpoint, measuring out-of-distribution generalization both in-domain (harder versions of the training environments) and cross-domain (environments unseen in training).
 ```bash
 # FrozenLake-Obscure ckpt → FrozenLake-hard (in-domain) + Alchemy-easy / Terminal (cross-domain)
 bash examples/research_cod/exp_plan_final/bench/run_eval.sh --train-tasks frozen_lake_obscure
 # Mixed ckpt → FrozenLake-hard + Alchemy-hard / Terminal
 bash examples/research_cod/exp_plan_final/bench/run_eval.sh --train-tasks mixed_flobs_alchran
 ```
 
+PDE Discovery, Optimal Control, and Grid Navigation share a checkpoint benchmark. Generate the PDE evaluation set; reuse the Control and Grid test sets above:
+
+```bash
+python examples/research_cod/get_pde_discovery_data.py \
+    --local_dir examples/research_cod/data/pde_discovery_eval_hard_stratified_4000_disjoint_testonly \
+    --train_size 1 --test_size 4000 --seed 20260902 \
+    --eval_pack_size 8 --eval_template_count 25 \
+    --eval_ground_truth_family physical_full_eval_4000.json --test_only
+```
+
+Set `EVAL_PROJECT`, `EVAL_GROUP`, and `EVAL_NAME` to the checkpoint run. `TRINITY_MODEL_PATH` points to the base model; each `global_step_*/actor/` must contain `model.safetensors`.
+
+```bash
+export TRINITY_CHECKPOINT_ROOT_DIR=/path/to/checkpoints
+
+# PDE RL checkpoints → all three environments
+EVAL_PROJECT=trinity-cod EVAL_GROUP=pde_discovery EVAL_NAME="your-pde-run" EVAL_TRAIN_DOMAIN=pde \
+    trinity run --config examples/research_cod/exp_plan_final/bench/eval_source_checkpoints_all_domains.yaml
+# Optimal Control RL checkpoints → all three environments
+EVAL_PROJECT=trinity-cod EVAL_GROUP=optimal_control EVAL_NAME="your-control-run" EVAL_TRAIN_DOMAIN=control \
+    trinity run --config examples/research_cod/exp_plan_final/bench/eval_source_checkpoints_all_domains.yaml
+# Grid Navigation RL checkpoints → all three environments
+EVAL_PROJECT=trinity-cod-final EVAL_GROUP=grid_navigation EVAL_NAME="your-grid-run" EVAL_TRAIN_DOMAIN=grid \
+    trinity run --config examples/research_cod/exp_plan_final/bench/eval_source_checkpoints_all_domains.yaml
+```
+
+Mixed OPD checkpoint evaluation (select a directory containing steps 5, 10, …, 60; step 0 evaluates the base model):
+
+```bash
+EVAL_PROJECT=trinity-cod EVAL_GROUP=mixed_multi_teacher_opd EVAL_NAME="your-mixed-opd-eval-run" \
+    trinity run --config examples/research_cod/exp_plan_final/bench/eval_mixed_opd_steps_0to60_all_domains.yaml
+```
+
+---
+
+## Key config knobs (`cod.cod_workflow_args`)
+
+| Field | Meaning |
+|---|---|
+| `activated_cod_methods` | CoD methods to enable; the main study uses `["iterative_hint_e2e"]` |
+| `hint_penalty_coef` / `length_penalty_coef` | Length penalty on hints / on a correct solution |
+| `task_pack_size` / `eval_task_pack_size` | Pack size for training / evaluation |
+
+---
+
+## Adding a CoD environment
+
+1. Subclass the CoD base workflow in `trinity/common/workflows/connect_the_dots/<env>/workflow.py` (see `frozen_lake/workflow_obscure.py` for a compact example).
+2. Register `"cod_<env>_workflow": "...workflow.CoD<Env>Workflow"` in the `default_mapping` dict of [`trinity/common/workflows/__init__.py`](../../trinity/common/workflows/__init__.py).
+3. Add a generator + config under `examples/research_cod/`, setting `default_workflow_type: 'cod_<env>_workflow'`.
+
+After subclassing `AsyncCoDMultiStepWorkflow` (`base_workflow.py`), implement:
+
+- `step_async(step_num)`: build the prompt for the task, call the model, apply its action / grade its answer, write the reward into `self.final_reward`, and return `(continue, experiences)`.
+- `_get_feedback()`: return the environment's feedback for this step (written into `exp.info["feedback"]`, used to generate the hint).
+- `max_step_num`: the max number of steps per task.
+
 ---
 
 ## Citation