diff --git a/.gitignore b/.gitignore index 27aec6eebfb..70e40523f52 100644 --- a/.gitignore +++ b/.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 diff --git a/examples/learn_to_ask/data_prepare/message_splitter.py b/examples/learn_to_ask/data_prepare/message_splitter.py index 06362b05b3b..b28017b7625 100644 --- a/examples/learn_to_ask/data_prepare/message_splitter.py +++ b/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} diff --git a/examples/research_cod/README.md b/examples/research_cod/README.md index 2f4219a4eea..b94c4a70dcc 100644 --- a/examples/research_cod/README.md +++ b/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. -

CoD overview
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. @@ -40,7 +38,12 @@ Key metric: `reward_iterative_hint_e2e_taskset_{ts}_pos_{pos}`, the mean reward

reward rising across pack positions -
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). +
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). +

+ +

+ CoD-trained Qwen3.6-27B compared with Qwen3.8 models across three environments +
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.

--- @@ -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 +└── / # 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 +cd conda create -n trinity python=3.12 && conda activate trinity pip install -e ".[vllm,flash_attn]" pip install gymnasium jinja2 pandas @@ -78,20 +106,63 @@ 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 @@ -99,6 +170,63 @@ bash examples/research_cod/exp_plan_final/bench/run_eval.sh --train-tasks frozen 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//workflow.py` (see `frozen_lake/workflow_obscure.py` for a compact example). +2. Register `"cod__workflow": "...workflow.CoDWorkflow"` 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__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 diff --git a/examples/research_cod/README_zh.md b/examples/research_cod/README_zh.md index 5743ba76fe7..0d122fb730d 100644 --- a/examples/research_cod/README_zh.md +++ b/examples/research_cod/README_zh.md @@ -10,8 +10,6 @@ CoD 把相关任务打成一个 pack（任务包），作为一条交替进行** 整个 pack 用 RL 端到端训练，细粒度的信用分配让“使后续任务更好解的 context 更新”获得奖励。 训练好的模型 reward 随 pack 位置递增，这正是 CoD 元能力被激发出来的标志。 -> **完整实现在 [`research/cod`](https://github.com/agentscope-ai/Trinity-RFT/tree/research/cod/examples/research_cod) 分支**，请切到该分支查看或运行完整代码。 -

CoD 方法总览
图 1：CoD-Deploy 与 CoD-Train 的示意（与标准的逐任务 RL 对比）。环境 A、B 用于训练，M 是部署或评测的新环境。每个 block 是一个 rollout 回合：求解任务 x_i（其本身可能是长程多轮任务），或更新 agent 对当前环境的 context z_i。 @@ -40,7 +38,12 @@ pack = [ task0, task1, task2, task3 ] # 同一 taskset 的一包相关任

reward 随 pack 位置递增 -
图 2：CoD 效应。reward 随 pack 位置上升，训练时如此，OOD 评测时也如此，包括 in-domain（更难的 FrozenLake）与 cross-domain（Alchemy、Terminal）。 +
图 2：Qwen3-8B 的 CoD 实验结果。reward 随 pack 位置上升，训练时如此，OOD 评测时也如此，包括 in-domain（更难的 FrozenLake）与 cross-domain（Alchemy、Terminal）。 +

+ +

+ CoD 训练后的 Qwen3.6-27B 与 Qwen3.8 系列在三个环境上的对比 +
CoD 训练后的 Qwen3.6-27B 与基础模型、Qwen3.8-27B、Qwen3.8-Max 在 PDE Discovery、Optimal Control 和 Grid Navigation 上的对比。

--- @@ -55,22 +58,47 @@ pack = [ task0, task1, task2, task3 ] # 同一 taskset 的一包相关任 | Alchemy-Random | `cod_random_alchemy_workflow` | 隐藏的合成配方，哪些元素能组合出哪个新元素 | | Terminal | `cod_terminal_workflow` | 命令、路径的用法与易踩的坑，以及文件大致在哪 | | Learn2Ask | `cod_learn2ask_workflow` | 何时继续追问、何时停下来给诊断 | +| Optimal Control | `cod_optimalcontrol_workflow` | 从交互反馈中了解系统的运动规律，更好地控制它到达不同目标 | +| PDE Discovery | `trinity.common.workflows.connect_the_dots.pde_discovery.workflow.CoDPDEDiscoveryWorkflow` | 从采样数据中发现未知的反应项，在后续任务中不断完善判断 | +| Grid Navigation | `cod_grid_navigation_workflow` | 逐步探索并记住地图中的移动代价，为后续任务选择低代价路线 | + +实现位于 [`trinity/common/workflows/connect_the_dots/`](../../trinity/common/workflows/connect_the_dots/)。 -各环境的 workflow 实现位于 `research/cod` 分支的 [`trinity/common/workflows/connect_the_dots/`](https://github.com/agentscope-ai/Trinity-RFT/tree/research/cod/trinity/common/workflows/connect_the_dots)。 +--- + +## 目录结构 + +``` +examples/research_cod/ +├── get_*_data.py # 各环境的数据生成器 +├── exp_plan_final/ # 主实验 +│ ├── train/ # 训练配置 +│ └── bench/ # 评测配置 +└── exp_plan_learn2ask/ # learn2ask的数据准备 + +trinity/common/workflows/connect_the_dots/ # CoD workflow 实现 +├── cod_workflow.py # 打包 / 迭代 hint / 任务奖励 +├── base_workflow.py # AsyncCoDMultiStepWorkflow 基类 +└── / # 每个环境各自有独立子目录 + ├── workflow.py # 任务渲染 / 评分 + └── prompts/ # system / user prompt + +trinity/algorithm/advantage_fn/cod_advantage.py # reward-to-go 信用分配（CoDAdvantageFn） +``` --- -## 运行 +## 快速开始 ```bash -git clone -b research/cod https://github.com/agentscope-ai/Trinity-RFT.git -cd Trinity-RFT +git clone +cd conda create -n trinity python=3.12 && conda activate trinity pip install -e ".[vllm,flash_attn]" pip install gymnasium jinja2 pandas ``` -**1. 生成数据** +**1. 生成数据。** ```bash # FrozenLake-Obscure python examples/research_cod/get_frozen_lake_data.py --local_dir examples/research_cod/data/frozen_lake_4567 \ @@ -78,20 +106,63 @@ 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. 训练** +**2. CoD 模型训练。** +将 `TRINITY_MODEL_PATH` 设为本地模型目录，并按实际硬件调整 YAML 中的 `cluster`（节点数 / 每节点 GPU 数）。在配置指定的 W&B 项目中，查看 `rollout/reward_iterative_hint_e2e_taskset_0_pos_{pos}/mean`，比较不同 pack 位置的平均奖励。 ```bash # FrozenLake-Obscure trinity run --config examples/research_cod/exp_plan_final/train/frozen_lake_obscure.yaml # Mixed（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 +``` + +Mixed OPD 每个环境使用一个教师模型，tokenizer 须与学生兼容： + +```bash +# 将教师 checkpoint 转换为 Hugging Face 格式 +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. 评测**（逐 checkpoint 衡量 OOD 泛化：同域更难版本 + 跨域未见环境） +**3. 评测。** +把训好的 checkpoint 逐个评测，衡量分布外（OOD）泛化：既包括 in-domain（训练环境的更难版本），也包括 cross-domain（训练中未见过的环境）。 ```bash # FrozenLake-Obscure ckpt → 同域 FrozenLake-hard，跨域 Alchemy-easy / Terminal bash examples/research_cod/exp_plan_final/bench/run_eval.sh --train-tasks frozen_lake_obscure @@ -99,10 +170,66 @@ bash examples/research_cod/exp_plan_final/bench/run_eval.sh --train-tasks frozen bash examples/research_cod/exp_plan_final/bench/run_eval.sh --train-tasks mixed_flobs_alchran ``` +PDE Discovery、Optimal Control 和 Grid Navigation 共用 checkpoint 评测配置。生成 PDE 评测集，Control 和 Grid 复用前面的 test split： + +```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 +``` + +`EVAL_PROJECT`、`EVAL_GROUP`、`EVAL_NAME` 对应实际训练目录。`TRINITY_MODEL_PATH` 指向基础模型，各 `global_step_*/actor/` 中须有 `model.safetensors`。 + +```bash +export TRINITY_CHECKPOINT_ROOT_DIR=/path/to/checkpoints + +# PDE RL checkpoint → 三个环境 +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 checkpoint → 三个环境 +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 checkpoint → 三个环境 +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 评测（目录中保留第 5、10、…、60 步，第 0 步评测基础模型）： + +```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 +``` + --- -## 引用 +## 关键配置项（`cod.cod_workflow_args`） + +| 字段 | 含义 | +|---|---| +| `activated_cod_methods` | 启用的 CoD 方法；主实验用 `["iterative_hint_e2e"]` | +| `hint_penalty_coef` / `length_penalty_coef` | 对 hint / 对正确解答的长度惩罚 | +| `task_pack_size` / `eval_task_pack_size` | 训练 / 评测时的 pack 大小 | +--- + +## 新增 CoD 环境 + +1. 在 `trinity/common/workflows/connect_the_dots//workflow.py` 继承 CoD 基类 workflow（可参考 `frozen_lake/workflow_obscure.py` 这个精简示例）。 +2. 在 [`trinity/common/workflows/__init__.py`](../../trinity/common/workflows/__init__.py) 的 `default_mapping` 里登记 `"cod__workflow": "...workflow.CoDWorkflow"`。 +3. 在 `examples/research_cod/` 下添加生成器 + 配置，并设 `default_workflow_type: 'cod__workflow'`。 + +继承 `AsyncCoDMultiStepWorkflow`（`base_workflow.py`）后，实现函数： + +- `step_async(step_num)`：为任务准备 prompt、调用模型推理、执行其动作 / 评判答案，将 reward 写入 `self.final_reward`，返回 `(是否继续, experiences)`。 +- `_get_feedback()`：返回环境对这步的反馈文本（写进 `exp.info["feedback"]`，用于生成 hint）。 +- `max_step_num`：单个任务的最大步数。 + +--- + +## 引用 ```bibtex @article{chen2026connect, title={Connect the Dots: Training LLMs for Long-Lifecycle Agents with Cross-Domain Generalization Via Reinforcement Learning}, diff --git a/examples/research_cod/assets/cod_qwen27b.png b/examples/research_cod/assets/cod_qwen27b.png new file mode 100644 index 00000000000..9126e752a4e Binary files /dev/null and b/examples/research_cod/assets/cod_qwen27b.png differ diff --git a/examples/research_cod/exp_plan_final/bench/eval_alchemy_easy.yaml b/examples/research_cod/exp_plan_final/bench/eval_alchemy_easy.yaml new file mode 100644 index 00000000000..aa51865ead7 --- /dev/null +++ b/examples/research_cod/exp_plan_final/bench/eval_alchemy_easy.yaml @@ -0,0 +1,89 @@ +# OOD eval: Alchemy-Random EASY +# Data prep: +# python examples/research_cod/get_alchemy_data.py --local_dir examples/research_cod/data/alchemy_random --train_size 50000 --test_size 4000 --seed 42 +project: "trinity-cod-final" +group: ${oc.env:EVAL_GROUP,frozen_lake_obscure} +name: ${oc.env:EVAL_NAME,qwen3-8b-flobs} +ray_namespace: "ood-alchemy-easy-4k-${oc.env:EVAL_GROUP,frozen_lake_obscure}-seed${oc.env:EVAL_SEED,1}" +mode: bench +checkpoint_root_dir: ${oc.env:TRINITY_CHECKPOINT_ROOT_DIR,./checkpoints} +continue_from_checkpoint: false +cod: + eval_task_pack_size: 8 + cod_workflow_args: + activated_cod_methods: ["iterative_hint_e2e"] + hint_example: false + reuse_workflow_instance: false + max_response_tokens_restraint: 1500 + hint_penalty_coef: 0.1 + cod_log_interval: 200 + log_dir: ${oc.env:TRINITY_LOG_DIR,./logs/research_cod} + exp_name: ${project}-${group}-${name}-ood-alchemy-easy-seed${oc.env:EVAL_SEED,1} + inject_pack_seed: true + pad_tasks_to_full_pack: true +model: + model_path: ${oc.env:TRINITY_MODEL_PATH} + max_prompt_tokens: 30000 + max_response_tokens: 2000 + max_model_len: 32000 +cluster: + node_num: 1 + gpu_per_node: 8 +buffer: + total_epochs: 1 + batch_size: 32 + explorer_input: + eval_tasksets: + - name: alchemy_random + storage_type: file + path: "examples/research_cod/data/alchemy_random" + split: test + workflow_args: + context_compression_mode: "keep_all" + length_penalty_coef: 0.1 + len_full_penalty: 2000000 + len_zero_penalty: 1000000 + max_rounds: 8 + material_mode: "unlimited" + scope: "per_pack" + show_recipes: false + show_elements: false + show_tiers: true + min_num_tiers: 3 + max_num_tiers: 4 + min_base_elements: 3 + max_base_elements: 4 + tier_shrink_min: 0.4 + tier_shrink_max: 0.6 + min_recipes_per_element: 1 + max_recipes_per_element: 1 + max_tier_gap: 1 + min_cross_tier_prob: 0.0 + max_cross_tier_prob: 0.1 + min_noise_node_ratio: 0.0 + max_noise_node_ratio: 0.0 + noise_chain_depth: 2 + rollout_args: + temperature: 1.0 + data_selector: + selector_type: shuffle + seed: ${oc.env:EVAL_SEED,1} + default_workflow_type: 'cod_random_alchemy_workflow' +explorer: + eval_on_startup: true + bench_on_latest_checkpoint: false # eval EVERY saved ckpt -> curve over ckpt step + runner_per_model: 32 + max_repeat_times_per_runner: 1 + max_timeout: 72000 + rollout_model: + enable_thinking: false + enable_history: false + engine_num: 8 + tensor_parallel_size: 1 + enable_prefix_caching: false + enforce_eager: false + dtype: bfloat16 + seed: ${oc.env:EVAL_SEED,1} + enable_openai_api: true +monitor: + monitor_type: wandb diff --git a/examples/research_cod/exp_plan_final/bench/eval_alchemy_hard.yaml b/examples/research_cod/exp_plan_final/bench/eval_alchemy_hard.yaml new file mode 100644 index 00000000000..cc48e928fcd --- /dev/null +++ b/examples/research_cod/exp_plan_final/bench/eval_alchemy_hard.yaml @@ -0,0 +1,89 @@ +# OOD eval: Alchemy-Random HARD +# Data prep: +# python examples/research_cod/get_alchemy_data.py --local_dir examples/research_cod/data/alchemy_random --train_size 50000 --test_size 4000 --seed 42 +project: "trinity-cod-final" +group: ${oc.env:EVAL_GROUP,mixed_flobs_alchran} +name: ${oc.env:EVAL_NAME,qwen3-8b-flobs} +ray_namespace: "ood-alchemy-hard-4k-${oc.env:EVAL_GROUP,mixed_flobs_alchran}-seed${oc.env:EVAL_SEED,1}" +mode: bench +checkpoint_root_dir: ${oc.env:TRINITY_CHECKPOINT_ROOT_DIR,./checkpoints} +continue_from_checkpoint: false +cod: + eval_task_pack_size: 8 + cod_workflow_args: + activated_cod_methods: ["iterative_hint_e2e"] + hint_example: false + reuse_workflow_instance: false + max_response_tokens_restraint: 1500 + hint_penalty_coef: 0.1 + cod_log_interval: 200 + log_dir: ${oc.env:TRINITY_LOG_DIR,./logs/research_cod} + exp_name: ${project}-${group}-${name}-ood-alchemy-hard-seed${oc.env:EVAL_SEED,1} + inject_pack_seed: true + pad_tasks_to_full_pack: true +model: + model_path: ${oc.env:TRINITY_MODEL_PATH} + max_prompt_tokens: 30000 + max_response_tokens: 2000 + max_model_len: 32000 +cluster: + node_num: 1 + gpu_per_node: 8 +buffer: + total_epochs: 1 + batch_size: 32 + explorer_input: + eval_tasksets: + - name: alchemy_random + storage_type: file + path: "examples/research_cod/data/alchemy_random" + split: test + workflow_args: + context_compression_mode: "keep_all" + length_penalty_coef: 0.1 + len_full_penalty: 2000000 + len_zero_penalty: 1000000 + max_rounds: 15 + material_mode: "limited" + scope: "per_pack" + show_recipes: false + show_elements: false + show_tiers: true + min_num_tiers: 3 + max_num_tiers: 4 + min_base_elements: 4 + max_base_elements: 6 + tier_shrink_min: 0.4 + tier_shrink_max: 0.8 + min_recipes_per_element: 1 + max_recipes_per_element: 2 + max_tier_gap: 1 + min_cross_tier_prob: 0.0 + max_cross_tier_prob: 0.05 + min_noise_node_ratio: 0.0 + max_noise_node_ratio: 0.3 + noise_chain_depth: 2 + rollout_args: + temperature: 1.0 + data_selector: + selector_type: shuffle + seed: ${oc.env:EVAL_SEED,1} + default_workflow_type: 'cod_random_alchemy_workflow' +explorer: + eval_on_startup: true + bench_on_latest_checkpoint: false # eval EVERY saved ckpt -> curve over ckpt step + runner_per_model: 32 + max_repeat_times_per_runner: 1 + max_timeout: 72000 + rollout_model: + enable_thinking: false + enable_history: false + engine_num: 8 + tensor_parallel_size: 1 + enable_prefix_caching: false + enforce_eager: false + dtype: bfloat16 + seed: ${oc.env:EVAL_SEED,1} + enable_openai_api: true +monitor: + monitor_type: wandb diff --git a/examples/research_cod/exp_plan_final/bench/eval_frozenlake_obscure_hard.yaml b/examples/research_cod/exp_plan_final/bench/eval_frozenlake_obscure_hard.yaml new file mode 100644 index 00000000000..fa240f4d265 --- /dev/null +++ b/examples/research_cod/exp_plan_final/bench/eval_frozenlake_obscure_hard.yaml @@ -0,0 +1,73 @@ +# OOD eval: FrozenLake-Obscure HARD (map "6767", env/agent steps 8/10) +# 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 +project: "trinity-cod-final" +group: ${oc.env:EVAL_GROUP,frozen_lake_obscure} +name: ${oc.env:EVAL_NAME,qwen3-8b-flobs} +ray_namespace: "ood-flobs-hard-4k-${oc.env:EVAL_GROUP,frozen_lake_obscure}-seed${oc.env:EVAL_SEED,1}" +mode: bench +checkpoint_root_dir: ${oc.env:TRINITY_CHECKPOINT_ROOT_DIR,./checkpoints} +continue_from_checkpoint: false +cod: + eval_task_pack_size: 8 + cod_workflow_args: + activated_cod_methods: ["iterative_hint_e2e"] + hint_example: false + reuse_workflow_instance: false + max_response_tokens_restraint: 1500 + hint_penalty_coef: 0.1 + cod_log_interval: 200 + log_dir: ${oc.env:TRINITY_LOG_DIR,./logs/research_cod} + exp_name: ${project}-${group}-${name}-ood-flobs-hard-seed${oc.env:EVAL_SEED,1} + inject_pack_seed: true + pad_tasks_to_full_pack: true +model: + model_path: ${oc.env:TRINITY_MODEL_PATH} + max_prompt_tokens: 30000 + max_response_tokens: 2000 + max_model_len: 32000 +cluster: + node_num: 1 + gpu_per_node: 8 +buffer: + total_epochs: 1 + batch_size: 32 + explorer_input: + eval_tasksets: + - name: frozen_lake_obscure_eval_hard + storage_type: file + path: "examples/research_cod/data/frozen_lake_6767" + split: test + format: + prompt_key: 'task_desc' + workflow_args: + context_compression_mode: "keep_all" + length_penalty_coef: 0.1 + len_full_penalty: 2000000 + len_zero_penalty: 1000000 + env_max_steps: 8 + agent_max_steps: 10 + mapping_mode: "per_pack" + rollout_args: + temperature: 1.0 + data_selector: + selector_type: shuffle + seed: ${oc.env:EVAL_SEED,1} + default_workflow_type: 'cod_frozenlake_obscure_workflow' +explorer: + eval_on_startup: true + bench_on_latest_checkpoint: false # eval EVERY saved ckpt -> curve over ckpt step + runner_per_model: 32 + max_repeat_times_per_runner: 1 + max_timeout: 72000 + rollout_model: + enable_thinking: false + enable_history: false + engine_num: 8 + tensor_parallel_size: 1 + enable_prefix_caching: false + enforce_eager: false + dtype: bfloat16 + seed: ${oc.env:EVAL_SEED,1} + enable_openai_api: true +monitor: + monitor_type: wandb diff --git a/examples/research_cod/exp_plan_final/bench/eval_mixed_opd_steps_0to60_all_domains.yaml b/examples/research_cod/exp_plan_final/bench/eval_mixed_opd_steps_0to60_all_domains.yaml new file mode 100644 index 00000000000..75aac2ad72f --- /dev/null +++ b/examples/research_cod/exp_plan_final/bench/eval_mixed_opd_steps_0to60_all_domains.yaml @@ -0,0 +1,156 @@ +# Evaluate the mixed-domain OPD student at steps 0, 5, ..., 60 on PDE, +# optimal control, and grid navigation in one benchmark job. +# Data prep: +# 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 +# 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 +# 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 +# +# Trinity bench evaluates every global_step_* directory it can see and does not +# provide a checkpoint-step filter. The checkpoint directory below is therefore +# an eval facade that must expose only global_step_{5,10,...,60}. Step 0 is the +# base model and is evaluated through explorer.eval_on_startup. +# +# The facade may contain symlinks to the model.safetensors files in the original +# training run, leaving that run untouched: +# checkpoints/trinity-cod/mixed_multi_teacher_opd/ +# / +# +# Run this single config as an 8-node x 4-GPU DLC job. Each target domain is a +# separate eval taskset, batch, and metric namespace. + +project: ${oc.env:EVAL_PROJECT,trinity-cod} +group: ${oc.env:EVAL_GROUP,mixed_multi_teacher_opd} +name: ${oc.env:EVAL_NAME,mixed-opd-eval-5to60} +ray_namespace: "mixed-opd-0to60-all-domains-seed${oc.env:EVAL_SEED,1}" +mode: bench +checkpoint_root_dir: ${oc.env:TRINITY_CHECKPOINT_ROOT_DIR,./checkpoints} +continue_from_checkpoint: true + +cod: + task_pack_size: 4 + eval_task_pack_size: 8 + packing_strategy: cod + cod_workflow_args: + activated_cod_methods: ["iterative_hint_e2e"] + hint_example: false + reuse_workflow_instance: false + context_compression_mode: "keep_all" + max_response_tokens_restraint: 4000 + hint_penalty_coef: 0.1 + cod_log_interval: 1 + log_dir: ${oc.env:TRINITY_LOG_DIR,./logs/research_cod} + exp_name: ${project}-${group}-${name}-seed${oc.env:EVAL_SEED,1} + inject_pack_seed: true + stable_pack_seed: true + pad_tasks_to_full_pack: true + +model: + model_path: ${oc.env:TRINITY_MODEL_PATH} + max_prompt_tokens: 27000 + max_response_tokens: 4000 + max_model_len: 32000 + +cluster: + node_num: 8 + gpu_per_node: 4 + +buffer: + total_epochs: 1 + batch_size: 64 + explorer_input: + eval_tasksets: + - name: pde_discovery + storage_type: file + path: ${oc.env:PDE_EVAL_DATA_DIR,examples/research_cod/data/pde_discovery_eval_hard_stratified_4000_disjoint_testonly} + split: test + format: + prompt_key: "task_desc" + response_key: "answer" + workflow_args: + update_context_impl: "pde_update_context" + context_compression_mode: "keep_all" + max_steps: 6 + point_budget: 6 + min_reaction_terms: 1 + max_reaction_terms: 3 + ground_truth_family: "physical_full_eval_4000.json" + pde_state_abs_limit: 3.0 + initial_condition_shape: + mode_count_range: [1, 2] + noise_level: 0.002 + kappa_threshold: 100.0 + dump_trajectories: false + rollout_args: + temperature: 1.0 + data_selector: + selector_type: sequential + repeat_times: 1 + default_workflow_type: "trinity.common.workflows.connect_the_dots.pde_discovery.workflow.CoDPDEDiscoveryWorkflow" + + - name: optimal_control_1d + storage_type: file + path: ${oc.env:CONTROL_1D_EVAL_DATA_DIR,examples/research_cod/data/optimal_control} + split: test + format: + prompt_key: "task_desc" + workflow_args: + update_context_impl: "update_context" + context_compression_mode: "keep_all" + include_previous_control_policy: true + a_env_range: [0.55, 1.20] + b_env_range: [-1.5, 1.5] + min_abs_b_env: 0.35 + enable_process_noise: false + rollout_args: + temperature: 1.0 + data_selector: + selector_type: sequential + repeat_times: 1 + default_workflow_type: "cod_optimalcontrol_workflow" + + - name: grid_navigation + storage_type: file + path: ${oc.env:GRID_NAVIGATION_EVAL_DATA_DIR,examples/research_cod/data/grid_navigation} + split: test + workflow_args: + update_context_impl: "update_context" + context_compression_mode: "keep_all" + length_penalty_coef: 0.0 + grid_min_size: 10 + grid_max_size: 14 + landscape_num_components: 5 + landscape_min_scale: 0.15 + landscape_max_scale: 0.45 + min_rounds: 3 + max_rounds: 5 + reveal_radius: 1 + rollout_args: + temperature: 1.0 + data_selector: + selector_type: sequential + repeat_times: 1 + default_workflow_type: "cod_grid_navigation_workflow" + +explorer: + name: mixed_opd_0to60_all_domains_seed${oc.env:EVAL_SEED,1} + eval_on_startup: true + bench_on_latest_checkpoint: false + runner_per_model: 32 + max_repeat_times_per_runner: 1 + max_timeout: 72000 + rollout_model: + enable_thinking: false + enable_history: false + engine_num: 32 + tensor_parallel_size: 1 + enable_prefix_caching: false + enforce_eager: false + dtype: bfloat16 + seed: ${oc.env:EVAL_SEED,1} + +synchronizer: + sync_method: "checkpoint" + sync_timeout: 72000 + +monitor: + monitor_type: wandb diff --git a/examples/research_cod/exp_plan_final/bench/eval_source_checkpoints_all_domains.yaml b/examples/research_cod/exp_plan_final/bench/eval_source_checkpoints_all_domains.yaml new file mode 100644 index 00000000000..3d357ba0e14 --- /dev/null +++ b/examples/research_cod/exp_plan_final/bench/eval_source_checkpoints_all_domains.yaml @@ -0,0 +1,168 @@ +# Evaluate every checkpoint from one source-domain training run on all three CoD +# domains. Submit this same config as one independent DLC job per source domain; +# the jobs may run in parallel without a shell launcher. Each job evaluates the +# base model first and then every global_step_* checkpoint in order. +# +# Generate the three fixed 4000-task eval datasets once before submitting jobs: +# 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 +# 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 +# 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 +# The generated parquet files are ignored by Git. On DLC, either prepare them in +# the checkout or point PDE_EVAL_DATA_DIR, CONTROL_1D_EVAL_DATA_DIR, and +# GRID_NAVIGATION_EVAL_DATA_DIR at fixed shared copies. +# +# Use EVAL_SEED=1,2,3 for independent rollout seeds. PDE selects its specialized +# update-context workflow through the taskset's domain-specific workflow_args. +# The defaults select the converted optimal-control bf16 facade. Override the +# EVAL_* variables for PDE or grid-navigation. Every source directory must be a +# safe facade in which each global_step_*/actor contains model.safetensors; do +# not point this benchmark at multi-shard huggingface/ directories. +# Run with: +# EVAL_SEED=1 \ +# trinity run \ +# --config examples/research_cod/exp_plan_final/bench/eval_source_checkpoints_all_domains.yaml \ +# --dlc + +project: ${oc.env:EVAL_PROJECT,trinity-cod} +group: ${oc.env:EVAL_GROUP,optimal_control} +name: ${oc.env:EVAL_NAME,qwen3.6-27b-optimal-control-cod-bf16-100to500} +ray_namespace: "domain-transfer-${oc.env:EVAL_TRAIN_DOMAIN,control}-seed${oc.env:EVAL_SEED,1}" +mode: bench +checkpoint_root_dir: ${oc.env:TRINITY_CHECKPOINT_ROOT_DIR,./checkpoints} +continue_from_checkpoint: true + +cod: + task_pack_size: 4 + eval_task_pack_size: 8 + packing_strategy: cod + cod_workflow_args: + activated_cod_methods: ["iterative_hint_e2e"] + hint_example: false + reuse_workflow_instance: false + context_compression_mode: "keep_all" + max_response_tokens_restraint: 4000 + hint_penalty_coef: 0.1 + cod_log_interval: 1 + log_dir: ${oc.env:TRINITY_LOG_DIR,./logs/research_cod} + exp_name: ${project}-${group}-${name}-domain-transfer-seed${oc.env:EVAL_SEED,1} + inject_pack_seed: true + stable_pack_seed: true + pad_tasks_to_full_pack: true + +model: + model_path: ${oc.env:TRINITY_MODEL_PATH} + max_prompt_tokens: 27000 + max_response_tokens: 4000 + max_model_len: 32000 + +cluster: + node_num: 6 + gpu_per_node: 4 + +buffer: + total_epochs: 1 + batch_size: 64 + explorer_input: + eval_tasksets: + - name: pde_discovery + storage_type: file + path: ${oc.env:PDE_EVAL_DATA_DIR,examples/research_cod/data/pde_discovery_eval_hard_stratified_4000_disjoint_testonly} + split: test + format: + prompt_key: "task_desc" + response_key: "answer" + workflow_args: + update_context_impl: "pde_update_context" + context_compression_mode: "keep_all" + max_steps: 6 + point_budget: 6 + min_reaction_terms: 1 + max_reaction_terms: 3 + ground_truth_family: "physical_full_eval_4000.json" + pde_state_abs_limit: 3.0 + initial_condition_shape: + mode_count_range: [1, 2] + noise_level: 0.002 + kappa_threshold: 100.0 + dump_trajectories: false + rollout_args: + temperature: 1.0 + data_selector: + selector_type: sequential + repeat_times: 1 + default_workflow_type: "trinity.common.workflows.connect_the_dots.pde_discovery.workflow.CoDPDEDiscoveryWorkflow" + + - name: optimal_control_1d + storage_type: file + path: ${oc.env:CONTROL_1D_EVAL_DATA_DIR,examples/research_cod/data/optimal_control} + split: test + format: + prompt_key: "task_desc" + workflow_args: + context_compression_mode: "keep_all" + include_previous_control_policy: true + a_env_range: [0.55, 1.20] + b_env_range: [-1.5, 1.5] + min_abs_b_env: 0.35 + enable_process_noise: false + rollout_args: + temperature: 1.0 + data_selector: + selector_type: sequential + repeat_times: 1 + default_workflow_type: "cod_optimalcontrol_workflow" + + - name: grid_navigation + storage_type: file + path: ${oc.env:GRID_NAVIGATION_EVAL_DATA_DIR,examples/research_cod/data/grid_navigation} + split: test + workflow_args: + context_compression_mode: "keep_all" + length_penalty_coef: 0.0 + grid_min_size: 10 + grid_max_size: 14 + landscape_num_components: 5 + landscape_min_scale: 0.15 + landscape_max_scale: 0.45 + min_rounds: 3 + max_rounds: 5 + reveal_radius: 1 + rollout_args: + temperature: 1.0 + data_selector: + selector_type: sequential + repeat_times: 1 + default_workflow_type: "cod_grid_navigation_workflow" + +explorer: + name: domain_transfer_${oc.env:EVAL_TRAIN_DOMAIN,control}_seed${oc.env:EVAL_SEED,1} + eval_on_startup: true + bench_on_latest_checkpoint: false + runner_per_model: 32 + max_repeat_times_per_runner: 1 + max_timeout: 72000 + rollout_model: + enable_thinking: false + enable_history: false + engine_num: 24 + tensor_parallel_size: 1 + enable_prefix_caching: false + enforce_eager: false + dtype: bfloat16 + seed: ${oc.env:EVAL_SEED,1} + +synchronizer: + sync_method: "checkpoint" + sync_timeout: 72000 + +monitor: + monitor_type: wandb diff --git a/examples/research_cod/exp_plan_final/bench/eval_strong_api_model_all_domains.yaml b/examples/research_cod/exp_plan_final/bench/eval_strong_api_model_all_domains.yaml new file mode 100644 index 00000000000..fc5968628e9 --- /dev/null +++ b/examples/research_cod/exp_plan_final/bench/eval_strong_api_model_all_domains.yaml @@ -0,0 +1,176 @@ +# Evaluate one stronger OpenAI-compatible API model on the fixed PDE, +# optimal-control, and grid-navigation test sets. +# Data prep: +# 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 +# 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 +# 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 +# +# This experiment does not commit any Trinity core changes. Before launching it, +# apply a temporary, uncommitted DSW adapter that makes ModelWrapper.chat_async +# call the OpenAI-compatible API for external models, returns model_version=0, +# tolerates base URLs ending in /v1, and reads QWEN_REASONING_EFFORT or +# QWEN_THINKING_BUDGET from the environment. The same temporary adapter makes +# Optimal Control delegate _chat to ModelWrapper.chat_async, so all three +# domains share one API path. +# +# Thinking controls differ by model: +# qwen3.8-max supports reasoning_effort=none/low/medium/xhigh; none disables +# thinking, while low/medium/xhigh select increasing inference intensity. +# qwen3.7-plus instead uses enable_thinking=true/false and a numeric +# thinking_budget; it has no native low/medium/xhigh reasoning_effort levels. +# +# Primary comparison (approximately 4096 reasoning tokens for both models): +# qwen3.8-max: reasoning_effort=low +# qwen3.7-plus: enable_thinking=true, thinking_budget=4096 +# both: max_prompt_tokens=54000, max_completion_tokens=8000 +# The larger prompt budget accommodates prior responses carried into later +# rounds/episodes and reduces truncation relative to the 27B 27000/4000 setup. +# In both runs, set preserve_thinking=false because the workflows do not retain +# reasoning_content in assistant-message history. The temporary DSW adapter +# supplies these OpenAI-compatible request parameters. +# +# Performance experiments to consider after the primary run: +# 1. token-matched comparison with the trained 27B model: response=4000; +# 2. qwen3.8-max reasoning effort: none / low / medium / xhigh; +# 3. qwen3.7-plus thinking: disabled or budgets 4096 / 16384 / 32768; +# 4. completion budget: 8000 / 16000 / 32000, guided by truncation rate. +# Budgets above 8000 require a larger max_completion_tokens value; otherwise +# the total completion cap truncates reasoning and/or leaves too little answer. + +project: ${oc.env:EVAL_PROJECT,trinity-cod} +group: ${oc.env:EVAL_GROUP,strong_api_model_all_domains} +name: ${oc.env:EVAL_NAME,qwen3.8-max-low-all-domains} +ray_namespace: "strong-api-${oc.env:EVAL_PROFILE,low}-seed${oc.env:EVAL_SEED,1}" +mode: bench +checkpoint_root_dir: ${oc.env:TRINITY_CHECKPOINT_ROOT_DIR,./checkpoints} +continue_from_checkpoint: false + +cod: + task_pack_size: 4 + eval_task_pack_size: 8 + packing_strategy: cod + cod_workflow_args: + activated_cod_methods: ["iterative_hint_e2e"] + hint_example: false + reuse_workflow_instance: false + context_compression_mode: "keep_all" + max_response_tokens_restraint: ${oc.env:EVAL_MAX_RESPONSE_TOKENS,8000} + hint_penalty_coef: 0.1 + cod_log_interval: 200 + log_dir: ${oc.env:TRINITY_LOG_DIR,./logs/research_cod} + exp_name: ${project}-${group}-${name}-seed${oc.env:EVAL_SEED,1} + inject_pack_seed: true + stable_pack_seed: true + pad_tasks_to_full_pack: true + +model: + model_path: ${oc.env:TRINITY_MODEL_NAME,qwen3.8-max} + max_prompt_tokens: ${oc.env:EVAL_MAX_PROMPT_TOKENS,54000} + max_response_tokens: ${oc.env:EVAL_MAX_RESPONSE_TOKENS,8000} + max_model_len: ${oc.env:EVAL_MAX_MODEL_LEN,64000} + enable_thinking: ${oc.decode:${oc.env:QWEN_ENABLE_THINKING,true}} + external_model: + enable: true + model_name: ${oc.env:TRINITY_MODEL_NAME,qwen3.8-max} + base_url_env: OPENAI_BASE_URL + api_key_env: OPENAI_API_KEY + +cluster: + ray_address: "auto" + node_num: 1 + gpu_per_node: 0 + +buffer: + total_epochs: 1 + batch_size: 64 + explorer_input: + eval_tasksets: + - name: pde_discovery + storage_type: file + path: ${oc.env:PDE_EVAL_DATA_DIR,examples/research_cod/data/pde_discovery_eval_hard_stratified_4000_disjoint_testonly} + split: test + format: + prompt_key: "task_desc" + response_key: "answer" + workflow_args: + update_context_impl: "pde_update_context" + context_compression_mode: "keep_all" + max_steps: 6 + point_budget: 6 + min_reaction_terms: 1 + max_reaction_terms: 3 + ground_truth_family: "physical_full_eval_4000.json" + pde_state_abs_limit: 3.0 + initial_condition_shape: + mode_count_range: [1, 2] + noise_level: 0.002 + kappa_threshold: 100.0 + dump_trajectories: false + rollout_args: + temperature: 1.0 + data_selector: + selector_type: sequential + repeat_times: 1 + default_workflow_type: "trinity.common.workflows.connect_the_dots.pde_discovery.workflow.CoDPDEDiscoveryWorkflow" + + - name: optimal_control_1d + storage_type: file + path: ${oc.env:CONTROL_1D_EVAL_DATA_DIR,examples/research_cod/data/optimal_control} + split: test + format: + prompt_key: "task_desc" + workflow_args: + update_context_impl: "update_context" + context_compression_mode: "keep_all" + include_previous_control_policy: true + a_env_range: [0.55, 1.20] + b_env_range: [-1.5, 1.5] + min_abs_b_env: 0.35 + enable_process_noise: false + rollout_args: + temperature: 1.0 + data_selector: + selector_type: sequential + repeat_times: 1 + default_workflow_type: "cod_optimalcontrol_workflow" + + - name: grid_navigation + storage_type: file + path: ${oc.env:GRID_NAVIGATION_EVAL_DATA_DIR,examples/research_cod/data/grid_navigation} + split: test + workflow_args: + update_context_impl: "update_context" + context_compression_mode: "keep_all" + length_penalty_coef: 0.0 + grid_min_size: 10 + grid_max_size: 14 + landscape_num_components: 5 + landscape_min_scale: 0.15 + landscape_max_scale: 0.45 + min_rounds: 3 + max_rounds: 5 + reveal_radius: 1 + rollout_args: + temperature: 1.0 + data_selector: + selector_type: sequential + repeat_times: 1 + default_workflow_type: "cod_grid_navigation_workflow" + +explorer: + name: strong_api_${oc.env:EVAL_PROFILE,low}_seed${oc.env:EVAL_SEED,1} + concurrent_mode: asynchronous + eval_on_startup: true + bench_on_latest_checkpoint: false + runner_per_model: ${oc.env:EVAL_RUNNER_PER_MODEL,32} + max_repeat_times_per_runner: 1 + max_timeout: 72000 + rollout_model: + engine_type: external + enable_history: false + engine_num: 1 + tensor_parallel_size: 1 + seed: ${oc.env:EVAL_SEED,1} + +monitor: + monitor_type: wandb diff --git a/examples/research_cod/exp_plan_final/bench/eval_terminal.yaml b/examples/research_cod/exp_plan_final/bench/eval_terminal.yaml new file mode 100644 index 00000000000..975864e9ce5 --- /dev/null +++ b/examples/research_cod/exp_plan_final/bench/eval_terminal.yaml @@ -0,0 +1,68 @@ +# OOD eval: Terminal +# Data prep: +# 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 +project: "trinity-cod-final" +group: ${oc.env:EVAL_GROUP,frozen_lake_obscure} +name: ${oc.env:EVAL_NAME,qwen3-8b-flobs} +ray_namespace: "ood-terminal-4k-${oc.env:EVAL_GROUP,frozen_lake_obscure}-seed${oc.env:EVAL_SEED,1}" +mode: bench +checkpoint_root_dir: ${oc.env:TRINITY_CHECKPOINT_ROOT_DIR,./checkpoints} +continue_from_checkpoint: false +cod: + eval_task_pack_size: 8 + cod_workflow_args: + activated_cod_methods: ["iterative_hint_e2e"] + hint_example: false + reuse_workflow_instance: false + max_response_tokens_restraint: 1500 + hint_penalty_coef: 0.1 + cod_log_interval: 200 + log_dir: ${oc.env:TRINITY_LOG_DIR,./logs/research_cod} + exp_name: ${project}-${group}-${name}-ood-terminal-seed${oc.env:EVAL_SEED,1} + inject_pack_seed: true + pad_tasks_to_full_pack: true +model: + model_path: ${oc.env:TRINITY_MODEL_PATH} + max_prompt_tokens: 30000 + max_response_tokens: 2000 + max_model_len: 32000 +cluster: + node_num: 1 + gpu_per_node: 8 +buffer: + total_epochs: 1 + batch_size: 32 + explorer_input: + eval_tasksets: + - name: terminal + storage_type: file + path: "examples/research_cod/data/terminal" + split: test + workflow_args: + agent_max_steps: 12 + composite_ratio: 0.5 + context_compression_mode: "keep_all" + rollout_args: + temperature: 1.0 + data_selector: + selector_type: shuffle + seed: ${oc.env:EVAL_SEED,1} + default_workflow_type: 'cod_terminal_workflow' +explorer: + eval_on_startup: true + bench_on_latest_checkpoint: false # eval EVERY saved ckpt -> curve over ckpt step + runner_per_model: 32 + max_repeat_times_per_runner: 1 + max_timeout: 72000 + rollout_model: + enable_thinking: false + enable_history: false + engine_num: 8 + tensor_parallel_size: 1 + enable_prefix_caching: false + enforce_eager: false + dtype: bfloat16 + seed: ${oc.env:EVAL_SEED,1} + enable_openai_api: true +monitor: + monitor_type: wandb diff --git a/examples/research_cod/exp_plan_final/bench/run_eval.sh b/examples/research_cod/exp_plan_final/bench/run_eval.sh new file mode 100644 index 00000000000..6a3116d76b7 --- /dev/null +++ b/examples/research_cod/exp_plan_final/bench/run_eval.sh @@ -0,0 +1,62 @@ +#!/usr/bin/env bash +# OOD-generalization eval: evaluate trained checkpoint(s) on a harder in-domain environment and on +# unseen cross-domain environments. Each job bench-evals EVERY saved checkpoint (a reward curve over +# checkpoint step). ckpt dir = /trinity-cod-final//; jobs use pack_size=8, 1 node x 8 GPU. +# +# frozen_lake_obscure ckpt: flobs (FrozenLake-hard, in-domain) + alchemy (easy) + terminal +# mixed_flobs_alchran ckpt: flobs (FrozenLake-hard, in-domain) + alchemy (hard) + terminal +# +# Usage (run from repo root; set TRINITY_CHECKPOINT_ROOT_DIR if ckpts aren't at the yaml default): +# bash .../run_eval.sh --train-tasks [--eval-tasks ] +# --train-tasks frozen_lake_obscure | mixed_flobs_alchran (one or more) +# --eval-tasks flobs | alchemy | terminal (default: all 3; comma- or space-separated) +# +# bash .../run_eval.sh --train-tasks frozen_lake_obscure +# bash .../run_eval.sh --train-tasks frozen_lake_obscure mixed_flobs_alchran --eval-tasks alchemy terminal +set -u +B="examples/research_cod/exp_plan_final/bench" +SEED="${EVAL_SEED:-1}" + +TRAIN_TASKS=() +EVAL_TASKS=() +mode="" +while [ $# -gt 0 ]; do + case "$1" in + --train-tasks|--train-task) mode=train ;; + --eval-tasks|--eval-task) mode=eval ;; + -h|--help) sed -n '2,13p' "$0"; exit 0 ;; + --*) echo "Unknown flag: $1"; exit 1 ;; + *) case "$mode" in + train) TRAIN_TASKS+=("$1") ;; + eval) EVAL_TASKS+=(${1//,/ }) ;; + *) echo "Pass --train-tasks before task names"; exit 1 ;; + esac ;; + esac + shift +done +[ ${#TRAIN_TASKS[@]} -ge 1 ] || { echo "usage: run_eval.sh --train-tasks ... [--eval-tasks ...]"; exit 1; } +[ ${#EVAL_TASKS[@]} -ge 1 ] || EVAL_TASKS=(flobs alchemy terminal) + +for t in "${TRAIN_TASKS[@]}"; do case "$t" in frozen_lake_obscure|mixed_flobs_alchran) ;; *) echo "unknown train-task: $t"; exit 1 ;; esac; done +for e in "${EVAL_TASKS[@]}"; do case "$e" in flobs|alchemy|terminal) ;; *) echo "unknown eval-task: $e"; exit 1 ;; esac; done + +name_of() { case "$1" in frozen_lake_obscure) echo qwen3-8b-flobs ;; mixed_flobs_alchran) echo qwen3-8b-mixed ;; esac; } +cfg_of() { # -> bench config name (alchemy difficulty follows the train task) + case "$2" in + flobs) echo eval_frozenlake_obscure_hard ;; + alchemy) case "$1" in frozen_lake_obscure) echo eval_alchemy_easy ;; mixed_flobs_alchran) echo eval_alchemy_hard ;; esac ;; + terminal) echo eval_terminal ;; + esac +} + +TOTAL=$(( ${#TRAIN_TASKS[@]} * ${#EVAL_TASKS[@]} )); DONE=0 +echo "############ run_eval: train=[${TRAIN_TASKS[*]}] eval=[${EVAL_TASKS[*]}] -> $TOTAL jobs ############" +for t in "${TRAIN_TASKS[@]}"; do + NAME=$(name_of "$t") + for e in "${EVAL_TASKS[@]}"; do + cfg=$(cfg_of "$t" "$e"); DONE=$((DONE + 1)) + echo ">>> [$DONE/$TOTAL] train=$t name=$NAME eval=$e seed=$SEED cfg=$cfg" + EVAL_GROUP="$t" EVAL_NAME="$NAME" EVAL_SEED="$SEED" trinity run --config "$B/$cfg.yaml" || echo " !!! FAILED: train=$t eval=$e" + done +done +echo "############ done: $DONE/$TOTAL ############" diff --git a/examples/research_cod/exp_plan_final/train/alchemy_abl_no_adaptive_red.yaml b/examples/research_cod/exp_plan_final/train/alchemy_abl_no_adaptive_red.yaml new file mode 100644 index 00000000000..62dcd713831 --- /dev/null +++ b/examples/research_cod/exp_plan_final/train/alchemy_abl_no_adaptive_red.yaml @@ -0,0 +1,134 @@ +# RL ablation: train on Alchemy-Random, pack=4. Removes adaptive RED-Weight +# (red_weight_adaptive_temp / red_weight_adaptive_version) to verify its effect. +# Data prep: python examples/research_cod/get_alchemy_data.py --local_dir examples/research_cod/data/alchemy_random --train_size 50000 --test_size 4000 --seed 42 +project: "trinity-cod-final" +group: "alchemy_abl_no_adaptive_red" +name: "qwen3-8b-alchemy" +mode: both +checkpoint_root_dir: ${oc.env:TRINITY_CHECKPOINT_ROOT_DIR,./checkpoints} +continue_from_checkpoint: false +cod: + task_pack_size: 4 + eval_task_pack_size: 8 + cod_workflow_args: + hint_example: false + activated_cod_methods: ["iterative_hint_e2e"] + reuse_workflow_instance: false + max_response_tokens_restraint: 1000 + hint_penalty_coef: 0.1 + cod_log_interval: 200 + log_dir: ${oc.env:TRINITY_LOG_DIR,./logs/research_cod} + exp_name: ${project}-${group}-${name} + inject_pack_seed: true + pad_tasks_to_full_pack: true +algorithm: + algorithm_type: multi_step_grpo + sample_strategy: staleness_control + sample_strategy_args: + max_staleness: 20 + repeat_times: 8 + optimizer: + lr: 1e-6 + advantage_fn: cod + advantage_fn_args: + std_cal_level: "none" + iterative_hint_e2e_causal: true + policy_loss_fn: rec + policy_loss_fn_args: # REC-OneSide-NoIS + epsilon_low: 0.2 + epsilon_high: 0.2 + clip_mode: "one-side" + weight: "none" + kl_loss_fn: k2 + kl_loss_fn_args: + kl_coef: 0.0 +model: + model_path: ${oc.env:TRINITY_MODEL_PATH} + max_prompt_tokens: 10000 + max_response_tokens: 2000 + max_model_len: 12000 +cluster: + node_num: 3 + gpu_per_node: 8 +buffer: + total_epochs: 1 + batch_size: 32 + train_batch_size: 1440 + explorer_input: + tasksets: + - name: alchemy_random + storage_type: file + path: "examples/research_cod/data/alchemy_random" + split: train + workflow_args: + context_compression_mode: "keep_all" + # length penalty in reward design + length_penalty_coef: 0.1 + len_full_penalty: 2000000 + len_zero_penalty: 1000000 + # specific to alchemy-random + max_rounds: 8 + material_mode: "unlimited" + scope: "per_pack" + show_recipes: false + show_elements: false + show_tiers: true + min_num_tiers: 3 + max_num_tiers: 4 + min_base_elements: 3 + max_base_elements: 4 + tier_shrink_min: 0.4 + tier_shrink_max: 0.6 + min_recipes_per_element: 1 + max_recipes_per_element: 1 + max_tier_gap: 1 + min_cross_tier_prob: 0.0 + max_cross_tier_prob: 0.1 + min_noise_node_ratio: 0.0 + max_noise_node_ratio: 0.0 + noise_chain_depth: 2 + rollout_args: + temperature: 1.0 + default_workflow_type: 'cod_random_alchemy_workflow' + trainer_input: + experience_buffer: + name: buffer + storage_type: queue + max_read_timeout: 72000 + replay_buffer: + enable: true + priority_fn: decay_limit_randomization + priority_fn_args: + sigma: 2.0 +explorer: + runner_per_model: 32 + max_repeat_times_per_runner: 1 + max_timeout: 72000 + rollout_model: + enable_thinking: false + enable_history: false + engine_num: 8 + tensor_parallel_size: 1 + enable_prefix_caching: false + enforce_eager: false + dtype: bfloat16 + seed: 42 +trainer: + save_interval: 100 + grad_clip: 1.0 + use_dynamic_bsz: true + max_token_len_per_gpu: 16384 + ulysses_sequence_parallel_size: 1 + trainer_config: + actor_rollout_ref: + actor: + checkpoint: + load_contents: ['model', 'extra'] + save_contents: ['model', 'extra'] +synchronizer: + sync_style: dynamic_by_explorer + sync_method: 'nccl' + sync_interval: 4 + sync_timeout: 72000 +monitor: + monitor_type: wandb diff --git a/examples/research_cod/exp_plan_final/train/alchemy_abl_no_adaptive_red_is.yaml b/examples/research_cod/exp_plan_final/train/alchemy_abl_no_adaptive_red_is.yaml new file mode 100644 index 00000000000..e3fa3517619 --- /dev/null +++ b/examples/research_cod/exp_plan_final/train/alchemy_abl_no_adaptive_red_is.yaml @@ -0,0 +1,135 @@ +# RL ablation: train on Alchemy-Random, pack=4. Builds on alchemy_abl_no_adaptive_red.yaml (adaptive +# RED-Weight already removed) and additionally sets policy_loss_fn_args.weight "none" -> "importance_sampling" +# to verify the rationale of discarding IS. +# Data prep: python examples/research_cod/get_alchemy_data.py --local_dir examples/research_cod/data/alchemy_random --train_size 50000 --test_size 4000 --seed 42 +project: "trinity-cod-final" +group: "alchemy_abl_no_adaptive_red_is" +name: "qwen3-8b-alchemy" +mode: both +checkpoint_root_dir: ${oc.env:TRINITY_CHECKPOINT_ROOT_DIR,./checkpoints} +continue_from_checkpoint: false +cod: + task_pack_size: 4 + eval_task_pack_size: 8 + cod_workflow_args: + hint_example: false + activated_cod_methods: ["iterative_hint_e2e"] + reuse_workflow_instance: false + max_response_tokens_restraint: 1000 + hint_penalty_coef: 0.1 + cod_log_interval: 200 + log_dir: ${oc.env:TRINITY_LOG_DIR,./logs/research_cod} + exp_name: ${project}-${group}-${name} + inject_pack_seed: true + pad_tasks_to_full_pack: true +algorithm: + algorithm_type: multi_step_grpo + sample_strategy: staleness_control + sample_strategy_args: + max_staleness: 20 + repeat_times: 8 + optimizer: + lr: 1e-6 + advantage_fn: cod + advantage_fn_args: + std_cal_level: "none" + iterative_hint_e2e_causal: true + policy_loss_fn: rec + policy_loss_fn_args: # REC-OneSide-IS + epsilon_low: 0.2 + epsilon_high: 0.2 + clip_mode: "one-side" + weight: "importance_sampling" + kl_loss_fn: k2 + kl_loss_fn_args: + kl_coef: 0.0 +model: + model_path: ${oc.env:TRINITY_MODEL_PATH} + max_prompt_tokens: 10000 + max_response_tokens: 2000 + max_model_len: 12000 +cluster: + node_num: 3 + gpu_per_node: 8 +buffer: + total_epochs: 1 + batch_size: 32 + train_batch_size: 1440 + explorer_input: + tasksets: + - name: alchemy_random + storage_type: file + path: "examples/research_cod/data/alchemy_random" + split: train + workflow_args: + context_compression_mode: "keep_all" + # length penalty in reward design + length_penalty_coef: 0.1 + len_full_penalty: 2000000 + len_zero_penalty: 1000000 + # specific to alchemy-random + max_rounds: 8 + material_mode: "unlimited" + scope: "per_pack" + show_recipes: false + show_elements: false + show_tiers: true + min_num_tiers: 3 + max_num_tiers: 4 + min_base_elements: 3 + max_base_elements: 4 + tier_shrink_min: 0.4 + tier_shrink_max: 0.6 + min_recipes_per_element: 1 + max_recipes_per_element: 1 + max_tier_gap: 1 + min_cross_tier_prob: 0.0 + max_cross_tier_prob: 0.1 + min_noise_node_ratio: 0.0 + max_noise_node_ratio: 0.0 + noise_chain_depth: 2 + rollout_args: + temperature: 1.0 + default_workflow_type: 'cod_random_alchemy_workflow' + trainer_input: + experience_buffer: + name: buffer + storage_type: queue + max_read_timeout: 72000 + replay_buffer: + enable: true + priority_fn: decay_limit_randomization + priority_fn_args: + sigma: 2.0 +explorer: + runner_per_model: 32 + max_repeat_times_per_runner: 1 + max_timeout: 72000 + rollout_model: + enable_thinking: false + enable_history: false + engine_num: 8 + tensor_parallel_size: 1 + enable_prefix_caching: false + enforce_eager: false + dtype: bfloat16 + seed: 42 +trainer: + save_interval: 100 + grad_clip: 1.0 + use_dynamic_bsz: true + max_token_len_per_gpu: 16384 + ulysses_sequence_parallel_size: 1 + trainer_config: + actor_rollout_ref: + actor: + checkpoint: + load_contents: ['model', 'extra'] + save_contents: ['model', 'extra'] +synchronizer: + sync_style: dynamic_by_explorer + sync_method: 'nccl' + sync_interval: 4 + sync_timeout: 72000 +monitor: + monitor_type: wandb diff --git a/examples/research_cod/exp_plan_final/train/frozen_lake_obscure.yaml b/examples/research_cod/exp_plan_final/train/frozen_lake_obscure.yaml new file mode 100644 index 00000000000..ddbbe4b7c4d --- /dev/null +++ b/examples/research_cod/exp_plan_final/train/frozen_lake_obscure.yaml @@ -0,0 +1,120 @@ +# Data prep: +# python examples/research_cod/get_frozen_lake_data.py --local_dir examples/research_cod/data/frozen_lake_4567 --train_size 50000 --test_size 4000 --map_min_size 4 --map_max_size 5 --tile_min_prob 0.6 --tile_max_prob 0.7 +project: "trinity-cod-final" +group: "frozen_lake_obscure" +name: "qwen3-8b-flobs" +mode: both +checkpoint_root_dir: ${oc.env:TRINITY_CHECKPOINT_ROOT_DIR,./checkpoints} +continue_from_checkpoint: false +cod: + task_pack_size: 4 + eval_task_pack_size: 8 + cod_workflow_args: + hint_example: false + activated_cod_methods: ["iterative_hint_e2e"] + reuse_workflow_instance: false + max_response_tokens_restraint: 1000 + hint_penalty_coef: 0.1 + cod_log_interval: 200 + log_dir: ${oc.env:TRINITY_LOG_DIR,./logs/research_cod} + exp_name: ${project}-${group}-${name} + inject_pack_seed: true + pad_tasks_to_full_pack: true +algorithm: + algorithm_type: multi_step_grpo + sample_strategy: staleness_control + sample_strategy_args: + max_staleness: 20 + repeat_times: 8 + optimizer: + lr: 1e-6 + advantage_fn: cod + advantage_fn_args: + std_cal_level: "none" + iterative_hint_e2e_causal: true + red_weight_adaptive_temp: true + red_weight_adaptive_version: "bisection" # "halving" / "bisection" / "step_function" + policy_loss_fn: rec + policy_loss_fn_args: # REC-OneSide-NoIS + epsilon_low: 0.2 + epsilon_high: 0.2 + clip_mode: "one-side" + weight: "none" + kl_loss_fn: k2 + kl_loss_fn_args: + kl_coef: 0.0 +model: + model_path: ${oc.env:TRINITY_MODEL_PATH} + max_prompt_tokens: 10000 + max_response_tokens: 2000 + max_model_len: 12000 +cluster: + node_num: 3 + gpu_per_node: 8 +buffer: + total_epochs: 1 + batch_size: 32 + train_batch_size: 1440 + explorer_input: + tasksets: + - name: frozenlake_obscure + storage_type: file + path: "examples/research_cod/data/frozen_lake_4567" + split: train + format: + prompt_key: 'task_desc' + workflow_args: + context_compression_mode: "keep_all" + # length penalty in reward design + length_penalty_coef: 0.1 + len_full_penalty: 2000000 + len_zero_penalty: 1000000 + # specific to frozenlake-obscure + env_max_steps: 6 + agent_max_steps: 8 + mapping_mode: "per_pack" + rollout_args: + temperature: 1.0 + default_workflow_type: 'cod_frozenlake_obscure_workflow' + trainer_input: + experience_buffer: + name: buffer + storage_type: queue + max_read_timeout: 72000 + replay_buffer: + enable: true + priority_fn: decay_limit_randomization + priority_fn_args: + sigma: 2.0 +explorer: + runner_per_model: 32 + max_repeat_times_per_runner: 1 + max_timeout: 72000 + rollout_model: + enable_thinking: false + enable_history: false + engine_num: 8 + tensor_parallel_size: 1 + enable_prefix_caching: false + enforce_eager: false + dtype: bfloat16 + seed: 42 +trainer: + save_interval: 100 + grad_clip: 1.0 + use_dynamic_bsz: true + max_token_len_per_gpu: 16384 + ulysses_sequence_parallel_size: 1 + trainer_config: + actor_rollout_ref: + actor: + checkpoint: + load_contents: ['model', 'extra'] + save_contents: ['model', 'extra'] +synchronizer: + sync_style: dynamic_by_explorer + sync_method: 'nccl' + sync_interval: 4 + sync_timeout: 72000 +monitor: + monitor_type: wandb diff --git a/examples/research_cod/exp_plan_final/train/grid_navigation.yaml b/examples/research_cod/exp_plan_final/train/grid_navigation.yaml new file mode 100644 index 00000000000..cd66264209c --- /dev/null +++ b/examples/research_cod/exp_plan_final/train/grid_navigation.yaml @@ -0,0 +1,122 @@ +# Data prep: +# 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 +project: ${oc.env:TRINITY_PROJECT,trinity-cod-final} +group: "grid_navigation" +name: "test-qwen3-8b" +mode: both +checkpoint_root_dir: ${oc.env:TRINITY_CHECKPOINT_ROOT_DIR,./checkpoints} +continue_from_checkpoint: false +cod: + task_pack_size: 4 + eval_task_pack_size: 8 + cod_workflow_args: + hint_example: false + activated_cod_methods: ["iterative_hint_e2e"] + reuse_workflow_instance: false + max_response_tokens_restraint: 1000 + hint_penalty_coef: 0.1 + cod_log_interval: 200 + log_dir: ${oc.env:TRINITY_LOG_DIR,./logs/research_cod} + exp_name: ${project}-${group}-${name} + inject_pack_seed: true + pad_tasks_to_full_pack: true +algorithm: + algorithm_type: multi_step_grpo + sample_strategy: staleness_control + sample_strategy_args: + max_staleness: 20 + repeat_times: 8 + optimizer: + lr: 1e-6 + lr_warmup_steps: 10 + lr_scheduler_type: constant + advantage_fn: cod + advantage_fn_args: + std_cal_level: "none" + iterative_hint_e2e_causal: true + red_weight_adaptive_temp: true + red_weight_adaptive_version: "bisection" # "halving" / "bisection" / "step_function" + policy_loss_fn: rec + policy_loss_fn_args: # REC-OneSide-NoIS + epsilon_low: 0.2 + epsilon_high: 0.2 + clip_mode: "one-side" + weight: "none" + kl_loss_fn: k2 + kl_loss_fn_args: + kl_coef: 0.0 +model: + model_path: ${oc.env:TRINITY_MODEL_PATH} + max_prompt_tokens: 10000 + max_response_tokens: 2000 + max_model_len: 12000 +cluster: + node_num: 3 + gpu_per_node: 8 +buffer: + total_epochs: 1 + batch_size: 32 + train_batch_size: 1440 + explorer_input: + tasksets: + - name: grid_navigation + storage_type: file + path: "examples/research_cod/data/grid_navigation" + split: train + workflow_args: + context_compression_mode: "keep_all" + length_penalty_coef: 0.0 + grid_min_size: 8 + grid_max_size: 12 + landscape_num_components: 4 + landscape_min_scale: 0.10 + landscape_max_scale: 0.35 + min_rounds: 4 + max_rounds: 6 + reveal_radius: 2 + rollout_args: + temperature: 1.0 + default_workflow_type: 'cod_grid_navigation_workflow' + trainer_input: + experience_buffer: + name: buffer + storage_type: queue + max_read_timeout: 72000 + replay_buffer: + enable: true + priority_fn: decay_limit_randomization + priority_fn_args: + sigma: 2.0 +explorer: + runner_per_model: 32 + max_repeat_times_per_runner: 1 + max_timeout: 72000 + rollout_model: + enable_thinking: false + enable_history: false + engine_num: 8 + tensor_parallel_size: 1 + enable_prefix_caching: false + enforce_eager: false + dtype: bfloat16 + seed: 42 +trainer: + save_interval: 100 + enable_preview: true + grad_clip: 1.0 + use_dynamic_bsz: true + max_token_len_per_gpu: 16384 + ulysses_sequence_parallel_size: 1 + trainer_config: + actor_rollout_ref: + actor: + checkpoint: + load_contents: ['model', 'extra'] + save_contents: ['model', 'extra'] +synchronizer: + sync_style: dynamic_by_explorer + sync_method: 'nccl' + sync_interval: 4 + sync_timeout: 72000 +monitor: + monitor_type: wandb diff --git a/examples/research_cod/exp_plan_final/train/grid_navigation_27b.yaml b/examples/research_cod/exp_plan_final/train/grid_navigation_27b.yaml new file mode 100644 index 00000000000..2aa43344c03 --- /dev/null +++ b/examples/research_cod/exp_plan_final/train/grid_navigation_27b.yaml @@ -0,0 +1,128 @@ +# Data prep: +# 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 +project: ${oc.env:TRINITY_PROJECT,trinity-cod-final} +group: "grid_navigation" +name: "qwen3.6-27b-20260901" +mode: both +checkpoint_root_dir: ${oc.env:TRINITY_CHECKPOINT_ROOT_DIR,./checkpoints} +continue_from_checkpoint: false +cod: + task_pack_size: 4 + eval_task_pack_size: 8 + cod_workflow_args: + hint_example: false + activated_cod_methods: ["iterative_hint_e2e"] + reuse_workflow_instance: false + max_response_tokens_restraint: 4000 + hint_penalty_coef: 0.1 + cod_log_interval: 200 + log_dir: ${oc.env:TRINITY_LOG_DIR,./logs/research_cod} + exp_name: ${project}-${group}-${name} + inject_pack_seed: true + pad_tasks_to_full_pack: true +algorithm: + algorithm_type: multi_step_grpo + sample_strategy: staleness_control + sample_strategy_args: + max_staleness: 20 + repeat_times: 8 + optimizer: + lr: 1e-6 + lr_warmup_steps: 10 + lr_scheduler_type: constant + advantage_fn: cod + advantage_fn_args: + std_cal_level: "none" + iterative_hint_e2e_causal: true + red_weight_adaptive_temp: true + red_weight_adaptive_version: "bisection" # "halving" / "bisection" / "step_function" + policy_loss_fn: rec + policy_loss_fn_args: # REC-OneSide-NoIS + epsilon_low: 0.2 + epsilon_high: 0.2 + clip_mode: "one-side" + weight: "none" + kl_loss_fn: k2 + kl_loss_fn_args: + kl_coef: 0.0 +model: + model_path: ${oc.env:TRINITY_MODEL_PATH} + max_prompt_tokens: 22000 + max_response_tokens: 3999 + max_model_len: 26000 +cluster: + node_num: 3 + gpu_per_node: 8 +buffer: + total_epochs: 1 + batch_size: 32 + train_batch_size: 1440 + explorer_input: + tasksets: + - name: grid_navigation + storage_type: file + path: "examples/research_cod/data/grid_navigation" + split: train + workflow_args: + context_compression_mode: "keep_all" + length_penalty_coef: 0.0 + grid_min_size: 8 + grid_max_size: 12 + landscape_num_components: 4 + landscape_min_scale: 0.15 + landscape_max_scale: 0.45 + min_rounds: 3 + max_rounds: 5 + reveal_radius: 1 + rollout_args: + temperature: 1.0 + default_workflow_type: 'cod_grid_navigation_workflow' + trainer_input: + experience_buffer: + name: buffer + storage_type: queue + max_read_timeout: 72000 + replay_buffer: + enable: true + priority_fn: decay_limit_randomization + priority_fn_args: + sigma: 2.0 +explorer: + runner_per_model: 32 + max_repeat_times_per_runner: 1 + max_timeout: 72000 + rollout_model: + enable_thinking: false + enable_history: false + engine_num: 8 + tensor_parallel_size: 1 + enable_prefix_caching: false + enforce_eager: false + dtype: bfloat16 + seed: 42 +trainer: + save_interval: 20 + enable_preview: true + grad_clip: 1.0 + use_dynamic_bsz: true + use_remove_padding: true + # max_token_len_per_gpu: 16384 + ulysses_sequence_parallel_size: 4 + # fix_actor_microbatch_loss_scale: true + trainer_config: + actor_rollout_ref: + model: + use_fused_kernels: ${oc.decode:${oc.env:FUSED_KERNELS,true}} + fused_kernel_options: + impl_backend: torch + actor: + checkpoint: + load_contents: ['model', 'extra'] + save_contents: ['model', 'extra'] +synchronizer: + sync_style: dynamic_by_explorer + sync_method: 'nccl' + sync_interval: 4 + sync_timeout: 72000 +monitor: + monitor_type: wandb diff --git a/examples/research_cod/exp_plan_final/train/mixed_flobs_alchran.yaml b/examples/research_cod/exp_plan_final/train/mixed_flobs_alchran.yaml new file mode 100644 index 00000000000..a69838b0399 --- /dev/null +++ b/examples/research_cod/exp_plan_final/train/mixed_flobs_alchran.yaml @@ -0,0 +1,155 @@ +# Data prep (joint training on FrozenLake-Obscure + Alchemy-Random): +# python examples/research_cod/get_frozen_lake_data.py --local_dir examples/research_cod/data/frozen_lake_4567 --train_size 50000 --test_size 4000 --map_min_size 4 --map_max_size 5 --tile_min_prob 0.6 --tile_max_prob 0.7 +# python examples/research_cod/get_alchemy_data.py --local_dir examples/research_cod/data/alchemy_random --train_size 50000 --test_size 4000 --seed 42 +project: "trinity-cod-final" +group: "mixed_flobs_alchran" +name: "qwen3-8b-mixed" +mode: both +checkpoint_root_dir: ${oc.env:TRINITY_CHECKPOINT_ROOT_DIR,./checkpoints} +continue_from_checkpoint: false +cod: + task_pack_size: 4 + eval_task_pack_size: 8 + cod_workflow_args: + hint_example: false + activated_cod_methods: ["iterative_hint_e2e"] + reuse_workflow_instance: false + max_response_tokens_restraint: 1000 + hint_penalty_coef: 0.1 + cod_log_interval: 200 + log_dir: ${oc.env:TRINITY_LOG_DIR,./logs/research_cod} + exp_name: ${project}-${group}-${name} + inject_pack_seed: true + pad_tasks_to_full_pack: true +algorithm: + algorithm_type: multi_step_grpo + sample_strategy: staleness_control + sample_strategy_args: + max_staleness: 20 + repeat_times: 8 + optimizer: + lr: 1e-6 + advantage_fn: cod + advantage_fn_args: + std_cal_level: "none" + iterative_hint_e2e_causal: true + red_weight_adaptive_temp: true + red_weight_adaptive_version: "bisection" # "halving" / "bisection" / "step_function" + policy_loss_fn: rec + policy_loss_fn_args: # REC-OneSide-NoIS + epsilon_low: 0.2 + epsilon_high: 0.2 + clip_mode: "one-side" + weight: "none" + kl_loss_fn: k2 + kl_loss_fn_args: + kl_coef: 0.0 +model: + model_path: ${oc.env:TRINITY_MODEL_PATH} + max_prompt_tokens: 10000 + max_response_tokens: 2000 + max_model_len: 12000 +cluster: + node_num: 4 + gpu_per_node: 8 +buffer: + total_epochs: 1 + batch_size: 32 + train_batch_size: 1440 + explorer_input: + tasksets: + - name: frozenlake_obscure + storage_type: file + path: "examples/research_cod/data/frozen_lake_4567" + split: train + format: + prompt_key: 'task_desc' + workflow_args: + context_compression_mode: "keep_all" + # length penalty in reward design + length_penalty_coef: 0.1 + len_full_penalty: 2000000 + len_zero_penalty: 1000000 + # specific to frozenlake-obscure + env_max_steps: 6 + agent_max_steps: 8 + mapping_mode: "per_pack" + rollout_args: + temperature: 1.0 + default_workflow_type: 'cod_frozenlake_obscure_workflow' + - name: alchemy_random + storage_type: file + path: "examples/research_cod/data/alchemy_random" + split: train + workflow_args: + context_compression_mode: "keep_all" + # length penalty in reward design + length_penalty_coef: 0.1 + len_full_penalty: 2000000 # 1500 + len_zero_penalty: 1000000 # 500 + # specific to alchemy-random + max_rounds: 8 + material_mode: "unlimited" + scope: "per_pack" + show_recipes: false + show_elements: false + show_tiers: true + min_num_tiers: 3 + max_num_tiers: 4 + min_base_elements: 3 + max_base_elements: 4 + tier_shrink_min: 0.4 + tier_shrink_max: 0.6 + min_recipes_per_element: 1 + max_recipes_per_element: 1 + max_tier_gap: 1 + min_cross_tier_prob: 0.0 + max_cross_tier_prob: 0.1 + min_noise_node_ratio: 0.0 + max_noise_node_ratio: 0.0 + noise_chain_depth: 2 + rollout_args: + temperature: 1.0 + default_workflow_type: 'cod_random_alchemy_workflow' + trainer_input: + experience_buffer: + name: buffer + storage_type: queue + max_read_timeout: 72000 + replay_buffer: + enable: true + priority_fn: decay_limit_randomization + priority_fn_args: + sigma: 2.0 +explorer: + runner_per_model: 32 + max_repeat_times_per_runner: 1 + max_timeout: 72000 + rollout_model: + enable_thinking: false + enable_history: false + engine_num: 8 + tensor_parallel_size: 1 + enable_prefix_caching: false + enforce_eager: false + dtype: bfloat16 + seed: 42 +trainer: + save_interval: 100 + grad_clip: 1.0 + use_dynamic_bsz: true + max_token_len_per_gpu: 16384 + ulysses_sequence_parallel_size: 1 + trainer_config: + actor_rollout_ref: + actor: + checkpoint: + load_contents: ['model', 'extra'] + save_contents: ['model', 'extra'] +synchronizer: + sync_style: dynamic_by_explorer + sync_method: 'nccl' + sync_interval: 4 + sync_timeout: 72000 +monitor: + monitor_type: wandb diff --git a/examples/research_cod/exp_plan_final/train/optimal_control_improve.yaml b/examples/research_cod/exp_plan_final/train/optimal_control_improve.yaml new file mode 100644 index 00000000000..bfb41fb0466 --- /dev/null +++ b/examples/research_cod/exp_plan_final/train/optimal_control_improve.yaml @@ -0,0 +1,130 @@ +# Data prep: +# 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 +# Before training, install the Qwen3.6 kernels as documented in examples/research_cod/README.md. +# export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True +project: "trinity-cod" +group: "optimal_control" +name: ${oc.env:EXP_NAME,qwen3.6-27b-optimal-control-cod} +mode: both +checkpoint_root_dir: ${oc.env:TRINITY_CHECKPOINT_ROOT_DIR,./checkpoints} +continue_from_checkpoint: false +cod: + task_pack_size: 4 + packing_strategy: cod + cod_workflow_args: + hint_example: false + activated_cod_methods: ["iterative_hint_e2e"] + reuse_workflow_instance: false + max_response_tokens_restraint: 4000 + hint_penalty_coef: 0.1 + cod_log_interval: 50 + log_dir: ${oc.env:TRINITY_LOG_DIR,./logs/research_cod} + exp_name: ${project}-${group}-${name} + inject_pack_seed: true + stable_pack_seed: true + pad_tasks_to_full_pack: true +algorithm: + algorithm_type: multi_step_grpo + sample_strategy: staleness_control + sample_strategy_args: + max_staleness: 20 + repeat_times: 8 + optimizer: + lr: 1e-6 + advantage_fn: cod + advantage_fn_args: + std_cal_level: "none" + iterative_hint_e2e_causal: true + red_weight_adaptive_temp: true + red_weight_adaptive_version: "bisection" + policy_loss_fn: rec + policy_loss_fn_args: + epsilon_low: 0.2 + epsilon_high: 0.2 + clip_mode: "one-side" + weight: "none" + kl_loss_fn: k2 + kl_loss_fn_args: + kl_coef: 0.0 +model: + model_path: ${oc.env:TRINITY_MODEL_PATH} + max_prompt_tokens: 12000 + max_response_tokens: 4000 + max_model_len: 40000 +cluster: + node_num: 3 + gpu_per_node: 8 +buffer: + total_epochs: 1 + batch_size: 32 + train_batch_size: 480 + explorer_input: + tasksets: + - name: optimal_control + storage_type: file + path: "examples/research_cod/data/optimal_control" + split: train + format: + prompt_key: 'task_desc' + workflow_args: + context_compression_mode: "keep_all" + include_previous_control_policy: true + a_env_range: [0.55, 1.20] + b_env_range: [-1.5, 1.5] + min_abs_b_env: 0.35 + enable_process_noise: false + rollout_args: + temperature: 1.0 + data_selector: + selector_type: sequential + default_workflow_type: 'cod_optimalcontrol_workflow' + trainer_input: + experience_buffer: + name: buffer + storage_type: queue + max_read_timeout: 72000 + replay_buffer: + enable: true + priority_fn: decay_limit_randomization + priority_fn_args: + sigma: 2.0 +explorer: + runner_per_model: 32 + max_repeat_times_per_runner: 1 + max_timeout: 72000 + rollout_model: + enable_thinking: false + enable_history: false + engine_num: 8 + tensor_parallel_size: 1 + enable_prefix_caching: false + enforce_eager: false + dtype: bfloat16 + seed: 42 +trainer: + save_interval: 100 + enable_preview: true + grad_clip: 1.0 + use_dynamic_bsz: ${oc.env:USE_DYNAMIC_BSZ,true} + use_remove_padding: ${oc.env:REMOVE_PADDING,true} + ulysses_sequence_parallel_size: ${oc.env:SP_SIZE,4} + fix_actor_microbatch_loss_scale: ${oc.env:FIX_LOSS_SCALE,true} + trainer_config: + actor_rollout_ref: + model: + use_fused_kernels: ${oc.decode:${oc.env:FUSED_KERNELS,true}} + fused_kernel_options: + impl_backend: torch + actor: + checkpoint: + load_contents: ['model', 'extra'] + save_contents: ['model', 'extra'] +synchronizer: + sync_style: dynamic_by_explorer + sync_method: 'nccl' + sync_interval: 4 + sync_timeout: 72000 +monitor: + monitor_type: wandb diff --git a/examples/research_cod/exp_plan_final/train/pde_control_grid_opd.yaml b/examples/research_cod/exp_plan_final/train/pde_control_grid_opd.yaml new file mode 100644 index 00000000000..6d9896b5bae --- /dev/null +++ b/examples/research_cod/exp_plan_final/train/pde_control_grid_opd.yaml @@ -0,0 +1,235 @@ +# Multi-teacher OPD on PDE, optimal-control, and grid-navigation tasksets. +# Data prep: +# python examples/research_cod/get_pde_discovery_data.py --local_dir examples/research_cod/data/pde_discovery_runtime_seed --train_size 50000 --test_size 4000 --seed 42 +# 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 +# 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 +# 教师路径指向已转换的 Hugging Face checkpoint，顺序为 PDE、Optimal Control、Grid Navigation。 +# export TRINITY_PDE_TEACHER_MODEL_PATH=/path/to/pde-teacher/actor/huggingface +# export TRINITY_OPTIMAL_CONTROL_TEACHER_MODEL_PATH=/path/to/control-teacher/actor/huggingface +# export TRINITY_GRID_NAVIGATION_TEACHER_MODEL_PATH=/path/to/grid-teacher/actor/huggingface +project: ${oc.env:OPD_PROJECT,trinity-cod} +group: "mixed_multi_teacher_opd" +name: ${oc.env:EXP_NAME,qwen3.6-27b-mixed-opd-pde75-control200-grid100} +mode: both +checkpoint_root_dir: ${oc.env:TRINITY_CHECKPOINT_ROOT_DIR,./checkpoints} +continue_from_checkpoint: false + +cod: + task_pack_size: 4 + eval_task_pack_size: 8 + packing_strategy: cod + cod_workflow_args: + hint_example: false + activated_cod_methods: ["iterative_hint_e2e"] + reuse_workflow_instance: false + enable_teacher_logprobs: true + max_response_tokens_restraint: 4000 + hint_penalty_coef: 0.1 + cod_log_interval: 200 + log_dir: ${oc.env:TRINITY_LOG_DIR,./logs/research_cod} + exp_name: ${project}-${group}-${name} + inject_pack_seed: true + stable_pack_seed: true + pad_tasks_to_full_pack: true + +algorithm: + algorithm_type: on_policy_distill + sample_strategy: staleness_control + sample_strategy_args: + max_staleness: 5 + repeat_times: 1 + optimizer: + lr: 1e-6 + lr_warmup_steps: 10 + lr_scheduler_type: constant + advantage_fn_args: + kl_coef: 1.0 + policy_loss_fn: rec + policy_loss_fn_args: + clip_mode: "none" + weight: "none" + fix_opd_advantage: true + kl_loss_fn: k2 + kl_loss_fn_args: + kl_coef: 0.0 + entropy_loss_fn: default + +model: + model_path: ${oc.env:TRINITY_MODEL_PATH} + max_prompt_tokens: 22000 + max_response_tokens: 4000 + max_model_len: 26001 + +cluster: + node_num: 3 + gpu_per_node: 8 + +buffer: + total_epochs: 1 + batch_size: 256 + train_batch_size: 720 + explorer_input: + tasksets: + - name: pde_discovery + storage_type: file + path: ${oc.env:TRINITY_PDE_TASKSET_PATH,examples/research_cod/data/pde_discovery_runtime_seed} + split: train + format: + prompt_key: "task_desc" + response_key: "answer" + workflow_args: + update_context_impl: "pde_update_context" + context_compression_mode: "keep_all" + max_steps: 6 + point_budget: 6 + min_reaction_terms: 1 + max_reaction_terms: 2 + ground_truth_family: "physical_full_train.json" + pde_state_abs_limit: 3.0 + initial_condition_shape: + mode_count_range: [1, 2] + noise_level: 0.002 + kappa_threshold: 100.0 + dump_trajectories: false + rollout_args: + temperature: 1.0 + data_selector: + selector_type: sequential + default_workflow_type: "trinity.common.workflows.connect_the_dots.pde_discovery.workflow.CoDPDEDiscoveryWorkflow" + + - name: optimal_control + storage_type: file + path: ${oc.env:TRINITY_OPTIMAL_CONTROL_TASKSET_PATH,examples/research_cod/data/optimal_control} + split: train + format: + prompt_key: "task_desc" + workflow_args: + update_context_impl: "update_context" + context_compression_mode: "keep_all" + include_previous_control_policy: true + a_env_range: [0.55, 1.20] + b_env_range: [-1.5, 1.5] + min_abs_b_env: 0.35 + enable_process_noise: false + rollout_args: + temperature: 1.0 + data_selector: + selector_type: sequential + default_workflow_type: "cod_optimalcontrol_workflow" + + - name: grid_navigation + storage_type: file + path: ${oc.env:TRINITY_GRID_NAVIGATION_TASKSET_PATH,examples/research_cod/data/grid_navigation} + split: train + workflow_args: + update_context_impl: "update_context" + context_compression_mode: "keep_all" + length_penalty_coef: 0.0 + grid_min_size: 8 + grid_max_size: 12 + landscape_num_components: 4 + landscape_min_scale: 0.15 + landscape_max_scale: 0.45 + min_rounds: 3 + max_rounds: 5 + reveal_radius: 1 + rollout_args: + temperature: 1.0 + data_selector: + selector_type: sequential + default_workflow_type: "cod_grid_navigation_workflow" + + trainer_input: + experience_buffer: + name: buffer + storage_type: queue + max_read_timeout: 72000 + +explorer: + runner_per_model: 32 + max_repeat_times_per_runner: 1 + max_timeout: 72000 + rollout_model: + enable_thinking: false + enable_history: false + engine_num: 4 + tensor_parallel_size: 1 + enable_prefix_caching: false + enforce_eager: true + dtype: bfloat16 + seed: 42 + auxiliary_models: + - name: pde_step75_teacher + model_path: ${oc.env:TRINITY_PDE_TEACHER_MODEL_PATH} + engine_type: vllm + engine_num: 4 + tensor_parallel_size: 1 + max_prompt_tokens: 22000 + max_response_tokens: 4000 + max_model_len: 26001 + enable_prefix_caching: false + enforce_eager: true + extra_engine_args: + max_num_seqs: 32 + dtype: bfloat16 + seed: 42 + - name: optimal_control_step200_teacher + model_path: ${oc.env:TRINITY_OPTIMAL_CONTROL_TEACHER_MODEL_PATH} + engine_type: vllm + engine_num: 4 + tensor_parallel_size: 1 + max_prompt_tokens: 22000 + max_response_tokens: 4000 + max_model_len: 26001 + enable_prefix_caching: false + enforce_eager: true + extra_engine_args: + max_num_seqs: 32 + dtype: bfloat16 + seed: 42 + - name: grid_navigation_step100_teacher + model_path: ${oc.env:TRINITY_GRID_NAVIGATION_TEACHER_MODEL_PATH} + engine_type: vllm + engine_num: 4 + tensor_parallel_size: 1 + max_prompt_tokens: 22000 + max_response_tokens: 4000 + max_model_len: 26001 + enable_prefix_caching: false + enforce_eager: true + extra_engine_args: + max_num_seqs: 32 + dtype: bfloat16 + seed: 42 + +trainer: + total_steps: 200 + save_interval: 5 + enable_preview: true + grad_clip: 1.0 + use_dynamic_bsz: ${oc.env:USE_DYNAMIC_BSZ,true} + max_token_len_per_gpu: 7500 + use_remove_padding: ${oc.env:REMOVE_PADDING,true} + ulysses_sequence_parallel_size: 4 + offload_policy: false + # Keep enabled for this run; revisit after inspecting the initial curves. + fix_actor_microbatch_loss_scale: ${oc.env:FIX_LOSS_SCALE,true} + trainer_config: + actor_rollout_ref: + model: + use_fused_kernels: ${oc.decode:${oc.env:FUSED_KERNELS,true}} + fused_kernel_options: + impl_backend: torch + actor: + checkpoint: + load_contents: ["model", "extra"] + save_contents: ["model", "extra"] + +synchronizer: + sync_style: dynamic_by_explorer + sync_method: nccl + sync_interval: 1 + sync_timeout: 72000 + +monitor: + monitor_type: wandb diff --git a/examples/research_cod/exp_plan_final/train/pde_discovery_cod_600steps_hard.yaml b/examples/research_cod/exp_plan_final/train/pde_discovery_cod_600steps_hard.yaml new file mode 100644 index 00000000000..b7f7422a071 --- /dev/null +++ b/examples/research_cod/exp_plan_final/train/pde_discovery_cod_600steps_hard.yaml @@ -0,0 +1,180 @@ +# Hard, step-bounded rerun of PDE CoD training. +# +# The 50,000-row training split contains about 1,562 explorer batches at +# batch_size=32, so 600 explorer steps consume only 19,200 raw tasks (about +# 38% of one epoch). Based on the previous run's convergence around the 100-step +# region, save every 25 trainer steps and evaluate the saved checkpoints with +# the hard PDE bench configs. Explorer/W&B rollout steps and trainer/checkpoint +# global steps are separate asynchronous counters and are not guaranteed to match. +# +# Data prep: +# 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 +# Before training, install the Qwen3.6 kernels as documented in examples/research_cod/README.md. +# export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True +project: "trinity-cod" +group: "pde_discovery" +name: ${oc.env:EXP_NAME,qwen3.6-27b-pde-discovery-cod-600steps-hard} +mode: both +checkpoint_root_dir: ${oc.env:TRINITY_CHECKPOINT_ROOT_DIR,./checkpoints} +continue_from_checkpoint: false + +cod: + task_pack_size: 4 + eval_task_pack_size: 8 + packing_strategy: cod + cod_workflow_args: + hint_example: false + activated_cod_methods: ["iterative_hint_e2e"] + update_context_impl: "pde_update_context" + reuse_workflow_instance: false + max_response_tokens_restraint: 4000 + hint_penalty_coef: 0.1 + cod_log_interval: 200 + log_dir: ${oc.env:TRINITY_LOG_DIR,./logs/research_cod} + exp_name: ${project}-${group}-${name} + inject_pack_seed: true + # Derive the runtime pack seed from dataset row identities. This preserves + # flexible runtime packing while keeping each eval pack's hidden PDE fixed + # across checkpoints, making the W&B eval curve directly comparable. + stable_pack_seed: true + pad_tasks_to_full_pack: true + +algorithm: + algorithm_type: multi_step_grpo + sample_strategy: staleness_control + sample_strategy_args: + max_staleness: 20 + repeat_times: 8 + optimizer: + lr: 1e-6 + advantage_fn: cod + advantage_fn_args: + std_cal_level: "none" + iterative_hint_e2e_causal: true + red_weight_adaptive_temp: true + red_weight_adaptive_version: "bisection" + policy_loss_fn: rec + policy_loss_fn_args: + epsilon_low: 0.2 + epsilon_high: 0.2 + clip_mode: "one-side" + weight: "none" + kl_loss_fn: k2 + kl_loss_fn_args: + kl_coef: 0.0 + # Compute and log policy entropy for collapse monitoring, but keep its + # optimization coefficient at zero so it does not change the training objective. + entropy_loss_fn: default + entropy_loss_fn_args: + entropy_coef: 0.0 + +model: + model_path: ${oc.env:TRINITY_MODEL_PATH} + # Prior PDE rollouts peaked around 16,111 prompt tokens. Keep prompt headroom + # and restore the 4,000-token response budget so the model has enough space + # to finish the mandatory XML action after its reasoning. + max_prompt_tokens: 18000 + max_response_tokens: 4000 + max_model_len: 22001 + +cluster: + node_num: 3 + gpu_per_node: 8 + +buffer: + # Caps task production at 600 explorer/rollout batches. In W&B this is the + # step axis used by rollout/* metrics. It takes precedence over + # total_epochs and is intentionally below one epoch. + total_steps: 600 + batch_size: 32 + train_batch_size: 1440 + explorer_input: + tasksets: + - name: pde_discovery + storage_type: file + # A distinct directory prevents accidentally reusing legacy rows that + # already contain pack_seed/task_idx/pack_size. + path: ${oc.env:TRINITY_PDE_TASKSET_PATH,examples/research_cod/data/pde_discovery_runtime_seed} + split: train + format: + prompt_key: 'task_desc' + response_key: 'answer' + workflow_args: + context_compression_mode: "keep_all" + max_steps: 6 + point_budget: 6 + min_reaction_terms: 1 + max_reaction_terms: 2 + ground_truth_family: "physical_full_train.json" + pde_state_abs_limit: 3.0 + initial_condition_shape: + mode_count_range: [1, 2] + noise_level: 0.002 + kappa_threshold: 100.0 + dump_trajectories: false + rollout_args: + temperature: 1.0 + data_selector: + selector_type: sequential + default_workflow_type: 'trinity.common.workflows.connect_the_dots.pde_discovery.workflow.CoDPDEDiscoveryWorkflow' + trainer_input: + experience_buffer: + name: buffer + storage_type: queue + max_read_timeout: 72000 + replay_buffer: + enable: true + priority_fn: decay_limit_randomization + priority_fn_args: + sigma: 2.0 + +explorer: + runner_per_model: 32 + max_repeat_times_per_runner: 1 + max_timeout: 72000 + rollout_model: + enable_thinking: false + enable_history: false + engine_num: 8 + tensor_parallel_size: 1 + enable_prefix_caching: false + enforce_eager: false + dtype: bfloat16 + seed: 42 + +trainer: + # This interval is measured in trainer global steps, not explorer steps. + # A final checkpoint is saved when training drains, and the default "last" + # policy exports that final checkpoint in Hugging Face format. + save_interval: 25 + enable_preview: true + grad_clip: 1.0 + use_dynamic_bsz: ${oc.env:USE_DYNAMIC_BSZ,true} + # A maximum-length sequence contributes ceil(22001 / SP4)=5501 tokens per + # trainer GPU. Sixteen trainer GPUs without offload peaked at about 51 GB + # with a 6500-token budget, leaving room for a larger dynamic micro-batch. + # The larger budget can reduce the micro-batch count and can be overridden + # at launch if later batches show a higher memory peak. + max_token_len_per_gpu: ${oc.env:MAX_TOKEN_LEN_PER_GPU,13000} + use_remove_padding: ${oc.env:REMOVE_PADDING,true} + ulysses_sequence_parallel_size: ${oc.env:SP_SIZE,4} + fix_actor_microbatch_loss_scale: ${oc.env:FIX_LOSS_SCALE,true} + trainer_config: + actor_rollout_ref: + model: + use_fused_kernels: ${oc.decode:${oc.env:FUSED_KERNELS,true}} + fused_kernel_options: + impl_backend: torch + actor: + checkpoint: + load_contents: ['model', 'extra'] + save_contents: ['model', 'extra'] + +synchronizer: + sync_style: dynamic_by_explorer + sync_method: 'nccl' + sync_interval: 4 + sync_timeout: 72000 + +monitor: + monitor_type: wandb diff --git a/examples/research_cod/exp_plan_learn2ask/data_prepare/1_info_extract_pipeline.py b/examples/research_cod/exp_plan_learn2ask/data_prepare/1_info_extract_pipeline.py new file mode 100644 index 00000000000..9fff6a9020d --- /dev/null +++ b/examples/research_cod/exp_plan_learn2ask/data_prepare/1_info_extract_pipeline.py @@ -0,0 +1,288 @@ +import argparse +import json +import os +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parents[4])) + +from examples.learn_to_ask.data_prepare.message_splitter import split_session_to_json_lines +from examples.research_cod.exp_plan_learn2ask.data_prepare.llm_info_extraction import ( + LLM_info_extraction_batch, + parse_llm_output, +) + + +def process_jsonl_file( + input_file, output_file, model_call_mode="online_api", max_retries=3, **kwargs +): + """Batched info extraction: collect every cid's prompt from every + session, feed the whole lot to the backend in one llm.generate call + (for local_vllm) so vLLM's continuous batching can do the work. + + Failed cids are retried in additional batched rounds; at most + max_retries + 1 rounds total. Surviving failures get info_set=None + (2_build_dataset drops those with decision=continue). + + Args: + input_file (str): path to input jsonl (one session per line). + output_file (str): path to output jsonl (one cid per line, with + `info_set` filled in). + model_call_mode (str): "online_api" or "local_vllm". + max_retries (int): extra rounds of batched retry on top of the + first attempt. + **kwargs: forwarded to the backend (model_path, enable_thinking, + tensor_parallel_size, etc.). + """ + # -- Load all sessions -- + sessions = [] + with open(input_file, "r", encoding="utf-8") as infile: + for line_num, line in enumerate(infile, 1): + line = line.strip() + if not line: + continue + try: + sessions.append(json.loads(line)) + except json.JSONDecodeError as e: + print(f"Warning: Skipping invalid JSON at line {line_num}: {e}") + print(f"[load] {len(sessions)} sessions from {input_file}") + + # -- Split every session into cids, keep flat + nested views -- + per_session_cids = [] # list[list[dict]] preserving session order + flat_data = [] # flat list of cid dicts (same objects as nested) + remaining_chats = [] # flat list of remaining_chat strings + + for session in sessions: + session_cids = [] + for cid_json in split_session_to_json_lines(session): + data = json.loads(cid_json) + data.setdefault("info_set", None) + session_cids.append(data) + flat_data.append(data) + remaining_chats.append(data.get("remaining_chat", "")) + per_session_cids.append(session_cids) + + total = len(flat_data) + print(f"[flatten] {total} cids across {len(sessions)} sessions") + if total == 0: + # Still emit an empty output so downstream doesn't trip over a missing file. + open(output_file, "w", encoding="utf-8").close() + return f"Nothing to process; wrote empty file {output_file}" + + # -- Batched generation with multi-round retry -- + pending = list(range(total)) + for round_idx in range(max_retries + 1): + if not pending: + break + print( + f"[batch] round {round_idx + 1}/{max_retries + 1}: " + f"generating for {len(pending)} cids" + ) + chats_to_run = [remaining_chats[i] for i in pending] + outputs = LLM_info_extraction_batch( + chats_to_run, model_call_mode, **kwargs + ) + + still_pending = [] + for orig_idx, text in zip(pending, outputs): + parsed = parse_llm_output(text) + if isinstance(parsed, list): + flat_data[orig_idx]["info_set"] = parsed + else: + still_pending.append(orig_idx) + print( + f"[batch] round {round_idx + 1}: " + f"{len(pending) - len(still_pending)} succeeded, " + f"{len(still_pending)} still pending" + ) + pending = still_pending + + if pending: + print( + f"[batch] WARNING: {len(pending)} cids still failed after " + f"{max_retries + 1} rounds; leaving info_set=None" + ) + + # -- Write output in session-major order (session 0's cids, then 1's, ...) -- + with open(output_file, "w", encoding="utf-8") as outf: + for session_cids in per_session_cids: + for data in session_cids: + outf.write(json.dumps(data, ensure_ascii=False) + "\n") + + return f"Successfully processed {total} cids. Results saved to {output_file}" + + +# Example usage: +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "--input_file", type=str, default="examples/learn_to_ask/data_raw/train_origin.jsonl" + ) + parser.add_argument( + "--output_file", type=str, default="examples/learn_to_ask/data_raw/train_processed.jsonl" + ) + parser.add_argument( + "--model_call_mode", type=str, choices=["online_api", "local_vllm"], default="local_vllm" + ) + parser.add_argument("--model_path", type=str, required=True) + parser.add_argument( + "--enable_thinking", + choices=["true", "false"], + default=None, + help=( + "For Qwen3-family local_vllm runs, explicitly toggle the chat " + "template's thinking mode. Omit to use the model default (which " + "is True for Qwen3, producing ... output that " + "breaks ast.literal_eval downstream — pass 'false' for extraction)." + ), + ) + parser.add_argument( + "--dataset_id", + type=str, + default=None, + help=( + "Dataset id on HF or ModelScope (e.g. 'JasonHaggard/RealMedConv'). When set, " + "--input_file is treated as a *destination path* for dumping the " + "HF split to jsonl before extraction, instead of an existing file." + ), + ) + parser.add_argument( + "--hf_split", + type=str, + default=None, + help="HF split name (required with --dataset_id), e.g. 'train' / 'test'. " + "Use --hf_splits instead to process several splits with a single engine.", + ) + parser.add_argument( + "--hf_splits", + nargs="+", + default=None, + help=( + "Multiple HF split names (e.g. 'train test') to process in ONE " + "process so vLLM engine is initialised only once. Overrides " + "--hf_split. When --dataset_id is given but neither --hf_split " + "nor --hf_splits is set, defaults to ['train', 'test']." + ), + ) + parser.add_argument( + "--output_template", + type=str, + default="examples/research_cod/data/learn2ask_artifacts/{split}_processed.jsonl", + help=( + "Output path template used in multi-split mode; must contain the " + "literal '{split}'. Default: " + "examples/research_cod/data/learn2ask_artifacts/{split}_processed.jsonl. " + "Kept out of examples/research_cod/data/learn2ask/ because HF " + "datasets 4.x treats that dir as the training base_path and its " + "recursive ** glob would sweep every train_*/test_*.jsonl into " + "the respective split." + ), + ) + parser.add_argument( + "--max_retries", + type=int, + default=3, + help="Extra rounds of batched retry on top of the first attempt (default: 3).", + ) + parser.add_argument( + "--tensor_parallel_size", + type=int, + default=1, + help="GPUs to shard one model across (vLLM tensor parallelism). For " + "single-process offline use, set to your GPU count (e.g. 8) to use all " + "cards; default: 1.", + ) + parser.add_argument( + "--data_parallel_size", + type=int, + default=1, + help="vLLM data parallelism (independent replicas). The single-process " + "offline LLM only supports 1; >1 raises ValueError and needs the " + "multi-process launcher. Default: 1.", + ) + parser.add_argument( + "--dataset_source", + choices=["hf", "modelscope"], + default="modelscope", + help="Download --dataset_id from HuggingFace or ModelScope (default: modelscope).", + ) + args = parser.parse_args() + + extra_kwargs = { + "tensor_parallel_size": args.tensor_parallel_size, + "data_parallel_size": args.data_parallel_size, + } + if args.enable_thinking is not None: + extra_kwargs["enable_thinking"] = args.enable_thinking == "true" + + # Default to ['train', 'test'] in one shot when the caller gave + # --dataset_id but did not pin down a single split. + if ( + args.dataset_id is not None + and args.hf_splits is None + and args.hf_split is None + ): + args.hf_splits = ["train", "test"] + + # -- Multi-split HF mode: one engine for all splits. -- + # Inputs (HF split dumps) go next to the outputs so the data prep + # directory holds everything. + if args.hf_splits: + assert args.dataset_id is not None, "--hf_splits requires --dataset_id" + assert ( + args.output_template and "{split}" in args.output_template + ), "--hf_splits requires --output_template containing '{split}'" + for split in args.hf_splits: + output_path = args.output_template.replace("{split}", split) + input_path = os.path.join( + os.path.dirname(os.path.abspath(output_path)), + f"{split}_raw.jsonl", + ) + os.makedirs(os.path.dirname(os.path.abspath(input_path)), exist_ok=True) + if args.dataset_source == "modelscope": + from modelscope.msdatasets import MsDataset # lazy import + ds = MsDataset.load(args.dataset_id, split=split) + n_rows = 0 + with open(input_path, "w", encoding="utf-8") as wf: + for item in ds: + wf.write(json.dumps(dict(item), ensure_ascii=False) + "\n") + n_rows += 1 + else: + from datasets import load_dataset # lazy import + ds = load_dataset(args.dataset_id, split=split) + ds.to_json(input_path, lines=True, force_ascii=False) + n_rows = len(ds) + print(f"[{args.dataset_source}] split={split}: {n_rows} rows -> {input_path}") + print( + process_jsonl_file( + input_file=input_path, + output_file=output_path, + model_call_mode=args.model_call_mode, + model_path=args.model_path, + max_retries=args.max_retries, + **extra_kwargs, + ) + ) + else: + # -- Single-split path (backward compatible). -- + if args.dataset_id is not None: + assert args.hf_split is not None, "--hf_split is required with --dataset_id" + from datasets import load_dataset # lazy import + ds = load_dataset(args.dataset_id, split=args.hf_split) + os.makedirs(os.path.dirname(os.path.abspath(args.input_file)), exist_ok=True) + ds.to_json(args.input_file, lines=True, force_ascii=False) + print( + f"[hf] loaded {len(ds)} rows from {args.dataset_id}:{args.hf_split} " + f"and dumped to {args.input_file}" + ) + + print( + process_jsonl_file( + input_file=args.input_file, + output_file=args.output_file, + model_call_mode=args.model_call_mode, + model_path=args.model_path, + max_retries=args.max_retries, + **extra_kwargs, + ) + ) diff --git a/examples/research_cod/exp_plan_learn2ask/data_prepare/2_build_dataset.py b/examples/research_cod/exp_plan_learn2ask/data_prepare/2_build_dataset.py new file mode 100644 index 00000000000..892784bb579 --- /dev/null +++ b/examples/research_cod/exp_plan_learn2ask/data_prepare/2_build_dataset.py @@ -0,0 +1,85 @@ +import argparse +import json +import os + + +def process_message(json_obj): + info_set = json_obj.get("info_set") + info_set_str = ", ".join(info_set) if isinstance(info_set, list) else "" + if "user: " not in json_obj["remaining_chat"]: + decision_str = "stop" + else: + decision_str = "continue" + # KNOWN BIAS: drops 2607 continue tasks (extractor ignores medication / allergy / chronic). + if not info_set_str and decision_str == "continue": + if_keep = False + else: + if_keep = True + return if_keep, info_set_str, decision_str + + +def main(input_file_path, output_file_path): + # Drop user-ending sessions (no doctor reply -> no ground-truth final decision) + # and ghost rounds (remaining_chat empty -> decision point already passed). + by_session = {} + with open(input_file_path, "r", encoding="utf-8") as infile: + for line in infile: + data = json.loads(line.strip()) + by_session.setdefault(data["session_id"], []).append(data) + + n_sess_dropped = n_ghost_dropped = n_info_dropped = n_kept = 0 + with open(output_file_path, "w", encoding="utf-8") as outfile: + for sid, rounds in by_session.items(): + rounds.sort(key=lambda r: r.get("round_number", 0)) + last_msgs = rounds[-1].get("messages") or [] + if last_msgs and last_msgs[-1].get("role") == "user": + n_sess_dropped += 1 + continue + for data in rounds: + if not data.get("remaining_chat"): + n_ghost_dropped += 1 + continue + if_keep, info_set, decision = process_message(data) + if not if_keep: + n_info_dropped += 1 + continue + + new_item = { + "cid": data["cid"], + "session_id": data["session_id"], + "diagn": data["diagn"], + "messages": data["messages"], + "decision_truth": decision, + "info_truth": info_set, + } + outfile.write(json.dumps(new_item, ensure_ascii=False) + "\n") + n_kept += 1 + print( + f"job done! kept={n_kept}, dropped: session={n_sess_dropped}, " + f"ghost_round={n_ghost_dropped}, no_info={n_info_dropped}" + ) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + + parser.add_argument( + "--input_template", + type=str, + default="examples/research_cod/data/learn2ask_artifacts/{split}_processed.jsonl", + ) + parser.add_argument( + "--output_template", + type=str, + default="examples/research_cod/data/learn2ask/{split}.jsonl", + ) + parser.add_argument("--splits", nargs="+", default=["train", "test"]) + + args = parser.parse_args() + + for split in args.splits: + input_path = args.input_template.replace("{split}", split) + output_path = args.output_template.replace("{split}", split) + os.makedirs(os.path.dirname(os.path.abspath(output_path)), exist_ok=True) + print(f"[{split}] {input_path} -> {output_path}") + main(input_path, output_path) diff --git a/examples/research_cod/exp_plan_learn2ask/data_prepare/llm_info_extraction.py b/examples/research_cod/exp_plan_learn2ask/data_prepare/llm_info_extraction.py new file mode 100644 index 00000000000..e218b06b083 --- /dev/null +++ b/examples/research_cod/exp_plan_learn2ask/data_prepare/llm_info_extraction.py @@ -0,0 +1,199 @@ +import ast +import os +import re +import traceback +from typing import List + +import openai + +llm = None + + +# System prompt shared by the single-shot and batched entry points. +_SYSTEM_MESSAGE = """ + # Task: + You are a medical information assistant. Given a dialogue between a physician (assistant) and a patient (user), extract the clinical attributes of interest to the physician based on their questions. The target fields include: symptom, symptom nature, symptom location, symptom severity, and symptom trigger. Then, identify the corresponding specific information from the patient's responses and pair it with the respective field. + # Requirements: + - Do not fabricate information or introduce new fields not listed above. Ignore patient-reported information regarding prior medication use, allergies, or underlying comorbidities; do not include such details in the output. + - Only include fields explicitly inquired about by the physician. Omit any fields not addressed in the dialogue. Avoid outputting vague terms (e.g., "unspecified" or "unknown"). + - Prevent duplication: if a symptom description already includes anatomical location, do not separately list the location field. + - Format each entry as a string enclosed in single quotes ('), separate multiple entries with commas, and enclose them in square brackets to form a list. Prefix the list with "output: " on the final line. If the dialogue is unrelated to the aforementioned clinical attributes, output "output: []". + - Do not include any commentary after the "output:" line. Condense colloquial patient expressions into concise, standardized, and clinically appropriate terminology. + # Example output format: + output: ['symptom: diarrhea', 'symptom nature: watery stool', 'symptom severity: 4-5 times per day'] + """ + + +def _build_messages(remaining_chat: str) -> list: + return [ + {"role": "system", "content": _SYSTEM_MESSAGE}, + {"role": "user", "content": "```\n" + remaining_chat + "\n```\n"}, + ] + + +def LLM_info_extraction(remaining_chat, model_call_mode, **kwargs): + """ + Extract information from a single remaining_chat using LLM. + + Kept for backward compatibility and ad-hoc debugging; the main + pipeline uses LLM_info_extraction_batch for throughput. + """ + messages = _build_messages(remaining_chat) + try: + if model_call_mode == "online_api": + return _call_online_api(messages, **kwargs) + elif model_call_mode == "local_vllm": + return _call_local_vllm(messages, **kwargs) + else: + return f"Error: Invalid model_call_mode '{model_call_mode}'. Must be 'online_api' or 'local_vllm'." + except Exception as e: + return f"Error occurred: {str(e)}" + + +def LLM_info_extraction_batch( + remaining_chats: List[str], model_call_mode: str, **kwargs +) -> List[str]: + """Batched counterpart of LLM_info_extraction. + + Builds the prompt for every remaining_chat, then issues a single + llm.generate(all_prompts) call so vLLM's continuous batching can + fully utilise the GPUs. Returns a list of response texts, one per + input chat, in the same order. + + online_api mode has no true batching (DashScope etc. are per-call), + so it just falls back to a sequential loop. + """ + messages_list = [_build_messages(rc) for rc in remaining_chats] + + if model_call_mode == "local_vllm": + return _call_local_vllm_batch(messages_list, **kwargs) + if model_call_mode == "online_api": + return [_call_online_api(m, **kwargs) for m in messages_list] + err = f"Error: Invalid model_call_mode '{model_call_mode}'. Must be 'online_api' or 'local_vllm'." + return [err] * len(remaining_chats) + + +def _call_online_api(messages, **kwargs): + """Handle OpenAI-style API calls""" + # Extract API parameters from kwargs or use defaults + api_key = kwargs.get("api_key", os.getenv("DASHSCOPE_API_KEY")) + api_base = kwargs.get("api_base", "https://dashscope.aliyuncs.com/compatible-mode/v1") + model = kwargs.get("model", "qwen2.5-72b-instruct") + temperature = kwargs.get("temperature", 0.7) + max_tokens = kwargs.get("max_tokens", 500) + + client = openai.OpenAI(api_key=api_key, base_url=api_base) + response = client.chat.completions.create( + model=model, messages=messages, temperature=temperature, max_tokens=max_tokens + ) + + return response.choices[0].message.content + + +def _vllm_setup(**kwargs): + """Lazy-init (and cache) the vLLM engine. Returns the global llm and + the SamplingParams to use.""" + from vllm import LLM, SamplingParams + + model_path = kwargs.get("model_path") + if not model_path: + raise ValueError("model_path is required for local vLLM inference") + + # Qwen3.6 thinking general, official: temp 1.0 / top_p 0.95 / top_k 20 / min_p 0 / + # presence_penalty 0. max_tokens 32768 per the card — needs room before "output:". + temperature = 1.0 + top_p = 0.95 + top_k = 20 + min_p = 0.0 + presence_penalty = 0.0 + max_tokens = 32768 + + tensor_parallel_size = kwargs.get("tensor_parallel_size", 1) + data_parallel_size = kwargs.get("data_parallel_size", 1) + gpu_memory_utilization = kwargs.get("gpu_memory_utilization", 0.9) + enforce_eager = kwargs.get("enforce_eager", False) + dtype = kwargs.get("dtype", "auto") + max_model_len = 40960 # 32768 output + headroom for the dialogue prompt + + global llm + if llm is None: + llm = LLM( + model=model_path, + tensor_parallel_size=tensor_parallel_size, + data_parallel_size=data_parallel_size, + gpu_memory_utilization=gpu_memory_utilization, + enforce_eager=enforce_eager, + dtype=dtype, + max_model_len=max_model_len, + ) + + sampling_params = SamplingParams( + temperature=temperature, + top_p=top_p, + top_k=top_k, + min_p=min_p, + presence_penalty=presence_penalty, + max_tokens=max_tokens, + ) + return llm, sampling_params + + +def _call_local_vllm(messages, **kwargs): + """Single-conversation local vLLM call via llm.chat: the engine renders + the chat template, and enable_thinking is passed via chat_template_kwargs.""" + try: + llm, sampling_params = _vllm_setup(**kwargs) + et = kwargs.get("enable_thinking") + ctk = {"enable_thinking": et} if et is not None else None + outputs = llm.chat(messages, sampling_params, chat_template_kwargs=ctk) + return outputs[0].outputs[0].text + except ImportError: + return "Error: vLLM library not installed. Please install it with 'pip install vllm'" + except Exception as e: + traceback.print_exc() # surface the real error instead of silently swallowing it + return f"Error in local vLLM inference: {str(e)}" + + +def _call_local_vllm_batch(messages_list: List[list], **kwargs) -> List[str]: + """Batched local vLLM call: hand all conversations to one llm.chat so + vLLM continuous-batches internally. Returns one output text per input, + in order. enable_thinking is passed via chat_template_kwargs. + """ + if not messages_list: + return [] + try: + llm, sampling_params = _vllm_setup(**kwargs) + et = kwargs.get("enable_thinking") + ctk = {"enable_thinking": et} if et is not None else None + outputs = llm.chat(messages_list, sampling_params, chat_template_kwargs=ctk) + return [o.outputs[0].text for o in outputs] + except ImportError: + return ["Error: vLLM library not installed. Please install it with 'pip install vllm'"] * len(messages_list) + except Exception as e: + traceback.print_exc() # surface the real error instead of silently swallowing it + return [f"Error in local vLLM inference: {str(e)}"] * len(messages_list) + + +def parse_llm_output(output_str): + """ + Convert the LLM info extraction output string to a list of strings. + + Args: + output_str (str): model output ending with "output: [...]" (a bracketed + list of "field: value" strings; "output: []" if nothing relevant). + + Returns: + list: List of strings if successful, error message string if failed + """ + # Lock onto the "output: [...]" format; take the LAST match so any + # reasoning or prose before it is ignored. + matches = re.findall(r"output:\s*(\[[^][]*\])", output_str, flags=re.DOTALL | re.IGNORECASE) + if not matches: + return f"Error parsing output: no 'output: [...]' found in [{repr(output_str)}]" + try: + result = ast.literal_eval(matches[-1]) + if not isinstance(result, list): + return f"Error: Expected a list, got {type(result)}" + return result + except Exception as e: + return f"Error parsing output: [{repr(matches[-1])}] error = {str(e)}" diff --git a/examples/research_cod/exp_plan_learn2ask/data_prepare/make_ood_split.py b/examples/research_cod/exp_plan_learn2ask/data_prepare/make_ood_split.py new file mode 100644 index 00000000000..2bddc0bb970 --- /dev/null +++ b/examples/research_cod/exp_plan_learn2ask/data_prepare/make_ood_split.py @@ -0,0 +1,139 @@ +# -*- coding: utf-8 -*- +"""Re-split learn2ask train+test into an OOD split by disease (diagn). + +Merges the existing train.jsonl + test.jsonl produced by +2_build_dataset.py, then routes whole diagn groups into the new +test split so train/test diagns are disjoint. Default picks diagns +greedily (random seeded order) until the cumulative sample count +reaches `--test_ratio * total`; pass `--test_diagns` to override +with an explicit list. + +Prints the diagn distribution at the start of every run so the +ratio target / explicit list can be sanity-checked against actual +data. +""" +import argparse +import json +import os +import random +from collections import Counter +from typing import Dict, List + + +def load_jsonl(path: str) -> List[dict]: + with open(path, "r", encoding="utf-8") as f: + return [json.loads(line) for line in f if line.strip()] + + +def write_jsonl(path: str, samples: List[dict]) -> None: + os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True) + with open(path, "w", encoding="utf-8") as f: + for s in samples: + f.write(json.dumps(s, ensure_ascii=False) + "\n") + + +def pick_test_diagns_by_ratio( + diagn_counts: Dict[str, int], + target_samples: int, + seed: int, +) -> List[str]: + """Greedy: shuffle diagns under `seed`, append until cumulative + sample count reaches the target. Overshoot up to one diagn is + accepted — use --test_diagns for finer control.""" + diagns = list(diagn_counts.keys()) + random.Random(seed).shuffle(diagns) + chosen, acc = [], 0 + for d in diagns: + if acc >= target_samples: + break + chosen.append(d) + acc += diagn_counts[d] + return chosen + + +def main(): + parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + parser.add_argument( + "--input_train", + default="examples/research_cod/data/learn2ask/train.jsonl", + ) + parser.add_argument( + "--input_test", + default="examples/research_cod/data/learn2ask/test.jsonl", + ) + parser.add_argument( + "--output_train", + default="examples/research_cod/data/learn2ask_ood/train.jsonl", + ) + parser.add_argument( + "--output_test", + default="examples/research_cod/data/learn2ask_ood/test.jsonl", + ) + parser.add_argument( + "--test_ratio", + type=float, + default=0.2, + help="Target fraction of samples routed to test, picked by diagn group.", + ) + parser.add_argument( + "--test_diagns", + type=str, + default=None, + help="Comma-separated explicit diagn list. Overrides --test_ratio.", + ) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument( + "--show_dist_only", + action="store_true", + help="Print the diagn distribution and exit without writing splits.", + ) + args = parser.parse_args() + + train_in = load_jsonl(args.input_train) + test_in = load_jsonl(args.input_test) + samples = train_in + test_in + print( + f"Loaded {len(samples)} samples " + f"({len(train_in)} train + {len(test_in)} test) from {args.input_train} / {args.input_test}" + ) + + counts = Counter(s.get("diagn", "(missing)") for s in samples) + total = sum(counts.values()) + + print(f"\nDiagn distribution ({len(counts)} unique diagns):") + for d, n in counts.most_common(): + print(f" {n:6d} ({100 * n / total:5.2f}%) {d}") + + if args.show_dist_only: + return + + if args.test_diagns: + test_diagns = {d.strip() for d in args.test_diagns.split(",") if d.strip()} + unknown = test_diagns - set(counts) + if unknown: + raise ValueError(f"--test_diagns contains unknown diagns: {sorted(unknown)}") + else: + target = int(total * args.test_ratio) + test_diagns = set(pick_test_diagns_by_ratio(counts, target, args.seed)) + + train_samples = [s for s in samples if s.get("diagn") not in test_diagns] + test_samples = [s for s in samples if s.get("diagn") in test_diagns] + + write_jsonl(args.output_train, train_samples) + write_jsonl(args.output_test, test_samples) + + train_diagns = set(counts) - test_diagns + print("\nSplit summary:") + print( + f" train: {len(train_samples)} samples / {len(train_diagns)} diagns " + f"-> {args.output_train}" + ) + print( + f" test: {len(test_samples)} samples / {len(test_diagns)} diagns " + f"({100 * len(test_samples) / total:.2f}% of total) -> {args.output_test}" + ) + print(f"\nTest diagns ({len(test_diagns)}): {sorted(test_diagns)}") + + +if __name__ == "__main__": + main() diff --git a/examples/research_cod/get_alchemy_data.py b/examples/research_cod/get_alchemy_data.py new file mode 100644 index 00000000000..735597e465e --- /dev/null +++ b/examples/research_cod/get_alchemy_data.py @@ -0,0 +1,59 @@ +"""Generate Alchemy datasets with configurable size.""" + +import argparse +import os + +import numpy as np +import pandas as pd + +DEFAULT_DATA_PATH = os.path.join( + os.path.dirname(os.path.abspath(__file__)), "data", "alchemy" +) + + +def save_dataset_to_local(data_path: str, data: list, split: str = "default") -> str: + os.makedirs(data_path, exist_ok=True) + data_df = pd.DataFrame(data) + dataset_path = os.path.join(data_path, f"{split}.parquet") + data_df.to_parquet(dataset_path) + print(f"Saved split '{split}' with {len(data)} examples at {dataset_path}") + return dataset_path + + +def prepare_alchemy_data(data_path: str, train_size: int, test_size: int, seed: int): + np.random.seed(seed) + + all_seeds = np.random.choice(10_000_000, size=train_size + test_size, replace=False) + train_seeds = all_seeds[:train_size] + test_seeds = all_seeds[train_size:] + + def process_fn(seed_val, idx): + return {"seed": int(seed_val), "index": idx, "uid": f"alchemy_{seed_val}"} + + train_data = [process_fn(s, i) for i, s in enumerate(train_seeds)] + test_data = [process_fn(s, i) for i, s in enumerate(test_seeds)] + + save_dataset_to_local(data_path, train_data, "train") + save_dataset_to_local(data_path, test_data, "test") + + return train_data, test_data + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Generate Alchemy dataset") + parser.add_argument("--local_dir", default=DEFAULT_DATA_PATH) + parser.add_argument("--train_size", type=int, default=50000) + parser.add_argument("--test_size", type=int, default=100) + parser.add_argument("--seed", type=int, default=42) + args = parser.parse_args() + + train_data, test_data = prepare_alchemy_data( + data_path=args.local_dir, + train_size=args.train_size, + test_size=args.test_size, + seed=args.seed, + ) + + print(f"\nTrain: {len(train_data)} examples") + print(f"Test: {len(test_data)} examples") + print(f"Sample: {train_data[0]}") diff --git a/examples/research_cod/get_frozen_lake_data.py b/examples/research_cod/get_frozen_lake_data.py new file mode 100644 index 00000000000..c21f9f0cb04 --- /dev/null +++ b/examples/research_cod/get_frozen_lake_data.py @@ -0,0 +1,85 @@ +""" +Generate FrozenLake datasets with configurable map size range. +Modified from examples/grpo_frozen_lake/get_frozen_lake_data.py +""" +import argparse +import os + +import numpy as np +import pandas as pd + +DEFAULT_DATA_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "data", "frozen_lake") + + +def save_dataset_to_local(data_path: str, data: list[dict], split: str = "default") -> str: + os.makedirs(data_path, exist_ok=True) + data_df = pd.DataFrame(data) + dataset_path = os.path.join(data_path, f"{split}.parquet") + data_df.to_parquet(dataset_path) + print(f"Saved split '{split}' with {len(data)} examples at {dataset_path}") + return dataset_path + + +def prepare_frozenlake_data( + data_path, + train_size=10000, + test_size=100, + map_min_size=4, + map_max_size=8, + tile_min_prob=0.6, + tile_max_prob=0.85, + ): + np.random.seed(42) + + train_seeds = np.random.randint(0, 100000, size=train_size) + test_seeds = np.random.randint(0, 100000, size=test_size) + # randint is [low, high), so use max_size+1 to include max_size + train_sizes = np.random.randint(map_min_size, map_max_size + 1, size=train_size) + test_sizes = np.random.randint(map_min_size, map_max_size + 1, size=test_size) + # p is the probability of frozen tile, i.e., 1 - p is the probability of hole + train_ps = np.random.uniform(tile_min_prob, tile_max_prob, size=train_size) + test_ps = np.random.uniform(tile_min_prob, tile_max_prob, size=test_size) + + def frozenlake_process_fn(seed, size, p, idx): + return {"seed": seed, "size": size, "p": p, "index": idx, "uid": f"{seed}_{size}_{p}"} + + train_data = [ + frozenlake_process_fn(seed, train_sizes[idx], train_ps[idx], idx) + for idx, seed in enumerate(train_seeds) + ] + test_data = [ + frozenlake_process_fn(seed, test_sizes[idx], test_ps[idx], idx) + for idx, seed in enumerate(test_seeds) + ] + + save_dataset_to_local(data_path, train_data, "train") + save_dataset_to_local(data_path, test_data, "test") + + return train_data, test_data + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--local_dir", default=DEFAULT_DATA_PATH) + parser.add_argument("--train_size", type=int, default=50000) + parser.add_argument("--test_size", type=int, default=100) + parser.add_argument("--map_min_size", type=int, default=4) + parser.add_argument("--map_max_size", type=int, default=5) + parser.add_argument("--tile_min_prob", type=float, default=0.6) # tile prob: larger is easier + parser.add_argument("--tile_max_prob", type=float, default=0.7) + args = parser.parse_args() + + train_data, test_data = prepare_frozenlake_data( + data_path=args.local_dir, + train_size=args.train_size, + test_size=args.test_size, + map_min_size=args.map_min_size, + map_max_size=args.map_max_size, + tile_min_prob=args.tile_min_prob, + tile_max_prob=args.tile_max_prob, + ) + + print(f"\nTrain: {len(train_data)} examples") + print(f"Test: {len(test_data)} examples") + print(f"Size range: [{args.map_min_size}, {args.map_max_size}]") + print(f"Sample: {train_data[0]}") diff --git a/examples/research_cod/get_grid_navigation_data.py b/examples/research_cod/get_grid_navigation_data.py new file mode 100644 index 00000000000..08811890d37 --- /dev/null +++ b/examples/research_cod/get_grid_navigation_data.py @@ -0,0 +1,68 @@ +"""Generate seed-only datasets for the CoD grid-navigation environment.""" + +import argparse +import os + +import numpy as np +import pandas as pd + +DEFAULT_DATA_PATH = os.path.join( + os.path.dirname(os.path.abspath(__file__)), "data", "grid_navigation" +) + + +def save_dataset_to_local(data_path: str, data: list[dict], split: str) -> str: + """Write one dataset split as a parquet file.""" + os.makedirs(data_path, exist_ok=True) + dataset_path = os.path.join(data_path, f"{split}.parquet") + pd.DataFrame(data).to_parquet(dataset_path) + print(f"Saved split '{split}' with {len(data)} examples at {dataset_path}") + return dataset_path + + +def prepare_grid_navigation_data( + data_path: str, + train_size: int, + test_size: int, + seed: int, +): + """Generate disjoint train and test task seeds and save both splits.""" + rng = np.random.default_rng(seed) + all_seeds = rng.choice(10_000_000, size=train_size + test_size, replace=False) + train_seeds = all_seeds[:train_size] + test_seeds = all_seeds[train_size:] + + def process_fn(task_seed: int, index: int) -> dict: + task_seed = int(task_seed) + return { + "seed": task_seed, + "index": index, + "uid": f"grid_navigation_{task_seed}", + } + + train_data = [process_fn(task_seed, i) for i, task_seed in enumerate(train_seeds)] + test_data = [process_fn(task_seed, i) for i, task_seed in enumerate(test_seeds)] + + save_dataset_to_local(data_path, train_data, "train") + save_dataset_to_local(data_path, test_data, "test") + return train_data, test_data + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Generate CoD grid-navigation data") + parser.add_argument("--local_dir", default=DEFAULT_DATA_PATH) + parser.add_argument("--train_size", type=int, default=50000) + parser.add_argument("--test_size", type=int, default=4000) + parser.add_argument("--seed", type=int, default=42) + args = parser.parse_args() + + train_data, test_data = prepare_grid_navigation_data( + data_path=args.local_dir, + train_size=args.train_size, + test_size=args.test_size, + seed=args.seed, + ) + + print(f"\nTrain: {len(train_data)} examples") + print(f"Test: {len(test_data)} examples") + print(f"Sample: {train_data[0]}") diff --git a/examples/research_cod/get_optimal_control_data.py b/examples/research_cod/get_optimal_control_data.py new file mode 100644 index 00000000000..aa34e68fffb --- /dev/null +++ b/examples/research_cod/get_optimal_control_data.py @@ -0,0 +1,231 @@ +"""Generate train/test datasets for the CoD Optimal Control workflow. + +The script supports reproducible random task generation with configurable +difficulty. A YAML config (e.g. the benchmark YAML) can be passed via +``--config``; any CLI arguments override the config values. +""" + +import argparse +import os +from typing import Any, Dict, List, Optional, Tuple + +import numpy as np +import pandas as pd +from omegaconf import OmegaConf + +from trinity.common.constants import TASKSET_PATH_ENV_VAR + +DEFAULT_DATA_PATH = os.path.join( + os.path.dirname(os.path.abspath(__file__)), "..", "data", "optimal_control" +) + +# Difficulty presets control the horizon and sampled state space. The hard +# preset keeps state scales bounded while moving targets farther away. +DIFFICULTY_PRESETS: Dict[str, Dict[str, Any]] = { + "easy": { + "max_horizon": 8, + "x0_min": -1.0, + "x0_max": 1.0, + "v0_min": -0.5, + "v0_max": 0.5, + "x_target_min": -3.0, + "x_target_max": 3.0, + }, + "medium": { + "max_horizon": 12, + "x0_min": -2.0, + "x0_max": 2.0, + "v0_min": -1.0, + "v0_max": 1.0, + "x_target_min": -5.0, + "x_target_max": 5.0, + }, + "hard": { + "max_horizon": 12, + "x0_min": -2.0, + "x0_max": 2.0, + "v0_min": -1.5, + "v0_max": 1.5, + "x_target_min": -8.0, + "x_target_max": 8.0, + }, +} + + +def save_dataset_to_local(data_path: str, data: list[dict], split: str = "default") -> str: + """Save dataset directly to local data_path.""" + os.makedirs(data_path, exist_ok=True) + data_df = pd.DataFrame(data) + dataset_path = os.path.join(data_path, f"{split}.parquet") + data_df.to_parquet(dataset_path) + print( + f"Saved dataset optimal_control split '{split}' with {len(data)} examples at " + f"{dataset_path}. Make sure to set the environment variable {TASKSET_PATH_ENV_VAR} " + f"to {data_path}." + ) + return dataset_path + + +def _sample_float(rng: np.random.Generator, low: float, high: float) -> float: + """Sample a single float in [low, high).""" + return float(rng.uniform(low, high)) + + +def _difficulty_params( + difficulty: str, + overrides: Optional[Dict[str, Any]] = None, +) -> Dict[str, Any]: + """Resolve difficulty preset and apply per-field overrides.""" + if difficulty not in DIFFICULTY_PRESETS: + raise ValueError( + f"Unknown difficulty '{difficulty}'. Choose from {list(DIFFICULTY_PRESETS.keys())}." + ) + params = dict(DIFFICULTY_PRESETS[difficulty]) + # The control penalty is independent of difficulty by default. + params.setdefault("control_penalty_coef", 0.03) + params.setdefault("v_target", 0.0) + # The workflow uses max_horizon directly as the rollout horizon. + params.setdefault("max_horizon", 30) + if overrides: + for key, value in overrides.items(): + if value is not None and key in params: + params[key] = value + return params + + +def _generate_random_tasks( + rng: np.random.Generator, + size: int, + params: Dict[str, Any], + seed_offset: int, +) -> List[dict]: + """Generate independent tasks whose hidden dynamics are resolved at runtime.""" + return [ + { + "x0": _sample_float(rng, params["x0_min"], params["x0_max"]), + "v0": _sample_float(rng, params["v0_min"], params["v0_max"]), + "x_target": _sample_float( + rng, + params["x_target_min"], + params["x_target_max"], + ), + "v_target": float(params["v_target"]), + "max_horizon": int(params["max_horizon"]), + "control_penalty_coef": float(params["control_penalty_coef"]), + "seed": seed_offset + task_idx, + } + for task_idx in range(size) + ] + + +def prepare_optimal_control_data( + data_path: str, + train_size: int = 128, + test_size: int = 16, + difficulty: str = "medium", + train_seed: int = 42, + test_seed: int = 2024, + **range_overrides: Any, +) -> Tuple[List[dict], List[dict]]: + """Generate pack-independent train and test splits.""" + params = _difficulty_params(difficulty, range_overrides) + train_data = _generate_random_tasks( + np.random.default_rng(train_seed), + train_size, + params, + seed_offset=0, + ) + test_data = _generate_random_tasks( + np.random.default_rng(test_seed), + test_size, + params, + seed_offset=train_size, + ) + save_dataset_to_local(data_path, train_data, "train") + save_dataset_to_local(data_path, test_data, "test") + return train_data, test_data + + +def _load_config(config_path: Optional[str]) -> Dict[str, Any]: + """Load an optional YAML config and extract ``task_generation`` settings.""" + if not config_path: + return {} + cfg = OmegaConf.load(config_path) + task_gen = cfg.get("task_generation", {}) + if not task_gen: + return {} + return OmegaConf.to_container(task_gen, resolve=True) + + +def _build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser() + parser.add_argument("--local_dir", default=DEFAULT_DATA_PATH) + parser.add_argument( + "--config", default=None, help="Path to a YAML config with a 'task_generation' section." + ) + parser.add_argument("--train_size", type=int, default=None) + parser.add_argument("--test_size", type=int, default=None) + parser.add_argument("--difficulty", type=str, default=None) + parser.add_argument("--train_seed", type=int, default=None) + parser.add_argument("--test_seed", type=int, default=None) + parser.add_argument("--max_horizon", type=int, default=None) + parser.add_argument("--control_penalty_coef", type=float, default=None) + parser.add_argument("--x0_min", type=float, default=None) + parser.add_argument("--x0_max", type=float, default=None) + parser.add_argument("--v0_min", type=float, default=None) + parser.add_argument("--v0_max", type=float, default=None) + parser.add_argument("--x_target_min", type=float, default=None) + parser.add_argument("--x_target_max", type=float, default=None) + return parser + + +def main() -> None: + parser = _build_parser() + args = parser.parse_args() + + # Config values are the base; CLI args override them. + config = _load_config(args.config) + + def get(name: str, default: Any) -> Any: + cli_value = getattr(args, name, None) + if cli_value is not None: + return cli_value + return config.get(name, default) + + train_size = int(get("train_size", 128)) + test_size = int(get("test_size", 16)) + difficulty = str(get("difficulty", "medium")) + train_seed = int(get("train_seed", 42)) + test_seed = int(get("test_seed", 2024)) + + range_overrides = { + "max_horizon": get("max_horizon", None), + "control_penalty_coef": get("control_penalty_coef", None), + "x0_min": get("x0_min", None), + "x0_max": get("x0_max", None), + "v0_min": get("v0_min", None), + "v0_max": get("v0_max", None), + "x_target_min": get("x_target_min", None), + "x_target_max": get("x_target_max", None), + } + # Filter out unset overrides so the difficulty preset is not overwritten. + range_overrides = {k: v for k, v in range_overrides.items() if v is not None} + + train_data, test_data = prepare_optimal_control_data( + data_path=args.local_dir, + train_size=train_size, + test_size=test_size, + difficulty=difficulty, + train_seed=train_seed, + test_seed=test_seed, + **range_overrides, + ) + + print(f"Train dataset: {len(train_data)} examples") + print(f"Test dataset: {len(test_data)} examples") + print("Sample train example:", train_data[0]) + print("Sample test example:", test_data[0]) + + +if __name__ == "__main__": + main() diff --git a/examples/research_cod/get_pde_discovery_data.py b/examples/research_cod/get_pde_discovery_data.py new file mode 100644 index 00000000000..33f5bfcaae8 --- /dev/null +++ b/examples/research_cod/get_pde_discovery_data.py @@ -0,0 +1,513 @@ +"""Generate CoD PDE discovery tasksets. + +The dataset intentionally does not pre-bind rows to CoD packs. Pack-level +randomness is injected at runtime by CoDWorkflow via ``pack_seed`` so changing +pack size or mixing tasksets does not require regenerating the dataset. +""" + +import argparse +import importlib.util +import os +import sys +import types +from collections import Counter +from pathlib import Path +from typing import Iterable + +import numpy as np +import pandas as pd + +DEFAULT_DATA_PATH = os.path.join( + os.path.dirname(os.path.abspath(__file__)), "data", "pde_discovery" +) + + +def save_dataset_to_local(data_path: str, data: list[dict], split: str) -> str: + os.makedirs(data_path, exist_ok=True) + dataset_path = os.path.join(data_path, f"{split}.parquet") + pd.DataFrame(data).to_parquet(dataset_path) + print(f"Saved split '{split}' with {len(data)} examples at {dataset_path}") + return dataset_path + + +def prepare_pde_discovery_data( + data_path: str, + train_size: int, + test_size: int, + seed: int, + eval_pack_size: int | None = None, + eval_template_count: int | None = None, + exclude_eval_paths: list[str] | None = None, + eval_ground_truth_family: str = "physical_full_eval_hard.json", + eval_min_reaction_terms: int = 1, + eval_max_reaction_terms: int = 3, + eval_state_abs_limit: float = 3.0, + eval_initial_mode_count_range: tuple[int, int] = (1, 2), + save_train_split: bool = True, + online_eval_data_path: str | None = None, + online_eval_instances_per_template: int | None = None, +) -> tuple[list[dict], list[dict]]: + counts = { + "train_size": train_size, + "test_size": test_size, + } + invalid = [name for name, value in counts.items() if value <= 0] + if invalid: + raise ValueError(f"PDE dataset counts must be positive: {', '.join(invalid)}") + if (online_eval_data_path is None) != (online_eval_instances_per_template is None): + raise ValueError( + "online_eval_data_path and online_eval_instances_per_template " + "must be provided together" + ) + if online_eval_data_path is not None: + if os.path.abspath(online_eval_data_path) == os.path.abspath(data_path): + raise ValueError( + "Online eval directory must differ from the full dataset directory" + ) + stale_online_train_path = os.path.join(online_eval_data_path, "train.parquet") + if os.path.exists(stale_online_train_path): + raise ValueError( + "Online test-only directory contains train.parquet: " + f"{stale_online_train_path}" + ) + + rng = np.random.default_rng(seed) + uint32_space_size = int(np.iinfo(np.uint32).max) + 1 + all_task_seeds = rng.choice( + uint32_space_size, + size=train_size + test_size, + replace=False, + ) + + eval_template_assignments = build_stratified_eval_template_assignments( + test_size=test_size, + eval_pack_size=eval_pack_size, + eval_template_count=eval_template_count, + ) + + def build_split(task_seeds, split_name, template_assignments=None): + rows = [] + for task_idx, task_seed in enumerate(task_seeds): + row = { + "uid": f"pde_{split_name}_{task_idx}", + "seed": int(task_seed), + "task_desc": ( + "Discover the nonlinear reaction term f(u) in " + "partial_t u = partial_xx u + f(u) using active " + "sampling, sparse regression, and scientific context." + ), + "answer": "", + } + if template_assignments is not None: + assignment = template_assignments[task_idx] + row.update(assignment) + pack_start = (task_idx // eval_pack_size) * eval_pack_size + row["pde_environment_seed"] = int(task_seeds[pack_start]) + rows.append(row) + return rows + + train_data = build_split( + all_task_seeds[:train_size], + "train", + ) + test_data = build_split( + all_task_seeds[train_size:], + "test", + template_assignments=eval_template_assignments, + ) + if exclude_eval_paths: + if eval_pack_size is None or eval_template_count is None: + raise ValueError( + "Equation-disjoint evaluation requires eval_pack_size and " + "eval_template_count" + ) + validate_eval_equation_disjointness( + new_rows=test_data, + exclude_eval_paths=exclude_eval_paths, + eval_pack_size=eval_pack_size, + ground_truth_family=eval_ground_truth_family, + min_reaction_terms=eval_min_reaction_terms, + max_reaction_terms=eval_max_reaction_terms, + state_abs_limit=eval_state_abs_limit, + initial_mode_count_range=eval_initial_mode_count_range, + ) + + if save_train_split: + save_dataset_to_local(data_path, train_data, "train") + else: + stale_train_path = os.path.join(data_path, "train.parquet") + if os.path.exists(stale_train_path): + raise ValueError( + f"--test_only requires a directory without train.parquet: {stale_train_path}" + ) + save_dataset_to_local(data_path, test_data, "test") + if online_eval_data_path is not None: + assert online_eval_instances_per_template is not None + if eval_pack_size is None or eval_template_count is None: + raise ValueError( + "Online eval export requires eval_pack_size and eval_template_count" + ) + online_test_data = select_eval_template_instances( + test_data, + eval_pack_size=eval_pack_size, + eval_template_count=eval_template_count, + instances_per_template=online_eval_instances_per_template, + ) + save_dataset_to_local(online_eval_data_path, online_test_data, "test") + return train_data, test_data + + +def _resolve_test_parquet(path: str) -> Path: + candidate = Path(path) + if candidate.is_dir(): + candidate = candidate / "test.parquet" + if not candidate.is_file(): + raise FileNotFoundError(f"Excluded eval test parquet does not exist: {candidate}") + return candidate + + +def _eval_pack_descriptors(rows: Iterable[dict], eval_pack_size: int) -> list[dict]: + """Return one hidden-environment descriptor for every fixed eval pack.""" + grouped: dict[int, list[dict]] = {} + for row in rows: + if "eval_pack_index" not in row: + raise ValueError("Stratified eval row is missing eval_pack_index") + grouped.setdefault(int(row["eval_pack_index"]), []).append(row) + + descriptors = [] + for pack_index in sorted(grouped): + pack = grouped[pack_index] + if len(pack) != eval_pack_size: + raise ValueError( + f"Eval pack {pack_index} has {len(pack)} rows, expected {eval_pack_size}" + ) + template_indices = {int(row["ground_truth_template_index"]) for row in pack} + environment_seeds = {int(row["pde_environment_seed"]) for row in pack} + if len(template_indices) != 1 or len(environment_seeds) != 1: + raise ValueError( + f"Eval pack {pack_index} does not share one template and environment seed" + ) + descriptors.append( + { + "pack_index": pack_index, + "template_index": template_indices.pop(), + "environment_seed": environment_seeds.pop(), + } + ) + return descriptors + + +def select_eval_template_instances( + rows: list[dict], + eval_pack_size: int, + eval_template_count: int, + instances_per_template: int, +) -> list[dict]: + """Select the first fixed equation instances of every eval template.""" + if instances_per_template <= 0: + raise ValueError("Online eval instances per template must be positive") + selected = [ + row + for row in rows + if int(row["ground_truth_template_instance"]) < instances_per_template + ] + descriptors = _eval_pack_descriptors(selected, eval_pack_size) + template_counts = Counter(descriptor["template_index"] for descriptor in descriptors) + expected_pack_count = eval_template_count * instances_per_template + if len(descriptors) != expected_pack_count: + raise ValueError( + f"Online eval selected {len(descriptors)} packs, expected {expected_pack_count}" + ) + if set(template_counts) != set(range(eval_template_count)) or set( + template_counts.values() + ) != {instances_per_template}: + raise ValueError("Online eval does not cover every template equally") + return selected + + +def _load_pde_ground_truth_modules(): + """Load equation-sampling modules without importing the vLLM workflow.""" + root = ( + Path(__file__).parents[2] + / "trinity/common/workflows/connect_the_dots/pde_discovery" + ) + package_name = "_pde_eval_generation_modules" + package = sys.modules.setdefault(package_name, types.ModuleType(package_name)) + package.__path__ = [str(root)] + + loaded = {} + for name in ("candidate", "pde_numeric", "ground_truth"): + module_name = f"{package_name}.{name}" + if module_name in sys.modules: + loaded[name] = sys.modules[module_name] + continue + spec = importlib.util.spec_from_file_location(module_name, root / f"{name}.py") + if spec is None or spec.loader is None: + raise ImportError(f"Cannot load PDE equation module: {name}") + module = importlib.util.module_from_spec(spec) + sys.modules[module_name] = module + spec.loader.exec_module(module) + loaded[name] = module + return loaded["candidate"], loaded["ground_truth"], loaded["pde_numeric"] + + +def _hidden_equation_signatures( + rows: Iterable[dict], + eval_pack_size: int, + ground_truth_family: str, + min_reaction_terms: int, + max_reaction_terms: int, + state_abs_limit: float, + initial_mode_count_range: tuple[int, int], +) -> list[tuple[tuple[tuple[str, float], ...], int]]: + """Reproduce the workflow's stable hidden equation for each eval pack. + + The equation signature is support plus rounded coefficients. Environment + seeds and initial conditions are deliberately not part of the signature, so + two packs with the same mathematical reaction are treated as duplicates. + """ + candidate, ground_truth, pde_numeric = _load_pde_ground_truth_modules() + + templates = ground_truth.load_hidden_gt_templates( + candidate.DEFAULT_DICTIONARY, + family=ground_truth_family, + ) + initial_shape = pde_numeric.InitialConditionShapeConfig.from_mapping( + {"mode_count_range": list(initial_mode_count_range)} + ) + initial_amplitude_upper = ground_truth.calibrate_family_initial_amplitude_upper( + templates, + min_reaction_terms, + max_reaction_terms, + state_abs_limit=state_abs_limit, + ) + + signatures = [] + for descriptor in _eval_pack_descriptors(rows, eval_pack_size): + environment_seed = descriptor["environment_seed"] + rng = np.random.default_rng( + np.random.SeedSequence( + [environment_seed, pde_numeric.GROUND_TRUTH_STREAM] + ) + ) + coefficients = {} + for _ in range(ground_truth.GT_STABILITY_MAX_ATTEMPTS): + _, coefficients = ground_truth.sample_hidden_reaction( + rng, + templates, + min_reaction_terms, + max_reaction_terms, + template_index=descriptor["template_index"], + ) + if ground_truth.is_stable_reaction_candidate( + coefficients=coefficients, + pde_grid_size=257, + pde_time_steps=2001, + trajectory_count=eval_pack_size, + pack_seed=environment_seed, + initial_amplitude_upper=initial_amplitude_upper, + initial_condition_shape=initial_shape, + state_abs_limit=state_abs_limit, + ): + break + signature = tuple( + (term, float(coefficients[term])) + for term in candidate.DEFAULT_DICTIONARY + if term in coefficients + ) + signatures.append((signature, descriptor["pack_index"])) + return signatures + + +def _format_equation_signature(signature: tuple[tuple[str, float], ...]) -> str: + return " + ".join(f"{coefficient:+.2f}*{term}" for term, coefficient in signature) + + +def validate_eval_equation_disjointness( + new_rows: list[dict], + exclude_eval_paths: list[str], + eval_pack_size: int, + ground_truth_family: str, + min_reaction_terms: int, + max_reaction_terms: int, + state_abs_limit: float, + initial_mode_count_range: tuple[int, int], +) -> None: + """Fail generation when exact hidden equations repeat internally or historically.""" + signature_args = { + "eval_pack_size": eval_pack_size, + "ground_truth_family": ground_truth_family, + "min_reaction_terms": min_reaction_terms, + "max_reaction_terms": max_reaction_terms, + "state_abs_limit": state_abs_limit, + "initial_mode_count_range": initial_mode_count_range, + } + new_records = _hidden_equation_signatures(new_rows, **signature_args) + new_counts = Counter(signature for signature, _ in new_records) + internal_duplicates = [signature for signature, count in new_counts.items() if count > 1] + if internal_duplicates: + rendered = "; ".join( + _format_equation_signature(signature) for signature in internal_duplicates[:5] + ) + raise ValueError(f"New eval set contains duplicate hidden equations: {rendered}") + + excluded_signatures = set() + excluded_pack_count = 0 + for excluded_path in exclude_eval_paths: + parquet_path = _resolve_test_parquet(excluded_path) + old_rows = pd.read_parquet(parquet_path).to_dict("records") + old_records = _hidden_equation_signatures(old_rows, **signature_args) + excluded_signatures.update(signature for signature, _ in old_records) + excluded_pack_count += len(old_records) + + overlaps = sorted(set(new_counts) & excluded_signatures, key=repr) + if overlaps: + rendered = "; ".join( + _format_equation_signature(signature) for signature in overlaps[:5] + ) + raise ValueError( + f"New eval set overlaps {len(overlaps)} excluded hidden equations: {rendered}" + ) + print( + "Verified hidden-equation disjointness: " + f"{len(new_records)} new packs are unique and disjoint from " + f"{excluded_pack_count} excluded packs" + ) + + +def build_stratified_eval_template_assignments( + test_size: int, + eval_pack_size: int | None, + eval_template_count: int | None, +) -> list[dict] | None: + """Assign an equal number of fixed eval packs to every template. + + Training rows remain pack-agnostic. This optional metadata is intended only + for a fixed scientific benchmark whose row order and pack size are held + constant across checkpoints. The metadata is consumed inside the PDE + environment and is never rendered into the model prompt. + """ + if eval_pack_size is None and eval_template_count is None: + return None + if eval_pack_size is None or eval_template_count is None: + raise ValueError( + "eval_pack_size and eval_template_count must be provided together" + ) + if eval_pack_size <= 0 or eval_template_count <= 0: + raise ValueError("eval pack size and template count must be positive") + if test_size % eval_pack_size != 0: + raise ValueError( + f"test_size={test_size} must be divisible by eval_pack_size={eval_pack_size}" + ) + + pack_count = test_size // eval_pack_size + if pack_count % eval_template_count != 0: + raise ValueError( + f"eval pack count {pack_count} must be divisible by " + f"eval_template_count={eval_template_count} for balanced coverage" + ) + + assignments = [] + for pack_index in range(pack_count): + template_index = pack_index % eval_template_count + template_instance_index = pack_index // eval_template_count + assignment = { + "eval_pack_index": pack_index, + "ground_truth_template_index": template_index, + "ground_truth_template_instance": template_instance_index, + } + assignments.extend(dict(assignment) for _ in range(eval_pack_size)) + return assignments + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Generate PDE discovery CoD dataset") + parser.add_argument("--local_dir", default=DEFAULT_DATA_PATH) + parser.add_argument("--train_size", type=int, default=50000) + parser.add_argument("--test_size", type=int, default=32) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument( + "--eval_pack_size", + type=int, + default=None, + help="Fixed CoD pack size for an optional stratified test split.", + ) + parser.add_argument( + "--eval_template_count", + type=int, + default=None, + help=( + "Number of hidden templates to cover equally in the stratified test " + "split. Must be used with --eval_pack_size." + ), + ) + parser.add_argument( + "--exclude_eval_path", + action="append", + default=[], + help=( + "Existing stratified eval directory or test.parquet whose exact hidden " + "equations must not occur in the new test split. May be repeated." + ), + ) + parser.add_argument( + "--online_eval_dir", + default=None, + help=( + "Optional directory for a smaller, deterministic online-eval subset " + "selected from the full stratified test split." + ), + ) + parser.add_argument( + "--online_eval_instances_per_template", + type=int, + default=None, + help=( + "Number of the first equation instances to retain per hidden template " + "in --online_eval_dir. Must be used with --online_eval_dir." + ), + ) + parser.add_argument( + "--eval_ground_truth_family", + default="physical_full_eval_hard.json", + help="Ground-truth family used to reproduce equations for overlap checks.", + ) + parser.add_argument("--eval_min_reaction_terms", type=int, default=1) + parser.add_argument("--eval_max_reaction_terms", type=int, default=3) + parser.add_argument("--eval_state_abs_limit", type=float, default=3.0) + parser.add_argument( + "--eval_initial_mode_count_range", + type=int, + nargs=2, + metavar=("MIN", "MAX"), + default=(1, 2), + ) + parser.add_argument( + "--test_only", + action="store_true", + help="Write only test.parquet, avoiding mixed-schema dataset directories.", + ) + args = parser.parse_args() + + train_data, test_data = prepare_pde_discovery_data( + data_path=args.local_dir, + train_size=args.train_size, + test_size=args.test_size, + seed=args.seed, + eval_pack_size=args.eval_pack_size, + eval_template_count=args.eval_template_count, + exclude_eval_paths=args.exclude_eval_path, + eval_ground_truth_family=args.eval_ground_truth_family, + eval_min_reaction_terms=args.eval_min_reaction_terms, + eval_max_reaction_terms=args.eval_max_reaction_terms, + eval_state_abs_limit=args.eval_state_abs_limit, + eval_initial_mode_count_range=tuple(args.eval_initial_mode_count_range), + save_train_split=not args.test_only, + online_eval_data_path=args.online_eval_dir, + online_eval_instances_per_template=args.online_eval_instances_per_template, + ) + + print(f"\nTrain rows: {len(train_data)}") + print(f"Test rows: {len(test_data)}") + if train_data: + print(f"Sample: {train_data[0]}") diff --git a/examples/research_cod/get_terminal_data.py b/examples/research_cod/get_terminal_data.py new file mode 100644 index 00000000000..cef70f3ad9e --- /dev/null +++ b/examples/research_cod/get_terminal_data.py @@ -0,0 +1,118 @@ +# -*- coding: utf-8 -*- +"""Generate training and evaluation datasets for the terminal file-ops task. + +Usage:: + + python examples/research_cod/get_terminal_data.py \ + --local_dir examples/research_cod/data/terminal \ + --train_size 50000 \ + --test_size 200 +""" + +import argparse +import os +import random + +import numpy as np +import pandas as pd + + +DEFAULT_DATA_PATH = "examples/research_cod/data/terminal" + +# Task types that can appear in the dataset +SINGLE_TYPES = [ + "upload", "download", "rename", "move", "chmod", + "delete", "copy", "pack", "mkdir", +] +COMPOSITE_TYPES = [ + "pack_upload", "download_extract", "mkdir_upload", + "upload_chmod", "upload_delete_source", "pack_upload_extract", + "download_rename", "backup_replace", +] +ALL_TYPES = SINGLE_TYPES + COMPOSITE_TYPES + + +def save_dataset_to_local(data_path: str, data: list, split: str) -> str: + os.makedirs(data_path, exist_ok=True) + df = pd.DataFrame(data) + path = os.path.join(data_path, f"{split}.parquet") + df.to_parquet(path) + print(f"Saved {len(data)} samples to {path}") + return path + + +def prepare_data( + data_path: str, + train_size: int = 50000, + test_size: int = 200, + seed: int = 42, + composite_ratio: float = 0.5, +): + np.random.seed(seed) + + # Generate unique seeds + all_seeds = set() + while len(all_seeds) < train_size + test_size: + all_seeds.add(np.random.randint(0, 10_000_000)) + all_seeds = list(all_seeds) + np.random.shuffle(all_seeds) + + train_seeds = all_seeds[:train_size] + test_seeds = all_seeds[train_size:train_size + test_size] + + def make_record(task_seed, idx): + # The task type is determined by the seed at runtime; + # but we precompute it here for metadata.source (pack grouping). + task_rng = random.Random(task_seed) + task_rng.choice(["windows", "mac", "linux"]) # skip OS pick + task_rng.randint(10, 99) # skip remote host + task_rng.choice(["admin", "deploy", "user", "webmaster", "devops", "ops", "ubuntu"]) # skip user + is_composite = task_rng.random() < composite_ratio + if is_composite: + task_type = task_rng.choice(COMPOSITE_TYPES) + else: + task_type = task_rng.choice(SINGLE_TYPES) + + return { + "seed": int(task_seed), + "prompt": str(task_seed), + "index": idx, + "uid": f"terminal_{task_seed}", + "metadata": { + "source": task_type, + "source_dataset": "terminal", + }, + } + + train_data = [make_record(s, i) for i, s in enumerate(train_seeds)] + test_data = [make_record(s, i) for i, s in enumerate(test_seeds)] + + save_dataset_to_local(data_path, train_data, "train") + save_dataset_to_local(data_path, test_data, "test") + + # Print distribution + from collections import Counter + train_types = Counter(r["metadata"]["source"] for r in train_data) + print(f"\nTask type distribution (train):") + for t, c in sorted(train_types.items(), key=lambda x: -x[1]): + print(f" {t}: {c} ({100 * c / train_size:.1f}%)") + + return train_data, test_data + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Generate terminal task dataset") + parser.add_argument("--local_dir", default=DEFAULT_DATA_PATH) + parser.add_argument("--train_size", type=int, default=50000) + parser.add_argument("--test_size", type=int, default=200) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--composite_ratio", type=float, default=0.5) + args = parser.parse_args() + + prepare_data( + data_path=args.local_dir, + train_size=args.train_size, + test_size=args.test_size, + seed=args.seed, + composite_ratio=args.composite_ratio, + ) diff --git a/pyproject.toml b/pyproject.toml index 008c00ab23f..a2469eb7264 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -50,6 +50,7 @@ dependencies = [ "datasets>=4.0.0", "typer>=0.23.0", "fsspec>=2023.10.0", + "gymnasium", ] [project.scripts] @@ -58,7 +59,7 @@ trinity = "trinity.cli.launcher:main" [project.optional-dependencies] vllm = [ # for routed_experts support, please install vllm 0.22.0 or above - "vllm>=0.19.1,<=0.23.0", + "vllm>=0.22.0,<=0.23.0", ] sglang = [ "sglang==0.5.13", @@ -127,7 +128,7 @@ flash_attn = [ qwen3_5 = [ "flash-linear-attention>=0.4.2", - "causal_conv1d>=1.6.0", + "causal_conv1d>=1.6.2.post1", ] [tool.setuptools.packages.find] diff --git a/tests/common/pde_discovery_test.py b/tests/common/pde_discovery_test.py new file mode 100644 index 00000000000..8fb0ccc9ca2 --- /dev/null +++ b/tests/common/pde_discovery_test.py @@ -0,0 +1,574 @@ +import importlib.util +import sys +import types +from collections import Counter +from pathlib import Path + +import numpy as np + +PDE_ROOT = ( + Path(__file__).parents[2] + / "trinity/common/workflows/connect_the_dots/pde_discovery" +) +PACKAGE = "pde_discovery_test_modules" + + +def _load_module(name: str, path: Path | None = None): + package = sys.modules.setdefault(PACKAGE, types.ModuleType(PACKAGE)) + package.__path__ = [str(PDE_ROOT)] + module_name = f"{PACKAGE}.{name}" + spec = importlib.util.spec_from_file_location( + module_name, path or PDE_ROOT / f"{name}.py" + ) + module = importlib.util.module_from_spec(spec) + sys.modules[module_name] = module + spec.loader.exec_module(module) + return module + + +candidate = _load_module("candidate") +pde_numeric = _load_module("pde_numeric") +ground_truth = _load_module("ground_truth") +regression = _load_module("regression") +cod_utils = _load_module("cod_utils", PDE_ROOT.parent / "utils.py") +prompts = _load_module("prompts", PDE_ROOT / "prompts/__init__.py") +data_generator = _load_module( + "data_generator", + Path(__file__).parents[2] / "examples/research_cod/get_pde_discovery_data.py", +) + +INITIAL_CONDITION_TEST_PACK_SEED = 1234 +INITIAL_CONDITION_TEST_GRID_SIZE = 129 +INITIAL_CONDITION_TEST_TRAJECTORY_COUNTS = (1, 3, 5, 8) +INITIAL_CONDITION_TEST_STATE_ABS_LIMIT = 3.0 +INITIAL_CONDITION_TEST_SHAPE = pde_numeric.InitialConditionShapeConfig( + mode_count_range=(1, 4), +) + + +def _parse_action(response: str): + payload, parse_error = cod_utils.parse_xml_answer( + response, {"dictionary", "candidate_equations", "uncertain_terms"} + ) + return types.SimpleNamespace(payload=payload, parse_error=parse_error or None) + + +def test_xml_answer_protocol_parses_reasoning_and_sampling(): + response = """I should sample broad regions first. + + + + + + + + + Additional reasoning outside the answer is ignored. + """ + + result = _parse_action(response) + + assert result.parse_error is None + assert result.payload == { + "action": "sample_pde_data", + "args": { + "point_grid": { + "point": [ + {"x": "0.2", "t": "0.1"}, + {"x": "0.8", "t": "0.9"}, + ], + } + }, + } + + +def test_xml_answer_protocol_parses_all_pde_tools(): + cases = [ + ( + "", + {"action": "summarize_pack_evidence", "args": {}}, + ), + ( + """ + merged_all0.05 + uu**3 + """, + { + "action": "run_sparse_regression", + "args": { + "dataset_id": "merged_all", + "alpha": "0.05", + "dictionary": ["u", "u**3"], + }, + }, + ), + ( + """ + 1.2*u - 0.8*u**3 + Cubic support is preferred. + sin(u) + """, + { + "action": "update_scientific_context", + "args": { + "preferred_equation": "1.2*u - 0.8*u**3", + "note": "Cubic support is preferred.", + "uncertain_terms": ["sin(u)"], + }, + }, + ), + ] + + for response, expected in cases: + result = _parse_action(response) + assert result.parse_error is None + assert result.payload == expected + + +def test_xml_answer_protocol_parses_candidate_list(): + equations = "".join(f"{idx}*u" for idx in range(8)) + response = ( + "" + f"{equations}" + ) + + result = _parse_action(response) + + assert result.parse_error is None + assert len(result.payload["args"]["candidate_equations"]) == 8 + + +def test_xml_answer_protocol_rejects_ambiguous_or_invalid_actions(): + cases = [ + ( + "" + "", + "expected_exactly_one_answer_tag", + ), + ( + "", + "invalid_answer_xml", + ), + ( + '{"action":"summarize_pack_evidence","args":{}}', + "invalid_answer_xml", + ), + ] + + for response, expected_error in cases: + result = _parse_action(response) + assert result.payload is None + assert result.parse_error == expected_error + + +def test_system_prompt_requires_one_terminal_xml_action(): + system_prompt = prompts.load_system_prompt( + dictionary_terms=["u", "u**3"], + support_size_hint="1 to 3", + ) + + assert system_prompt.count("## Response protocol") == 1 + assert "one ... block wrapping exactly one action" in system_prompt + assert "Never write" in system_prompt + assert "after the closing " in system_prompt + + +def test_candidate_coefficients_match_rendered_equation(): + coefficients = { + "u": -1.25, + "exp(u)-1": 0.5, + "log(1+u**2)": -0.25, + } + equation = candidate.format_equation(coefficients) + + assert candidate.candidate_coefficients(equation, candidate.DEFAULT_DICTIONARY) == coefficients + assert candidate.candidate_coefficients("f(u)=1*u+2*u", ["u"]) == {"u": 3.0} + assert candidate.candidate_coefficients( + "1*u+999*unknown", candidate.DEFAULT_DICTIONARY + ) == {} + + +def test_expanded_dictionary_terms_are_finite_and_round_trip(): + terms = [ + "u**7", + "sin(2*u)", + "u**4/(1+u**4)", + "u/(1+u+u**2)", + ] + u = np.linspace(-2.0, 2.0, 101) + coefficients = {term: float(index + 1) for index, term in enumerate(terms)} + + for term in terms: + assert np.all(np.isfinite(candidate.basis_values(term, u))) + equation = candidate.format_equation(coefficients) + assert candidate.candidate_coefficients(equation, terms) == coefficients + + +def test_physical_hard_ground_truth_family_uses_supported_two_term_templates(): + templates = ground_truth.load_hidden_gt_templates( + candidate.DEFAULT_DICTIONARY, + family="physical_hard", + ) + + assert templates + assert all(len(template["terms"]) == 2 for template in templates) + assert all( + term in candidate.DEFAULT_DICTIONARY + for template in templates + for term in template["terms"] + ) + + +def test_physical_full_ground_truth_family_is_deduplicated_one_to_two_terms(): + def template_key(template): + return ( + tuple(sorted(template["terms"])), + tuple(sorted(template["signs"].items())), + ) + + full = ground_truth.load_hidden_gt_templates( + candidate.DEFAULT_DICTIONARY, + family="physical_full", + ) + full_keys = [template_key(template) for template in full] + + assert len(full_keys) == len(set(full_keys)) + assert len(full_keys) == 55 + assert {len(template["terms"]) for template in full} == {1, 2} + assert any(set(template["terms"]) == {"u**3", "u**7"} for template in full) + + +def test_physical_full_train_eval_splits_are_support_disjoint(): + full = ground_truth.load_hidden_gt_templates( + candidate.DEFAULT_DICTIONARY, + family="physical_full", + ) + train = ground_truth.load_hidden_gt_templates( + candidate.DEFAULT_DICTIONARY, + family="physical_full_train", + ) + evaluation = ground_truth.load_hidden_gt_templates( + candidate.DEFAULT_DICTIONARY, + family="physical_full_eval", + ) + + def template_key(template): + return ( + tuple(sorted(template["terms"])), + tuple(sorted(template["signs"].items())), + ) + + train_supports = {frozenset(template["terms"]) for template in train} + eval_supports = {frozenset(template["terms"]) for template in evaluation} + train_terms = {term for template in train for term in template["terms"]} + eval_terms = {term for template in evaluation for term in template["terms"]} + + assert len(train) == 40 + assert len(evaluation) == 15 + assert train_supports.isdisjoint(eval_supports) + assert train_terms == set(candidate.DEFAULT_DICTIONARY) + assert eval_terms <= train_terms + assert {template_key(template) for template in full} == { + template_key(template) for template in [*train, *evaluation] + } + assert {len(template["terms"]) for template in train} == {1, 2} + assert {len(template["terms"]) for template in evaluation} == {2} + + +def test_physical_full_eval_hard_has_unique_two_term_and_varied_three_term_supports(): + train = ground_truth.load_hidden_gt_templates( + candidate.DEFAULT_DICTIONARY, + family="physical_full_train", + ) + evaluation = ground_truth.load_hidden_gt_templates( + candidate.DEFAULT_DICTIONARY, + family="physical_full_eval_hard", + ) + train_supports = {frozenset(template["terms"]) for template in train} + eval_supports = {frozenset(template["terms"]) for template in evaluation} + train_terms = {term for template in train for term in template["terms"]} + eval_terms = {term for template in evaluation for term in template["terms"]} + + assert len(evaluation) == 18 + assert train_supports.isdisjoint(eval_supports) + assert eval_terms <= train_terms + assert sum(len(template["terms"]) == 2 for template in evaluation) == 12 + assert sum(len(template["terms"]) == 3 for template in evaluation) == 6 + assert len( + { + frozenset(template["terms"]) + for template in evaluation + if len(template["terms"]) == 2 + } + ) == 12 + assert len( + { + frozenset(template["terms"]) + for template in evaluation + if len(template["terms"]) == 3 + } + ) == 6 + + +def test_physical_full_eval_4000_has_held_out_balanced_supports(): + train = ground_truth.load_hidden_gt_templates( + candidate.DEFAULT_DICTIONARY, + family="physical_full_train", + ) + evaluation = ground_truth.load_hidden_gt_templates( + candidate.DEFAULT_DICTIONARY, + family="physical_full_eval_4000", + ) + train_supports = {frozenset(template["terms"]) for template in train} + eval_supports = {frozenset(template["terms"]) for template in evaluation} + train_terms = {term for template in train for term in template["terms"]} + eval_terms = {term for template in evaluation for term in template["terms"]} + + assert len(evaluation) == 25 + assert len(eval_supports) == 25 + assert train_supports.isdisjoint(eval_supports) + assert eval_terms <= train_terms + assert sum(len(template["terms"]) == 2 for template in evaluation) == 17 + assert sum(len(template["terms"]) == 3 for template in evaluation) == 8 + + +def test_stratified_eval_assignments_balance_every_hidden_template(): + assignments = data_generator.build_stratified_eval_template_assignments( + test_size=288, + eval_pack_size=8, + eval_template_count=18, + ) + + assert assignments is not None + assert len(assignments) == 288 + template_row_counts = {template_index: 0 for template_index in range(18)} + for pack_index in range(36): + pack = assignments[pack_index * 8 : (pack_index + 1) * 8] + assert {row["eval_pack_index"] for row in pack} == {pack_index} + assert len({row["ground_truth_template_index"] for row in pack}) == 1 + assert len({row["ground_truth_template_instance"] for row in pack}) == 1 + template_row_counts[pack[0]["ground_truth_template_index"]] += len(pack) + + assert set(template_row_counts.values()) == {16} + assert { + row["ground_truth_template_instance"] for row in assignments + } == {0, 1} + + +def test_larger_stratified_eval_assignments_use_four_packs_per_template(): + assignments = data_generator.build_stratified_eval_template_assignments( + test_size=576, + eval_pack_size=8, + eval_template_count=18, + ) + + assert assignments is not None + assert len(assignments) == 576 + template_row_counts = Counter( + row["ground_truth_template_index"] for row in assignments + ) + assert set(template_row_counts) == set(range(18)) + assert set(template_row_counts.values()) == {32} + assert { + row["ground_truth_template_instance"] for row in assignments + } == {0, 1, 2, 3} + + +def test_online_eval_subset_keeps_one_fixed_pack_per_template(): + assignments = data_generator.build_stratified_eval_template_assignments( + test_size=576, + eval_pack_size=8, + eval_template_count=18, + ) + + assert assignments is not None + for row in assignments: + row["pde_environment_seed"] = row["eval_pack_index"] + online_assignments = data_generator.select_eval_template_instances( + assignments, + eval_pack_size=8, + eval_template_count=18, + instances_per_template=1, + ) + + assert len(online_assignments) == 144 + assert { + row["ground_truth_template_instance"] for row in online_assignments + } == {0} + template_row_counts = Counter( + row["ground_truth_template_index"] for row in online_assignments + ) + assert set(template_row_counts) == set(range(18)) + assert set(template_row_counts.values()) == {8} + + +def test_stratified_eval_assignments_reject_unbalanced_pack_counts(): + try: + data_generator.build_stratified_eval_template_assignments( + test_size=144, + eval_pack_size=8, + eval_template_count=12, + ) + except ValueError as exc: + assert "must be divisible" in str(exc) + else: + raise AssertionError("Expected unbalanced eval template coverage to fail") + + +def test_fixed_ground_truth_template_index_selects_requested_support(): + templates = ground_truth.load_hidden_gt_templates( + candidate.DEFAULT_DICTIONARY, + family="physical_full_eval_hard", + ) + + for template_index, template in enumerate(templates): + terms, coefficients = ground_truth.sample_hidden_reaction( + np.random.default_rng(1000 + template_index), + templates, + min_terms=1, + max_terms=3, + template_index=template_index, + ) + assert set(terms) == set(template["terms"]) + assert set(coefficients) == set(template["terms"]) + + try: + ground_truth.sample_hidden_reaction( + np.random.default_rng(0), + templates, + min_terms=1, + max_terms=3, + template_index=len(templates), + ) + except ValueError as exc: + assert "outside the eligible template range" in str(exc) + else: + raise AssertionError("Expected an out-of-range template index to fail") + + +def test_initial_condition_amplitudes_follow_trajectory_count(): + amplitude_upper = ground_truth.calibrate_family_initial_amplitude_upper( + ground_truth.load_hidden_gt_templates( + candidate.DEFAULT_DICTIONARY, + family="physical_full_train", + ), + min_terms=1, + max_terms=2, + state_abs_limit=INITIAL_CONDITION_TEST_STATE_ABS_LIMIT, + ) + for trajectory_count in INITIAL_CONDITION_TEST_TRAJECTORY_COUNTS: + amplitudes = pde_numeric.initial_condition_amplitudes( + INITIAL_CONDITION_TEST_PACK_SEED, + trajectory_count, + amplitude_upper=amplitude_upper, + ) + band_edges = np.linspace( + 0.0, + amplitude_upper, + trajectory_count + 1, + ) + + assert len(amplitudes) == trajectory_count + for amplitude, lower, upper in zip( + sorted(amplitudes), band_edges[:-1], band_edges[1:] + ): + assert lower <= amplitude <= upper + assert amplitudes == pde_numeric.initial_condition_amplitudes( + INITIAL_CONDITION_TEST_PACK_SEED, + trajectory_count, + amplitude_upper=amplitude_upper, + ) + + +def test_initial_conditions_obey_sampled_amplitude_and_boundaries(): + x_grid = np.linspace(0.0, 1.0, INITIAL_CONDITION_TEST_GRID_SIZE) + trajectory_count = INITIAL_CONDITION_TEST_TRAJECTORY_COUNTS[-1] + amplitude_upper = ground_truth.calibrate_family_initial_amplitude_upper( + ground_truth.load_hidden_gt_templates( + candidate.DEFAULT_DICTIONARY, + family="physical_full_eval_hard", + ), + min_terms=1, + max_terms=3, + state_abs_limit=INITIAL_CONDITION_TEST_STATE_ABS_LIMIT, + ) + amplitudes = pde_numeric.initial_condition_amplitudes( + INITIAL_CONDITION_TEST_PACK_SEED, + trajectory_count, + amplitude_upper=amplitude_upper, + ) + + for trajectory_index, amplitude in enumerate(amplitudes): + u0 = pde_numeric.initial_condition( + x_grid, + INITIAL_CONDITION_TEST_PACK_SEED, + trajectory_index, + trajectory_count=trajectory_count, + amplitude_upper=amplitude_upper, + shape_config=INITIAL_CONDITION_TEST_SHAPE, + ) + assert u0[0] == 0.0 + assert u0[-1] == 0.0 + assert np.all(u0 >= 0.0) + assert np.isclose(float(np.max(u0)), amplitude) + + +def test_initial_condition_shape_config_parses_mode_range(): + config = pde_numeric.InitialConditionShapeConfig.from_mapping( + { + "mode_count_range": [2, 2], + } + ) + assert config.mode_count_range == (2, 2) + + +def test_reaction_residual_matches_generated_pde_at_terminal_time(): + x_grid = np.linspace(0.0, 1.0, 33) + t_grid = np.linspace(0.0, 0.1, 101) + dx = float(x_grid[1] - x_grid[0]) + dt = float(t_grid[1] - t_grid[0]) + ratio = dt / dx**2 + interior = len(x_grid) - 2 + reaction_fn = lambda u: 2.0 * u + + field = pde_numeric.simulate_dense_trajectory( + x_grid=x_grid, + t_grid=t_grid, + dx=dx, + dt=dt, + lower=-ratio * np.ones(interior - 1), + diag=(1.0 + 2.0 * ratio) * np.ones(interior), + upper=-ratio * np.ones(interior - 1), + trajectory_index=0, + reaction_fn=reaction_fn, + pack_seed=1, + initial_amplitude_upper=1.0, + initial_condition_shape=INITIAL_CONDITION_TEST_SHAPE, + state_abs_limit=INITIAL_CONDITION_TEST_STATE_ABS_LIMIT, + ) + + expected = reaction_fn(field["u"][:, 1:-1]) + np.testing.assert_allclose(field["y"][:, 1:-1], expected, atol=1e-10) + + +def test_regression_error_matches_rendered_candidate(): + u = np.linspace(0.1, 1.0, 100) + y = 1.23456 * u + result = regression.sparse_regression_result( + u=u, + y=y, + u_objective=u, + y_objective=y, + dictionary=["u"], + alpha=0.0, + threshold=0.0, + max_reaction_terms=1, + default_blind_penalty=100.0, + kappa_threshold=100.0, + ) + coefficients = candidate.candidate_coefficients(result["equation"], ["u"]) + actual_error = float( + np.max(np.abs(y - candidate.reaction_value(u, coefficients))) + ) + + assert result["max_error"] == actual_error diff --git a/trinity/algorithm/advantage_fn/__init__.py b/trinity/algorithm/advantage_fn/__init__.py index 6b8aa6d3986..135761aaad7 100644 --- a/trinity/algorithm/advantage_fn/__init__.py +++ b/trinity/algorithm/advantage_fn/__init__.py @@ -22,6 +22,7 @@ "jsd": "trinity.algorithm.advantage_fn.jsd_advantage.JSDAdvantage", "clipb": "trinity.algorithm.advantage_fn.clipb_advantage.ClipBAdvantageFn", "clipv": "trinity.algorithm.advantage_fn.clipv_advantage.ClipVAdvantageFn", + "cod": "trinity.algorithm.advantage_fn.cod_advantage.CoDAdvantageFn", }, ) diff --git a/trinity/algorithm/advantage_fn/cod_advantage.py b/trinity/algorithm/advantage_fn/cod_advantage.py new file mode 100644 index 00000000000..5a3fcd08292 --- /dev/null +++ b/trinity/algorithm/advantage_fn/cod_advantage.py @@ -0,0 +1,461 @@ +"""CoD advantage computation.""" +from typing import Dict, List, Optional, Tuple +import math + +import torch + +from trinity.algorithm.advantage_fn.advantage_fn import AdvantageFn +from trinity.buffer.operators import ExperienceOperator +from trinity.common.experience import Experience, group_by +from trinity.utils.metrics import aggregate_metrics + + +def helper_get_reward(exp: Experience, reward_field: str): + if reward_field == "reward": + return exp.reward + elif reward_field == "total_reward": + return exp.info["total_reward"] + else: + raise ValueError(f"Invalid reward_field {reward_field}") + + +class CoDAdvantageFn(AdvantageFn, ExperienceOperator): + """An advantage function dedicated for CoD.""" + + def __init__( + self, + epsilon: float = 1e-6, + enable_step_norm: bool = False, + std_cal_level: str = "group", # 'group' (task-level) or 'batch' or 'none' + std_threshold: Optional[float] = None, + iterative_hint_e2e_causal: bool = False, + e2e_causal_returns_style: str = "mean", + e2e_causal_returns_window: int = -1, + e2e_causal_returns_gamma: float = 1.0, + e2e_causal_baseline: str = "group-mean", + mask_format_issue_exp: bool = False, + red_weight_temp: Optional[float] = None, + red_weight_adaptive_temp: bool = False, + red_weight_adaptive_version: Optional[str] = None, + red_weight_adaptive_level: Optional[str] = None, + red_weight_adv_shift: bool = False, + **kwargs, + ) -> None: + """Initialize the CoD advantage function. + + Args: + --- original multi-step grpo advantage --- + epsilon (float): A small value to avoid division by zero. + enable_step_norm (bool): If True, normalize advantages by trajectory length. + std_cal_level (str): The scope for calculating reward standard deviation. + 'group' (default): Std is calculated per task group. + 'batch': Std is calculated across all last-step rewards in the entire batch. + 'none': no Std calculation or advantage normalization. + The mean is always calculated per task group. + std_threshold (Optional[float]): If provided, task groups with a reward standard deviation + equal or below this threshold will be skipped. + + --- CoD-specific --- + + iterative_hint_e2e_causal (bool): If True, use iterative hint with end-to-end meta-rl and advantages that respect causality. Options: + - {"style": "mean"} + - {"style": "sliding_window_mean", "window": int} + - {"style": "discounted", "gamma": float} + - e2e_causal_baseline: "group-mean" + - mask_format_issue_exp (bool): if true, mask exp with format issue & positive score + """ + self.epsilon = epsilon + self.enable_step_norm = enable_step_norm + self.std_cal_level = std_cal_level + self.std_threshold = std_threshold + if self.std_cal_level not in ["group", "batch", "none"]: + raise ValueError("std_cal_level must be either 'group' or 'batch' or 'none'") + self.iterative_hint_e2e_causal = iterative_hint_e2e_causal + self.e2e_causal_returns_style = e2e_causal_returns_style + self.e2e_causal_returns_window = e2e_causal_returns_window + self.e2e_causal_returns_gamma = e2e_causal_returns_gamma + self.e2e_causal_baseline = e2e_causal_baseline + self.mask_format_issue_exp = mask_format_issue_exp # DEPRECATED, doesn't take effect + self.red_weight_temp = red_weight_temp + self.red_weight_adaptive_temp = red_weight_adaptive_temp + self.red_weight_adaptive_version = red_weight_adaptive_version + self.red_weight_adaptive_level = red_weight_adaptive_level + self.red_weight_adv_shift = red_weight_adv_shift + + + def calculate_last_step_advantage( + self, + exps: Dict[str, Experience], + precomputed_std: Optional[torch.Tensor] = None, + reward_field: str = "reward", + ) -> Tuple[Dict[str, float], Dict[str, float], bool]: + """Calculate group advantage for a given group of experiences. + + Args: + exps (Dict[str, Experience]): One experience per run, keyed by run ID. + precomputed_std (Optional[torch.Tensor]): Precomputed standard deviation for batch-level calculation. + reward_field: which field to use as reward (for score and advantage calculation). + + Returns: + Dict[str, float]: Scores for each run. + Dict[str, float]: Metrics for logging. + bool: Whether this group should be skipped. + """ + with torch.no_grad(): + if len(exps) == 1: + group_reward_mean = torch.tensor(0.0) + group_reward_std = torch.tensor(1.0) + else: + rewards = torch.tensor([helper_get_reward(exp, reward_field) for exp in exps.values()], dtype=torch.float32) + group_reward_mean = torch.mean(rewards) + group_reward_std = torch.std(rewards) + + # Determine if this group should be skipped based on std_threshold + should_skip = False + if self.std_threshold is not None: + if len(exps) == 1 or group_reward_std <= self.std_threshold: + should_skip = True + + scores = {} + for rid, exp in exps.items(): + exp_reward = helper_get_reward(exp, reward_field) + if self.std_cal_level == "batch" and precomputed_std is not None: + score = (exp_reward - group_reward_mean) / (precomputed_std + self.epsilon) + elif self.std_cal_level == "group": + score = (exp_reward - group_reward_mean) / (group_reward_std + self.epsilon) + elif self.std_cal_level == "none": + score = exp_reward - group_reward_mean + else: + raise ValueError(f"Invalid std_cal_level '{self.std_cal_level}'.") + scores[rid] = score.item() + + # Use standard task rewards for metrics + standard_rewards = torch.tensor([exp.reward for exp in exps.values()], dtype=torch.float32) + standard_reward_mean = torch.mean(standard_rewards) + standard_reward_std = torch.std(standard_rewards) + metrics = { + "reward_mean": standard_reward_mean.item(), + "reward_std": standard_reward_std.item(), + } + + return scores, metrics, should_skip + + def broadcast_advantages( + self, run_exps: Dict[str, List[Experience]], scores: Dict[str, float] + ) -> Dict[str, List[Experience]]: + """Broadcast the calculated advantages to all previous steps in each run. + + Args: + run_exps (Dict[str, List[Experience]]): Experiences grouped by run ID. + scores (Dict[str, float]): Calculated scores for each run. + + Returns: + Dict[str, List[Experience]]: Updated experiences with advantages broadcasted. + """ + for run_id, exps in run_exps.items(): + score = scores[run_id] + traj_length = len(exps) + for exp in exps: + exp.advantages = exp.action_mask * score # type: ignore [operator] + if self.enable_step_norm: + exp.advantages /= traj_length + exp.returns = exp.advantages.clone() + return run_exps + + def process_standard_grpo_custom_score( + self, + exps: List[Experience], + reward_field: str = "reward", + ) -> Tuple[List[Experience], Dict]: + """Standard GRPO with custom field as reward, e.g., per-task reward or total reward in end-to-end meta-rl. + + reward_field: which field to use as reward (for score and advantage calculation) + - "reward": use exp.reward + - "total_reward": use exp.info["total_reward"], tailored to iterative_hint_e2e + """ + if len(exps) == 0: + return [], {} + cnt = 0 + metric_list = [] + filtered_count = 0 + # Step 1: split the experiences into sub-groups by task + task_exps = group_by(exps, "task") + + # --- Pre-computation step for batch-level standard deviation --- + precomputed_std = None + if self.std_cal_level == "batch": + all_laststep_rewards = [] + for task_exp in task_exps.values(): + # First, group all experiences by run to find the last step of each run + task_run_exps = group_by(task_exp, "run") + # Collect rewards from the last step of every run in the entire batch + last_step_rewards = [ + helper_get_reward(run_steps[-1], reward_field) for run_steps in task_run_exps.values() if run_steps + ] + all_laststep_rewards.extend(last_step_rewards) + + if len(all_laststep_rewards) <= 1: + precomputed_std = torch.tensor(1.0) + else: + precomputed_std = torch.std(torch.tensor(all_laststep_rewards, dtype=torch.float32)) + # --- End of pre-computation --- + + # Step 2: further split each task's experiences into sub-groups by run + result_exps = [] + total_task_groups = len(task_exps) + skipped_task_groups = 0 + + for task_exp in task_exps.values(): + run_exps = group_by(task_exp, "run") + + # Step3: extract the last experience (last step) from each run and calculate scores + last_step_exps = {run_id: step_exps[-1] for run_id, step_exps in run_exps.items()} + scores, metrics, should_skip = self.calculate_last_step_advantage( + last_step_exps, + precomputed_std=precomputed_std, + reward_field=reward_field, + ) + + # Skip this task group if std is below threshold + if should_skip: + # Count all experiences in this task group as filtered + task_exp_count = sum(len(step_exps) for step_exps in run_exps.values()) + filtered_count += task_exp_count + skipped_task_groups += 1 + metric_list.append(metrics) + continue + + metric_list.append(metrics) + + # Step 4: broadcast the advantages to all previous steps + run_exps = self.broadcast_advantages(run_exps, scores) + for exps in run_exps.values(): + cnt += len(exps) + result_exps.extend(exps) + + metrics = aggregate_metrics(metric_list, prefix="group_advantages") + metrics["experience_count"] = cnt + metrics["filtered_count"] = filtered_count + + # Calculate the ratio of skipped task groups + if total_task_groups > 0: + metrics["skipped_group_ratio"] = skipped_task_groups / total_task_groups + else: + metrics["skipped_group_ratio"] = 0.0 + + return result_exps, metrics + + def process_standard_grpo(self, exps: List[Experience]) -> Tuple[List[Experience], Dict]: + """Standard GRPO advantage, using exp.reward as reward field.""" + return self.process_standard_grpo_custom_score(exps, "reward") + + def process_iterative_hint_e2e_causal(self, exps: List[Experience]) -> Tuple[List[Experience], Dict]: + """Refined advantage for cross-task e2e case, respecting causal relationship. + + Assume each exp has exp.info fields: + - "exp_type": "solve_task" or "gen_hint" + - "task_idx" (0, 1, 2, ...) or "hint_idx" (0, 1, ...) + - "task_rewards": List[float], list of task/episode-wise outcome rewards within the same meta-trajectory + """ + metrics = dict() + metrics["experience_count"] = len(exps) + + # Drop trajectories the judge flagged (own- or trajectory-level). + def _is_bad(e: Experience) -> bool: + return ( + e.info.get("judge_format_error", False) + or e.info.get("judge_format_error_in_trajectory", False) + ) + + bad_count = sum(1 for e in exps if _is_bad(e)) + exps = [e for e in exps if not _is_bad(e)] + metrics["judge_format_error_count"] = bad_count + + # Step 1. For each exp, calculate exp.info["returns_norm"], plus concise record for Step 2 + record_returns = dict() + + for exp in exps: + task_rewards = exp.info["task_rewards"] # List[float] + exp_type = exp.info["exp_type"] + + # Get rewards-to-go (for gen_hint exp, also include current exp's own reward) + if exp_type == "solve_task": + idx = int(exp.info["task_idx"]) + rewards_to_go = task_rewards[idx:] + elif exp_type == "gen_hint": + idx = int(exp.info["hint_idx"]) + rewards_to_go = [exp.reward] + task_rewards[(idx + 1) : ] + else: + raise ValueError(f"Invalid exp_type {exp_type}, something went wrong.") + rewards_to_go = [float(r) for r in rewards_to_go] # mitigate dtype error in torch.tensor/sum below + + # Calculate returns + e2e_causal_returns_style = self.e2e_causal_returns_style + if e2e_causal_returns_style == "mean": + exp_returns_norm = torch.mean(torch.tensor(rewards_to_go)).item() + elif e2e_causal_returns_style == "sliding_window_mean": + window = int(self.e2e_causal_returns_window) + assert window >= 1, f"Expect sliding-window size >= 1, get {window}." + rewards_to_go_window = rewards_to_go[ : min(window, len(rewards_to_go))] + exp_returns_norm = torch.mean(torch.tensor(rewards_to_go_window)).item() + elif e2e_causal_returns_style == "discounted": + gamma = float(self.e2e_causal_returns_gamma) + assert 0.0 <= gamma <= 1.0, f"Expect gamma within range [0, 1], get {gamma}." + rewards_weights = [gamma ** i for i in range(len(rewards_to_go))] + exp_returns_norm = torch.sum(torch.tensor(rewards_to_go) * torch.tensor(rewards_weights)).item() + else: + raise ValueError(f"Invalid e2e_causal_returns_style: {e2e_causal_returns_style}.") + + exp.info["returns_norm"] = exp_returns_norm + + key_for_record = "_".join([str(exp.eid.task), exp_type, str(idx)]) + exp.info["key_for_record"] = key_for_record + if key_for_record not in record_returns: + record_returns[key_for_record] = dict() + if exp.eid.run not in record_returns[key_for_record]: + record_returns[key_for_record][exp.eid.run] = exp_returns_norm + + # Step 2. Calculate episode-wise baseline within each meta-task, then calculate advantage for each exp + episode_wise_baselines = dict() + for key, value in record_returns.items(): + if self.e2e_causal_baseline == "group-mean": + baseline = torch.mean(torch.tensor(list(value.values()))).item() + else: + raise ValueError(f"Invalid e2e_causal_baseline: {self.e2e_causal_baseline}") + episode_wise_baselines[key] = baseline + + # Compute per-exp score (= returns_norm - baseline) + for exp in exps: + key = exp.info["key_for_record"] + baseline = episode_wise_baselines[key] + exp.info["e2e_causal_score"] = exp.info["returns_norm"] - baseline + exp.info["e2e_causal_resp_len"] = torch.sum(exp.action_mask).item() + + # Helper functions for re-weighting + def estimate_adv_sum(lst_scores, lst_resp_len, temp): + if temp is None: + return sum([lst_scores[i] * lst_resp_len[i] for i in range(len(lst_scores))]) + return sum( + [lst_scores[i] * math.exp(lst_scores[i] / temp) * lst_resp_len[i] for i in range(len(lst_scores))] + ) + + def helper_process_lst_exps(lst_exps): + # Decide the re-weight temperature: adaptive (set so that token-wise mean + # advantage is non-negative) or the fixed red_weight_temp. + if self.red_weight_adaptive_temp: + weight_temp = None + assert self.red_weight_adaptive_version == "bisection", f"Invalid red_weight_adaptive_version: {self.red_weight_adaptive_version}" + lst_scores = [exp.info["e2e_causal_score"] for exp in lst_exps] + lst_resp_len = [exp.info["e2e_causal_resp_len"] for exp in lst_exps] + if estimate_adv_sum(lst_scores, lst_resp_len, None) < 0: + temp_low = 0.8 # avoid too aggressive weighting + temp_high = 100.0 + for _ in range(20): + temp_mid = (temp_low + temp_high) / 2 + if estimate_adv_sum(lst_scores, lst_resp_len, temp_mid) < 0: + temp_high = temp_mid + else: + temp_low = temp_mid + if temp_high - temp_low < 0.05: + break + weight_temp = temp_mid + + else: + weight_temp = self.red_weight_temp + if weight_temp: + assert weight_temp > 0.0, "red_weight_temp must be positive float." + + # Calculate advantage shift to achieve zero token-mean + adv_mean = None + if self.red_weight_adv_shift: + tensor_scores = torch.tensor([exp.info["e2e_causal_score"] for exp in lst_exps], dtype=torch.float32) + tensor_resp_len = torch.tensor([exp.info["e2e_causal_resp_len"] for exp in lst_exps], dtype=torch.float32) + adv_mean = torch.sum(tensor_scores * tensor_resp_len).item() / torch.sum(tensor_resp_len).item() + + # Adaptively re-weight or shift scores, then set advantage and returns + for exp in lst_exps: + score = exp.info["e2e_causal_score"] + if weight_temp: + score *= math.exp(score / weight_temp) + if adv_mean: + if adv_mean < 0: + score -= adv_mean + exp.advantages = exp.action_mask * score + exp.returns = exp.action_mask * exp.info["returns_norm"] + + return lst_exps, weight_temp + + # Group exps by specified level, then process exps for each group + if self.red_weight_adaptive_level == "task": + grouped_exps = group_by(exps, "task") + elif self.red_weight_adaptive_level == "taskset": + def group_by_taskset_id(experiences: List[Experience]) -> Dict[str, List[Experience]]: + """Group experiences by taskset_id. (Modified from `group_by`)""" + grouped = {} + for exp in experiences: + group_id = str(exp.info["taskset_id"]) # taskset_id: int + if group_id not in grouped: + grouped[group_id] = [] + grouped[group_id].append(exp) + return grouped + grouped_exps = group_by_taskset_id(exps) + else: + grouped_exps = {"0": exps} + + num_groups = len(grouped_exps) + print(f"!!! CoD-advantage debug: number of groups = {num_groups} !!!") + + result_exps = [] + weight_temps = [] + for lst_exps in grouped_exps.values(): + processed_lst_exps, weight_temp = helper_process_lst_exps(lst_exps) + result_exps.extend(processed_lst_exps) + weight_temps.append(weight_temp) + + # Diagnostic metric for re-weighting mechanism + active_weight_temps = [temp for temp in weight_temps if temp is not None] + if len(active_weight_temps) > 0: + metrics["active_weight_temp_mean"] = sum(active_weight_temps) / len(active_weight_temps) + else: + metrics["active_weight_temp_mean"] = 0.0 + + return result_exps, metrics + + def process(self, exps: List[Experience]) -> Tuple[List[Experience], Dict]: + if self.iterative_hint_e2e_causal: + return self.process_iterative_hint_e2e_causal(exps) + else: + return self.process_standard_grpo(exps) + + def __call__(self, exps, **kwargs): + """Not used, `process` is the method being called.""" + if self.iterative_hint_e2e_causal: + return self.process_iterative_hint_e2e_causal(exps) + else: + return self.process_standard_grpo(exps) + + @classmethod + def compute_in_trainer(cls) -> bool: + """Whether the advantage should be computed in the trainer loop.""" + return False + + @classmethod + def default_args(cls) -> Dict: + """Return the default configuration for this strategy.""" + return { + "epsilon": 1e-6, + "enable_step_norm": False, + "std_threshold": None, + "std_cal_level": "group", + "iterative_hint_e2e_causal": False, + "e2e_causal_returns_style": "mean", + "e2e_causal_returns_window": -1, + "e2e_causal_returns_gamma": 1.0, + "e2e_causal_baseline": "group-mean", + "mask_format_issue_exp": False, + "red_weight_temp": None, + "red_weight_adaptive_temp": False, + "red_weight_adaptive_version": None, + "red_weight_adaptive_level": None, + "red_weight_adv_shift": False, + } diff --git a/trinity/algorithm/advantage_fn/multi_step_grpo_advantage.py b/trinity/algorithm/advantage_fn/multi_step_grpo_advantage.py index 80b4d568c52..5f6b5f044e7 100644 --- a/trinity/algorithm/advantage_fn/multi_step_grpo_advantage.py +++ b/trinity/algorithm/advantage_fn/multi_step_grpo_advantage.py @@ -20,7 +20,7 @@ def __init__( self, epsilon: float = 1e-6, enable_step_norm: bool = False, - std_cal_level: str = "group", # 'group' (task-level) or 'batch' + std_cal_level: str = "group", # 'group' (task-level) or 'batch' or 'none' std_threshold: Optional[float] = None, **kwargs, ) -> None: @@ -32,6 +32,7 @@ def __init__( std_cal_level (str): The scope for calculating reward standard deviation. 'group' (default): Std is calculated per task group. 'batch': Std is calculated across all last-step rewards in the entire batch. + 'none': no Std calculation or advantage normalization. The mean is always calculated per task group. std_threshold (Optional[float]): If provided, task groups with a reward standard deviation equal or below this threshold will be skipped. @@ -40,8 +41,8 @@ def __init__( self.enable_step_norm = enable_step_norm self.std_cal_level = std_cal_level self.std_threshold = std_threshold - if self.std_cal_level not in ["group", "batch"]: - raise ValueError("std_cal_level must be either 'group' or 'batch'") + if self.std_cal_level not in ["group", "batch", "none"]: + raise ValueError("std_cal_level must be either 'group' or 'batch' or 'none'") def calculate_last_step_advantage( self, @@ -78,8 +79,12 @@ def calculate_last_step_advantage( for rid, exp in exps.items(): if self.std_cal_level == "batch" and precomputed_std is not None: score = (exp.reward - group_reward_mean) / (precomputed_std + self.epsilon) - else: + elif self.std_cal_level == "group": score = (exp.reward - group_reward_mean) / (group_reward_std + self.epsilon) + elif self.std_cal_level == "none": + score = exp.reward - group_reward_mean + else: + raise ValueError(f"Invalid std_cal_level '{self.std_cal_level}'.") scores[rid] = score.item() metrics = { "reward_mean": group_reward_mean.item(), diff --git a/trinity/algorithm/algorithm.py b/trinity/algorithm/algorithm.py index 2cdf4cc08e1..252dd695679 100644 --- a/trinity/algorithm/algorithm.py +++ b/trinity/algorithm/algorithm.py @@ -584,7 +584,7 @@ class OnPolicyDistillAlgorithm(AlgorithmType): """ use_critic: bool = False - use_reference: bool = False + use_reference: bool = True # !!! PATCH FOR COD !!! compute_advantage_in_trainer: bool = True # advantage_fn computes from teacher_logprobs can_balance_batch: bool = True schema: str = "experience" diff --git a/trinity/algorithm/policy_loss_fn/rec_policy_loss.py b/trinity/algorithm/policy_loss_fn/rec_policy_loss.py index 6f00ac86c17..0905c74c443 100644 --- a/trinity/algorithm/policy_loss_fn/rec_policy_loss.py +++ b/trinity/algorithm/policy_loss_fn/rec_policy_loss.py @@ -1,7 +1,7 @@ """REC-token policy loss function. """ -from typing import Dict, Tuple +from typing import Dict, Optional, Tuple import torch @@ -22,6 +22,8 @@ def __init__( regularizer: str = "none", regularizer_coef: float = 0.0, temp: float = 1.0, + loss_agg_mode: Optional[str] = "token-mean", + fix_opd_advantage: bool = False, ) -> None: super().__init__(backend=backend) @@ -49,6 +51,7 @@ def __init__( assert self.weight in [ "none", "importance_sampling", + "truncated_importance_sampling", "gspo_importance_sampling", "advantage", ], f"Invalid weight: {self.weight}" @@ -64,6 +67,9 @@ def __init__( self.temp = temp assert self.temp > 0.0, f"Invalid temp: {self.temp}" + self.loss_agg_mode = loss_agg_mode + self.fix_opd_advantage = fix_opd_advantage + def __call__( # type: ignore self, logprob: torch.Tensor, # [batch_size, seq_len] @@ -105,12 +111,21 @@ def __call__( # type: ignore if self.weight == "importance_sampling": advantages = advantages * ratio # importance sampling + elif self.weight == "truncated_importance_sampling": + advantages = advantages * torch.clamp(ratio, 1 - self.epsilon_low_prime, 1 + self.epsilon_high_prime) elif self.weight == "gspo_importance_sampling": advantages = advantages * normalized_seq_ratio elif self.weight == "advantage": weight = torch.exp(advantages / self.temp) advantages = advantages * weight # advantage weighting (unnormalized version) + if self.fix_opd_advantage: + # Fix advantage calculation for running OPD with off-policyness. Rationale: + # - Implementation of OPD advantage in Trinity gives teacher_logprob - old_logprob (assuming kl_coef = 1.0) + # - Targeted OPD advantage = teacher_logprob - logprob.detach() (by current policy) + # - Hence the fix is to add to advantage: old_logprob - logprob.detach() + advantages = advantages + old_logprob - logprob.detach() + pg_losses = -advantages * logprob * is_in_range.float() if self.regularizer == "forward-kl": @@ -121,15 +136,14 @@ def __call__( # type: ignore regularizer_losses = self.regularizer_coef * (logprob - old_logprob).square() pg_losses = pg_losses + regularizer_losses + loss_agg_mode = self.loss_agg_mode if self.clip_mode == "gspo-one-side": - # [EXPERIMENTAL] specialized for gspo-style rec variant for now - pg_loss = aggregate_loss( - values=pg_losses, - mask=action_mask, - loss_agg_mode="seq-mean-token-mean", - ) - else: - pg_loss = masked_mean(pg_losses, action_mask) + loss_agg_mode = "seq-mean-token-mean" + pg_loss = aggregate_loss( + values=pg_losses, + mask=action_mask, + loss_agg_mode=loss_agg_mode, + ) pg_clipfrac = masked_mean(is_clipped_mask.float(), action_mask) metrics = { @@ -150,4 +164,6 @@ def default_args(cls) -> Dict: "regularizer": "none", "regularizer_coef": 0.0, "temp": 1.0, + "loss_agg_mode": "token-mean", + "fix_opd_advantage": False, } diff --git a/trinity/buffer/operators/__init__.py b/trinity/buffer/operators/__init__.py index 35baa7dad21..bdb443745ce 100644 --- a/trinity/buffer/operators/__init__.py +++ b/trinity/buffer/operators/__init__.py @@ -16,6 +16,8 @@ "pass_rate_calculator": "trinity.buffer.operators.mappers.pass_rate_calculator.PassRateCalculator", "data_juicer": "trinity.buffer.operators.data_juicer_operator.DataJuicerOperator", "invalid_reward_filter": "trinity.buffer.operators.filters.reward_filter.InvalidRewardFilter", + "cod_advantage_fn": "trinity.algorithm.advantage_fn.cod_advantage.CoDAdvantageFn", + "cod_overwrite_reward_with_returns": "trinity.buffer.operators.cod_operator.OverwriteRewardWithReturns", }, ) diff --git a/trinity/buffer/operators/cod_operator.py b/trinity/buffer/operators/cod_operator.py new file mode 100644 index 00000000000..575e8d2f997 --- /dev/null +++ b/trinity/buffer/operators/cod_operator.py @@ -0,0 +1,14 @@ +from typing import List, Tuple + +from trinity.buffer.operators import ExperienceOperator +from trinity.common.experience import Experience + + +class OverwriteRewardWithReturns(ExperienceOperator): + """Overwrite exp.reward with exp.info["returns_norm"], dedicated to CoD-PPO algorithm.""" + + def process(self, exps: List[Experience]) -> Tuple[List[Experience], dict]: + result_exps = [exp for exp in exps] + for exp in result_exps: + exp.reward = exp.info["returns_norm"] + return result_exps, {} diff --git a/trinity/common/config.py b/trinity/common/config.py index 6d27cb3af87..ac2c5eee742 100644 --- a/trinity/common/config.py +++ b/trinity/common/config.py @@ -808,7 +808,7 @@ class TrainerConfig: trainer_type: str = "verl" trainer_strategy: str = "fsdp2" # "fsdp", "fsdp2" or "megatron" save_interval: int = 0 - enable_preview: bool = False # enable rollout preview in wandb + enable_preview: bool = True # enable rollout preview in wandb total_steps: Optional[ int ] = None # total training steps, training stops when reaching this step, None means no limit @@ -936,6 +936,19 @@ class StageConfig: trainer: Optional[TrainerConfig] = None +@dataclass +class CoDConfig: + """!!! Config for CoD !!!""" + + # number of original tasks to be packed into one task + task_pack_size: int = 4 + eval_task_pack_size: Optional[int] = None + packing_strategy: str = "cod" + + # cod workflow args: log_dir, exp_name, etc. + cod_workflow_args: dict = field(default_factory=dict) + + @dataclass class Config: """Global Configuration""" @@ -967,6 +980,9 @@ class Config: service: ServiceConfig = field(default_factory=ServiceConfig) log: LogConfig = field(default_factory=LogConfig) + # !!! cod config !!! + cod: CoDConfig = field(default_factory=CoDConfig) + # configurations for different training stages stages: List[StageConfig] = field(default_factory=list) diff --git a/trinity/common/experience.py b/trinity/common/experience.py index ef6c91429f7..aed20d2a6d8 100644 --- a/trinity/common/experience.py +++ b/trinity/common/experience.py @@ -222,9 +222,13 @@ def __init__( # noqa: C901 ), f"Token ids must be larger than the prompt length. Got len(tokens)={len(tokens)}, prompt_length={prompt_length}." action_mask = torch.ones(len(tokens) - prompt_length, dtype=torch.bool) else: + # mask out exp if prompt truncated action_mask = torch.zeros(len(logprobs), dtype=torch.bool) elif experience_type == "dpo": prompt_length = len(tokens) + # Note: For multi_turn, action_mask is passed in from outside. + # Truncation handling for multi_turn (concatenated) is done in convert_messages_to_experience, + # which only masks the incomplete turn, preserving complete turns. if eid is None: self.eid = EID() elif isinstance(eid, dict): diff --git a/trinity/common/models/model.py b/trinity/common/models/model.py index bc3050906cb..1a00cca3fe0 100644 --- a/trinity/common/models/model.py +++ b/trinity/common/models/model.py @@ -353,6 +353,9 @@ async def convert_messages_to_experience( prompt_length=prompt_length, action_mask=action_mask[prompt_length:], # Exclude the prompt tokens messages=messages, + # !!! MODIFICATION FOR COD START !!! + response_text=self.tokenizer.decode(token_ids[prompt_length:]) if prompt_length < len(token_ids) else None, + # !!! MODIFICATION FOR COD END !!! truncate_status=truncate_status, ) diff --git a/trinity/common/models/vllm_model.py b/trinity/common/models/vllm_model.py index b30e83e8287..6d055ab6bd8 100644 --- a/trinity/common/models/vllm_model.py +++ b/trinity/common/models/vllm_model.py @@ -8,7 +8,7 @@ import numpy as np import torch from packaging.version import parse as parse_version -from transformers import AutoProcessor +from transformers import AutoProcessor, AutoTokenizer from trinity.common.config import InferenceModelConfig from trinity.common.constants import SyncMethod @@ -95,10 +95,19 @@ def __init__( async def _initialize_tokenizer(self): if self.tokenizer is None: + # !!! PATCH FOR COD START !!! if self.vllm_version >= parse_version("0.15.0"): - self.tokenizer = self.async_llm.get_tokenizer() + engine_tokenizer = self.async_llm.get_tokenizer() else: - self.tokenizer = await self.async_llm.get_tokenizer() + engine_tokenizer = await self.async_llm.get_tokenizer() + # Standalone tokenizer to avoid the "Already borrowed" race, keeping + # the engine's model_max_length so truncation stays aligned. + self.tokenizer = AutoTokenizer.from_pretrained( + self.config.model_path, + trust_remote_code=self.config.trust_remote_code, + ) + self.tokenizer.model_max_length = engine_tokenizer.model_max_length + # !!! PATCH FOR COD END !!! self.tokenizer.truncation_side = "left" async def _initialize_processor(self): diff --git a/trinity/common/workflows/__init__.py b/trinity/common/workflows/__init__.py index 7627cdca296..cbda8e195c8 100644 --- a/trinity/common/workflows/__init__.py +++ b/trinity/common/workflows/__init__.py @@ -52,6 +52,16 @@ "on_policy_distill_math_workflow": "trinity.common.workflows.on_policy_distill_workflow.OnPolicyDistillMathWorkflow", # custom workflows "sudoku_workflow": "trinity.common.workflows.envs.sudoku.sudoku_workflow.SudokuWorkflow", + # connect-the-dots (cod) + "cod_workflow": "trinity.common.workflows.connect_the_dots.cod_workflow.CoDWorkflow", + # cod workflows (unified prompt management) + "cod_frozenlake_obscure_workflow": "trinity.common.workflows.connect_the_dots.frozen_lake.workflow_obscure.CoDFrozenLakeObscureWorkflow", + "cod_frozenlake_obscure_react_workflow": "trinity.common.workflows.connect_the_dots.frozen_lake.workflow_obscure_react.CoDFrozenLakeObscureReActWorkflow", + "cod_grid_navigation_workflow": "trinity.common.workflows.connect_the_dots.grid_navigation.workflow.CoDGridNavigationWorkflow", + "cod_random_alchemy_workflow": "trinity.common.workflows.connect_the_dots.alchemy.workflow_random.CoDRandomAlchemyWorkflow", + "cod_terminal_workflow": "trinity.common.workflows.connect_the_dots.terminal.workflow.CoDTerminalWorkflow", + "cod_learn2ask_workflow": "trinity.common.workflows.connect_the_dots.learn2ask.workflow.CoDLearn2AskWorkflow", + "cod_optimalcontrol_workflow": "trinity.common.workflows.connect_the_dots.optimalcontrol.workflow.CoDOptimalControlWorkflow", # !!! }, ) diff --git a/trinity/common/workflows/connect_the_dots/agentscope_utils.py b/trinity/common/workflows/connect_the_dots/agentscope_utils.py new file mode 100644 index 00000000000..257441a68c1 --- /dev/null +++ b/trinity/common/workflows/connect_the_dots/agentscope_utils.py @@ -0,0 +1,224 @@ +"""AgentScope adapters used by CoD workflows.""" + +import json +from typing import Any, List, Optional + +try: + from agentscope.agent import AgentBase, ReActAgent + from agentscope.formatter import OpenAIChatFormatter + from agentscope.memory import InMemoryMemory + from agentscope.message import Msg + from agentscope.model import OpenAIChatModel + from agentscope.tool import Toolkit +except Exception as e: + _AGENTSCOPE_IMPORT_ERROR = e +else: + _AGENTSCOPE_IMPORT_ERROR = None + + +def _ensure_agentscope() -> None: + if _AGENTSCOPE_IMPORT_ERROR is not None: + raise ImportError( + "AgentScope is not installed or failed to import. Please install it " + "before running CoD AgentScope workflows." + ) from _AGENTSCOPE_IMPORT_ERROR + + +def _build_model_generate_kwargs(model) -> dict: + """Build AgentScope generate kwargs from rollout model config.""" + config = model.config + kwargs = { + "temperature": config.temperature, + "top_p": config.top_p, + "max_tokens": config.max_response_tokens, + } + extra_body = {} + if config.top_k is not None: + extra_body["top_k"] = config.top_k + if config.min_response_tokens is not None: + extra_body["min_tokens"] = config.min_response_tokens + if config.repetition_penalty is not None: + extra_body["repetition_penalty"] = config.repetition_penalty + if extra_body: + kwargs["extra_body"] = extra_body + return kwargs + + +def _build_rollout_generate_kwargs(rollout_args) -> dict: + """Build AgentScope generate kwargs from task rollout_args.""" + kwargs = { + "temperature": rollout_args.temperature, + "top_p": rollout_args.top_p, + } + if rollout_args.max_tokens is not None: + kwargs["max_tokens"] = rollout_args.max_tokens + if rollout_args.top_k != -1: + kwargs["extra_body"] = {"top_k": rollout_args.top_k} + return kwargs + + +def _build_agentscope_chat_model(model_path: str, client, generate_kwargs: dict): + """Build an AgentScope OpenAIChatModel with Trinity's recording client.""" + _ensure_agentscope() + + chat_model = OpenAIChatModel( + api_key="EMPTY", + model_name=model_path, + stream=False, + generate_kwargs=generate_kwargs, + ) + chat_model.client = client + return chat_model + + +async def _build_agentscope_chat_model_from_trinity(model, generate_kwargs: dict): + """Build an AgentScope chat model from a Trinity ModelWrapper.""" + model_path = model.model_path + client = model.get_openai_async_client() + return _build_agentscope_chat_model( + model_path=model_path, + client=client, + generate_kwargs=generate_kwargs, + ) + + +async def build_agentscope_single_turn_agent(*, name: str, model, rollout_args): + """Build an AgentScope agent that maps messages to one model call.""" + _ensure_agentscope() + + class _SingleTurnAgent(AgentBase): + """Single-turn agent: messages -> one model call -> assistant Msg.""" + + def __init__(self, model, formatter): + super().__init__() + self.name = name + self.model = model + self.formatter = formatter + + async def reply(self, messages: List[dict]) -> "Msg": + msgs = [Msg(m["role"], m["content"], m["role"]) for m in messages] + res = await self.model(await self.formatter.format(msgs=msgs)) + return Msg(self.name, res.content, "assistant") + + async def observe(self, msg=None) -> None: + pass + + async def handle_interrupt(self, *args, **kwargs) -> "Msg": + return Msg(self.name, "Interrupted.", "assistant") + + generate_kwargs = _build_rollout_generate_kwargs(rollout_args) + chat_model = await _build_agentscope_chat_model_from_trinity(model, generate_kwargs) + return _SingleTurnAgent(chat_model, OpenAIChatFormatter()) + + +async def build_agentscope_react_agent( + *, + name: str, + model, + system_prompt: str, + compress_assistant_fn, + toolkit=None, + max_iters: int = 1, +): + """Build an AgentScope ReActAgent with compressed assistant text memory.""" + _ensure_agentscope() + + class _CompressingMemory(InMemoryMemory): + """Compress text-only assistant messages without touching tool blocks.""" + + async def add(self, memories, marks=None, allow_duplicates=False, **kwargs): + if memories is not None: + msgs = memories if isinstance(memories, list) else [memories] + out: List[Optional[Msg]] = [] + for msg in msgs: + has_tool_blocks = ( + msg is not None + and ( + msg.has_content_blocks("tool_use") + or msg.has_content_blocks("tool_result") + ) + ) + if ( + msg is not None + and msg.role == "assistant" + and not has_tool_blocks + ): + text = msg.get_text_content() or "" + msg = Msg( + msg.name, + compress_assistant_fn(text), + "assistant", + ) + out.append(msg) + memories = out if isinstance(memories, list) else out[0] + await super().add( + memories, marks=marks, allow_duplicates=allow_duplicates, **kwargs + ) + + generate_kwargs = _build_model_generate_kwargs(model) + chat_model = await _build_agentscope_chat_model_from_trinity(model, generate_kwargs) + return ReActAgent( + name=name, + sys_prompt=system_prompt, + model=chat_model, + formatter=OpenAIChatFormatter(), + toolkit=toolkit if toolkit is not None else Toolkit(), + memory=_CompressingMemory(), + max_iters=max_iters, + ) + + +async def run_agentscope_agent_step(agent, user_content: str) -> str: + """Run one AgentScope reply and return final text.""" + _ensure_agentscope() + + reply = await agent.reply(Msg("user", user_content, role="user")) + return reply.get_text_content() or "" + + +def _format_tool_output(output: Any) -> str: + if isinstance(output, str): + return output + if isinstance(output, list): + parts = [] + for item in output: + if isinstance(item, dict) and item.get("type") == "text": + parts.append(str(item.get("text", ""))) + else: + parts.append(json.dumps(item, ensure_ascii=False)) + return "\n".join(parts) + return json.dumps(output, ensure_ascii=False) + + +def agentscope_msg_to_text(msg) -> str: + """Render one AgentScope message for CoD logs and trajectories.""" + parts = [] + for block in msg.get_content_blocks(): + block_type = block.get("type") + if block_type == "text": + parts.append(block.get("text", "")) + elif block_type == "tool_use": + tool_input = json.dumps(block.get("input", {}), ensure_ascii=False) + parts.append(f"[tool_call] {block.get('name')}({tool_input})") + elif block_type == "tool_result": + output = _format_tool_output(block.get("output", "")) + parts.append(f"[tool_result] {block.get('name')}: {output}") + elif block_type == "thinking": + parts.append(f"[thinking] {block.get('thinking', '')}") + else: + parts.append(f"[{block_type}]") + content = "\n".join(part for part in parts if part) + if not content: + return "" + label = "Tool" if msg.has_content_blocks("tool_result") else msg.role.capitalize() + return f"{label}: {content}" + + +def agentscope_msgs_to_text(messages: List[Any]) -> str: + """Render AgentScope messages as a compact transcript.""" + lines = [] + for msg in messages: + rendered = agentscope_msg_to_text(msg) + if rendered.strip(): + lines.append(rendered) + return "\n".join(lines) diff --git a/trinity/common/workflows/connect_the_dots/alchemy/__init__.py b/trinity/common/workflows/connect_the_dots/alchemy/__init__.py new file mode 100644 index 00000000000..57ef39a4422 --- /dev/null +++ b/trinity/common/workflows/connect_the_dots/alchemy/__init__.py @@ -0,0 +1,5 @@ +# -*- coding: utf-8 -*- +from trinity.common.workflows.connect_the_dots.alchemy.env import AlchemyEnv +from trinity.common.workflows.connect_the_dots.alchemy.common import GameInstance + +__all__ = ["AlchemyEnv", "GameInstance"] diff --git a/trinity/common/workflows/connect_the_dots/alchemy/common.py b/trinity/common/workflows/connect_the_dots/alchemy/common.py new file mode 100644 index 00000000000..584d3a580f9 --- /dev/null +++ b/trinity/common/workflows/connect_the_dots/alchemy/common.py @@ -0,0 +1,57 @@ +# -*- coding: utf-8 -*- +"""Shared types and data structures for Alchemy game variants.""" + +from dataclasses import dataclass +from typing import Dict, List, Optional, Set, Tuple + +# Recipe key: tuple(sorted([e1, e2])) to handle both order and self-combination +RecipeKey = Tuple[str, str] + + +@dataclass +class GameInstance: + """A complete definition of one Alchemy game.""" + + target: str + recipes: Dict[RecipeKey, str] # {("air","air"): "pressure", ...} + starting_inventory: Dict[str, int] # {element: count} available at start + all_elements: Set[str] # All elements in this subgraph + solution_path: List[Tuple[str, str, str]] # [(e1, e2, result), ...] + depth: int # Depth of target element + local_tiers: Dict[str, int] = None # Tiers computed within this subgraph + + def __post_init__(self): + if self.local_tiers is None: + self.local_tiers = self._compute_local_tiers() + # Filter out recipes that violate local tier ordering + self.recipes = { + key: result for key, result in self.recipes.items() + if not (key[0] in self.local_tiers and key[1] in self.local_tiers and result in self.local_tiers) + or self.local_tiers[result] > max(self.local_tiers[key[0]], self.local_tiers[key[1]]) + } + + def _compute_local_tiers(self) -> Dict[str, int]: + """Compute tiers from recipe graph structure. + + Leaf elements (not produced by any recipe) are tier 0. + Other elements: tier = max(tier(input1), tier(input2)) + 1. + """ + all_ingredients: Set[str] = set() + all_results: Set[str] = set() + for key, result in self.recipes.items(): + all_ingredients.update(key) + all_results.add(result) + leaves = (all_ingredients | set(self.starting_inventory)) - all_results + + tiers = {elem: 0 for elem in leaves} + changed = True + while changed: + changed = False + for key, result in self.recipes.items(): + if result in tiers: + continue + e1, e2 = key + if e1 in tiers and e2 in tiers: + tiers[result] = max(tiers[e1], tiers[e2]) + 1 + changed = True + return tiers diff --git a/trinity/common/workflows/connect_the_dots/alchemy/env.py b/trinity/common/workflows/connect_the_dots/alchemy/env.py new file mode 100644 index 00000000000..547e13662ec --- /dev/null +++ b/trinity/common/workflows/connect_the_dots/alchemy/env.py @@ -0,0 +1,159 @@ +# -*- coding: utf-8 -*- +"""Gymnasium environment for the Alchemy crafting game.""" + +from typing import Any, Dict, List, Optional, Tuple + +import gymnasium as gym + +from trinity.common.workflows.connect_the_dots.alchemy.common import GameInstance, RecipeKey + + +class AlchemyEnv(gym.Env): + """Alchemy crafting game environment. + + Supports two material modes: + - "limited": elements are consumed on every attempt (success or failure). + - "unlimited": elements are never consumed; new discoveries are added to inventory. + + Game ends when target is synthesized, round limit is reached, or (limited only) + materials run out. + """ + + metadata = {"render_modes": ["text"]} + + def __init__(self, game_instance: GameInstance, max_rounds: int = 35, + material_mode: str = "limited", render_mode: str = "text"): + super().__init__() + self.game = game_instance + self.max_rounds = max_rounds + self.material_mode = material_mode + self.render_mode = render_mode + self.inventory: Dict[str, int] = {} + self.discovered_recipes: List[Tuple[str, str, str]] = [] + self.failed_combinations: List[Tuple[str, str]] = [] + self.current_round: int = 0 + self._last_result: Optional[str] = None + self._invalid_action: bool = False + self._success: bool = False + + def reset(self, seed=None, options=None) -> Tuple[str, Dict[str, Any]]: + super().reset(seed=seed) + if self.material_mode == "unlimited": + self.inventory = {k: 1 for k in self.game.starting_inventory} + else: + self.inventory = dict(self.game.starting_inventory) + self.discovered_recipes = [] + self.failed_combinations = [] + self.current_round = 0 + self._last_result = None + self._invalid_action = False + self._success = False + return self.render(), self._get_info() + + def step(self, action: Tuple[str, str]) -> Tuple[str, float, bool, bool, Dict[str, Any]]: + raw1, raw2 = action + self.current_round += 1 + self._last_result = None + self._invalid_action = False + + # Normalize to inventory keys (case-insensitive) + inv_lower = {k.lower(): k for k in self.inventory} + elem1 = inv_lower.get(raw1.lower(), raw1) + elem2 = inv_lower.get(raw2.lower(), raw2) + + # Validate elements exist in inventory + if self.material_mode == "unlimited": + # Unlimited: just check both elements exist (no quantity requirement) + if self.inventory.get(elem1, 0) < 1 or self.inventory.get(elem2, 0) < 1: + self._invalid_action = True + else: + if elem1 == elem2: + if self.inventory.get(elem1, 0) < 2: + self._invalid_action = True + else: + if self.inventory.get(elem1, 0) < 1 or self.inventory.get(elem2, 0) < 1: + self._invalid_action = True + + if self._invalid_action: + return self._finish_step() + + # Consume ingredients (limited mode only) + if self.material_mode == "limited": + if elem1 == elem2: + self.inventory[elem1] -= 2 + else: + self.inventory[elem1] -= 1 + self.inventory[elem2] -= 1 + self.inventory = {k: v for k, v in self.inventory.items() if v > 0} + + # Lookup recipe + key: RecipeKey = tuple(sorted([elem1.lower(), elem2.lower()])) + result = self.game.recipes.get(key) + + if result is not None: + self._last_result = result + if self.material_mode == "unlimited": + self.inventory[result] = 1 + else: + self.inventory[result] = self.inventory.get(result, 0) + 1 + self.discovered_recipes.append((elem1, elem2, result)) + else: + self.failed_combinations.append((elem1, elem2)) + + return self._finish_step() + + def _finish_step(self) -> Tuple[str, float, bool, bool, Dict[str, Any]]: + self._success = self.game.target in self.inventory + terminated = self._success + truncated = ( + not terminated + and (self.current_round >= self.max_rounds or not self.get_feasible_actions()) + ) + reward = 1.0 if self._success else 0.0 + return self.render(), reward, terminated, truncated, self._get_info() + + def get_feasible_actions(self) -> List[Tuple[str, str]]: + elems = sorted(self.inventory.keys()) + pairs = [] + for i, e1 in enumerate(elems): + for e2 in elems[i:]: + if self.material_mode == "unlimited": + # Unlimited: any pair is feasible as long as elements exist + pairs.append((e1, e2)) + else: + if e1 == e2: + if self.inventory[e1] >= 2: + pairs.append((e1, e2)) + else: + pairs.append((e1, e2)) + return pairs + + def render(self) -> str: + inv_str = ", ".join(f"{e} x{c}" for e, c in sorted(self.inventory.items())) + lines = [ + f"Target: {self.game.target}", + f"Round: {self.current_round}/{self.max_rounds}", + f"Inventory: {inv_str}" if inv_str else "Inventory: (empty)", + ] + if self.discovered_recipes: + lines.append("Discovered recipes:") + for e1, e2, r in self.discovered_recipes: + lines.append(f" {e1} + {e2} = {r}") + if self.failed_combinations: + lines.append("Failed combinations:") + for e1, e2 in self.failed_combinations: + lines.append(f" {e1} + {e2}") + return "\n".join(lines) + + def _get_info(self) -> Dict[str, Any]: + return { + "inventory": dict(self.inventory), + "discovered_recipes": list(self.discovered_recipes), + "failed_combinations": list(self.failed_combinations), + "feasible_actions_count": len(self.get_feasible_actions()), + "target": self.game.target, + "current_round": self.current_round, + "success": self._success, + "last_result": self._last_result, + "invalid_action": self._invalid_action, + } diff --git a/trinity/common/workflows/connect_the_dots/alchemy/prompts/__init__.py b/trinity/common/workflows/connect_the_dots/alchemy/prompts/__init__.py new file mode 100644 index 00000000000..ea26a9597c8 --- /dev/null +++ b/trinity/common/workflows/connect_the_dots/alchemy/prompts/__init__.py @@ -0,0 +1,43 @@ +# -*- coding: utf-8 -*- +"""Prompt management for CoD Alchemy workflow using Jinja2 templates.""" + +from pathlib import Path +from typing import Dict, Optional + +from jinja2 import Environment, FileSystemLoader + +PROMPTS_DIR = Path(__file__).parent + + +def get_jinja_env() -> Environment: + return Environment( + loader=FileSystemLoader(PROMPTS_DIR), + trim_blocks=True, + lstrip_blocks=True, + ) + + +def load_system_prompt(**kwargs) -> str: + env = get_jinja_env() + template = env.get_template("system.jinja2") + return template.render(**kwargs) + + +def load_user_prompt( + current_round: int, + max_rounds: int, + inventory: Dict[str, int], + target: str, + action_feedback: Optional[str] = None, + **kwargs, +) -> str: + env = get_jinja_env() + template = env.get_template("user.jinja2") + return template.render( + current_round=current_round, + max_rounds=max_rounds, + inventory=inventory, + target=target, + action_feedback=action_feedback, + **kwargs, + ) diff --git a/trinity/common/workflows/connect_the_dots/alchemy/prompts/system.jinja2 b/trinity/common/workflows/connect_the_dots/alchemy/prompts/system.jinja2 new file mode 100644 index 00000000000..99d82e4eb06 --- /dev/null +++ b/trinity/common/workflows/connect_the_dots/alchemy/prompts/system.jinja2 @@ -0,0 +1,61 @@ +You are in an Alchemy crafting environment. Your goal is to synthesize the target element "{{ target }}" by combining elements in your inventory. + +## Game Mechanics + +There exists a fixed set of recipes. Each recipe takes exactly two input elements and produces one output element. A recipe's two inputs can be the same or different elements, and can be from different tiers. + +- **Combination**: Each round, choose two elements from your inventory to combine. If the pair matches a valid recipe, the result is added to your inventory. If not, nothing is produced. +- **Tier system**: Every element has a tier. Every recipe produces a result whose tier is strictly higher than both inputs' tiers. +{% if material_mode == "unlimited" %} +- **No consumption**: Elements are NOT consumed when used. You keep all elements after each attempt, whether it succeeds or fails. Newly discovered elements are permanently added to your inventory. +{% else %} +- **Consumption**: Every attempt consumes 1 copy of each input element, whether it succeeds or fails. If both inputs are the same element, 2 copies are consumed. Failed attempts lose materials with no result. +{% endif %} +- **Winning**: You win when "{{ target }}" appears in your inventory. +{% if material_mode == "unlimited" %} +- **Ending**: The game ends when you succeed or reach the round limit. +{% else %} +- **Ending**: The game ends when you succeed, run out of materials, or reach the round limit. +{% endif %} +{% if obfuscate_names %} +- All element names in this Alchemy crafting environment have no correspondence to real-world concepts. Do not try to guess recipes based on element names. +{% endif %} +{% if not show_recipes %} + +## Hidden Clues About the Environment +Recipes (i.e., which elements can be combined and which cannot) are unknown initially. You can try to discover them through experimentation, or based on additional hints provided in the context.{% if not obfuscate_names %} You may also use your knowledge of how elements naturally combine to make educated guesses.{% endif %} + +{% endif %} +{% if show_recipes %} + +## Available Recipes +{% for key, result in recipes.items() %} +- {{ key[0] }} + {{ key[1] }} = {{ result }} +{% endfor %} +{% endif %} +{% if show_elements %} + +## Elements in This World +{{ all_elements | join(', ') }} +{% endif %} +{% if show_tiers %} + +## Element Tiers + +{{ tier_info }} +{% endif %} + +## Response Format +First think about which combination of elements to try out in the current round, then provide your action using ... tags at the end. Format your complete response as follows: +``` +[THINKING] +Element1 + Element2 +``` +where [THINKING] should be replaced with your thinking process, while Element1 and Element2 are elements (same or different) chosen from your inventory. + +Other requirements: +- Be concise and avoid overthinking in your thinking process. +- Take one and only one action for each round. +- You should synthesize the target element within a limited number of rounds. Try to find a short path towards the goal, e.g., by leveraging recipes that have been revealed and avoiding combinations that are known to fail. +- Balance exploration and exploitation: as you strive to reach the goal, you may acquire information about the environment along the way, which can be potentially helpful for solving other similar tasks later on. + diff --git a/trinity/common/workflows/connect_the_dots/alchemy/prompts/user.jinja2 b/trinity/common/workflows/connect_the_dots/alchemy/prompts/user.jinja2 new file mode 100644 index 00000000000..b984b827405 --- /dev/null +++ b/trinity/common/workflows/connect_the_dots/alchemy/prompts/user.jinja2 @@ -0,0 +1,22 @@ +Round {{ current_round }}/{{ max_rounds }} +{% if action_feedback %} + +{{ action_feedback }} +{% endif %} + +Target: {{ target }} +{% if material_mode == "unlimited" %} + +Available elements in your inventory: +{% for elem in inventory.keys() | sort %} +- {{ elem }} +{% endfor %} +{% else %} + +Inventory ({{ inventory.values() | sum }} total copies): +{% for elem, count in inventory.items() | sort %} +- {{ elem }} x{{ count }} +{% endfor %} +{% endif %} + +Choose two elements available in your inventory to combine. diff --git a/trinity/common/workflows/connect_the_dots/alchemy/random_graph.py b/trinity/common/workflows/connect_the_dots/alchemy/random_graph.py new file mode 100644 index 00000000000..e7b92d1ad77 --- /dev/null +++ b/trinity/common/workflows/connect_the_dots/alchemy/random_graph.py @@ -0,0 +1,432 @@ +# -*- coding: utf-8 -*- +"""Procedural graph generator for Random Alchemy game. + +Generates DAG-structured recipe graphs from scratch without external data. +Each graph has tiered elements, configurable branching, cross-tier connections, +and noise (distractor) elements. +""" + +import random +import string +from collections import defaultdict +from typing import Dict, List, Optional, Set, Tuple + +from trinity.common.workflows.connect_the_dots.alchemy.common import ( + GameInstance, + RecipeKey, +) + + +def _generate_name(rng: random.Random, existing: Set[str], length_min: int = 3, length_max: int = 5) -> str: + """Generate a unique random element name (lowercase to match env recipe lookup).""" + while True: + length = rng.randint(length_min, length_max) + name = "".join(rng.choices(string.ascii_lowercase, k=length)) + if name not in existing: + return name + + +class RandomGraphGenerator: + """Procedural DAG generator for alchemy-style games. + + Generates a complete recipe graph with: + - Tiered elements (tier 0 = base, higher = synthesized) + - Configurable branching and cross-tier connections + - Noise (distractor) elements that are dead ends + """ + + def generate_graph( + self, + rng: random.Random, + num_tiers: int, + base_elements: int, + tier_shrink: float, + min_recipes_per_element: int, + max_recipes_per_element: int, + max_tier_gap: int, + cross_tier_prob: float, + noise_node_ratio: float, + noise_chain_depth: int, + ) -> Tuple[Dict[int, List[str]], Dict[RecipeKey, str], Dict[str, int], Set[str]]: + """Generate a complete recipe graph. + + Returns: + (tier_elements, recipes, local_tiers, noise_elements) + - tier_elements: {tier: [element_names]} (includes noise) + - recipes: {(e1, e2): result} + - local_tiers: {element: tier} + - noise_elements: set of element names that are distractors + """ + all_names: Set[str] = set() + + # Step 1: Determine elements per tier + tier_sizes = [base_elements] + for t in range(1, num_tiers): + size = max(1, round(tier_sizes[t - 1] * tier_shrink)) + tier_sizes.append(size) + + # Step 2: Generate element names per tier + tier_elements: Dict[int, List[str]] = {} + for t in range(num_tiers): + tier_elements[t] = [] + for _ in range(tier_sizes[t]): + name = _generate_name(rng, all_names) + all_names.add(name) + tier_elements[t].append(name) + + # Step 3: Generate recipes for each non-leaf element + recipes: Dict[RecipeKey, str] = {} + used_keys: Set[RecipeKey] = set() + + for t in range(1, num_tiers): + for elem in tier_elements[t]: + num_recipes = rng.randint(min_recipes_per_element, max_recipes_per_element) + for _ in range(num_recipes): + key = self._make_recipe( + rng, elem, t, tier_elements, max_tier_gap, + cross_tier_prob, used_keys, + ) + if key is not None: + recipes[key] = elem + used_keys.add(key) + + # Step 4: Add noise elements and recipes + noise_elements: Set[str] = set() + if noise_node_ratio > 0: + noise_elements = self._add_noise( + rng, tier_elements, recipes, used_keys, all_names, + noise_node_ratio, noise_chain_depth, max_tier_gap, cross_tier_prob, + ) + + # Step 5: Compute local tiers + local_tiers: Dict[str, int] = {} + for t, elems in tier_elements.items(): + for e in elems: + local_tiers[e] = t + + return tier_elements, recipes, local_tiers, noise_elements + + def _make_recipe( + self, + rng: random.Random, + result_elem: str, + result_tier: int, + tier_elements: Dict[int, List[str]], + max_tier_gap: int, + cross_tier_prob: float, + used_keys: Set[RecipeKey], + max_attempts: int = 50, + ) -> Optional[RecipeKey]: + """Generate one recipe for a given result element. + + - input1: always from tier T-1 + - input2: from tier T-1 (normal) or lower tier (cross-tier, with probability) + """ + for _ in range(max_attempts): + # input1: must be from tier T-1 + input1 = rng.choice(tier_elements[result_tier - 1]) + + # input2: cross-tier or same tier + if rng.random() < cross_tier_prob and result_tier >= 2: + min_tier = max(0, result_tier - max_tier_gap) + max_input_tier = result_tier - 2 # must be lower than T-1 + if min_tier <= max_input_tier: + input2_tier = rng.randint(min_tier, max_input_tier) + input2 = rng.choice(tier_elements[input2_tier]) + else: + input2 = rng.choice(tier_elements[result_tier - 1]) + else: + input2 = rng.choice(tier_elements[result_tier - 1]) + + key: RecipeKey = tuple(sorted([input1, input2])) + if key not in used_keys: + return key + + return None # Could not find a unique recipe + + def _add_noise( + self, + rng: random.Random, + tier_elements: Dict[int, List[str]], + recipes: Dict[RecipeKey, str], + used_keys: Set[RecipeKey], + all_names: Set[str], + noise_node_ratio: float, + noise_chain_depth: int, + max_tier_gap: int, + cross_tier_prob: float, + ) -> Set[str]: + """Add noise (distractor) elements that are dead ends. + + Returns the set of all noise element names. + """ + num_tiers = len(tier_elements) + # Snapshot original sizes before adding noise + original_sizes = {t: len(elems) for t, elems in tier_elements.items()} + all_noise: Set[str] = set() + + for depth_round in range(noise_chain_depth): + new_noise: Dict[int, List[str]] = {} + + for t in range(1, num_tiers): + # Number of noise nodes based on ORIGINAL tier size, not inflated + num_noise = max(0, round(original_sizes[t] * noise_node_ratio)) + if depth_round > 0: + # Fewer noise nodes in deeper rounds + num_noise = max(0, round(num_noise * 0.5)) + + for _ in range(num_noise): + noise_name = _generate_name(rng, all_names) + all_names.add(noise_name) + + # Generate one recipe for this noise node + key = self._make_recipe( + rng, noise_name, t, tier_elements, max_tier_gap, + cross_tier_prob, used_keys, + ) + if key is not None: + recipes[key] = noise_name + used_keys.add(key) + new_noise.setdefault(t, []).append(noise_name) + all_noise.add(noise_name) + + # Add noise nodes to tier_elements for next round + for t, names in new_noise.items(): + tier_elements[t].extend(names) + + return all_noise + + def generate_instance( + self, + seed: int, + min_num_tiers: int = 3, + max_num_tiers: int = 4, + min_base_elements: int = 10, + max_base_elements: int = 15, + tier_shrink_min: float = 0.4, + tier_shrink_max: float = 0.6, + min_recipes_per_element: int = 2, + max_recipes_per_element: int = 3, + max_tier_gap: int = 2, + min_cross_tier_prob: float = 0.2, + max_cross_tier_prob: float = 0.35, + min_noise_node_ratio: float = 0.15, + max_noise_node_ratio: float = 0.25, + noise_chain_depth: int = 3, + material_mode: str = "unlimited", + min_material_mult: float = 1.5, + max_material_mult: float = 2.0, + ) -> GameInstance: + """Generate a game instance for per-task scope.""" + rng = random.Random(seed) + + num_tiers = rng.randint(min_num_tiers, max_num_tiers) + base_elements = rng.randint(min_base_elements, max_base_elements) + tier_shrink = rng.uniform(tier_shrink_min, tier_shrink_max) + cross_tier_prob = rng.uniform(min_cross_tier_prob, max_cross_tier_prob) + noise_node_ratio = rng.uniform(min_noise_node_ratio, max_noise_node_ratio) + + tier_elements, recipes, local_tiers, noise_elements = self.generate_graph( + rng, num_tiers, base_elements, tier_shrink, + min_recipes_per_element, max_recipes_per_element, + max_tier_gap, cross_tier_prob, noise_node_ratio, noise_chain_depth, + ) + + # Target from highest tier (noise excluded) + valid_targets = [e for e in tier_elements[num_tiers - 1] if e not in noise_elements] + if not valid_targets: + # Fallback: try lower tiers + for t in range(num_tiers - 2, 0, -1): + valid_targets = [e for e in tier_elements[t] if e not in noise_elements] + if valid_targets: + break + target = rng.choice(valid_targets) + + return self._build_instance( + rng, target, tier_elements, recipes, local_tiers, + material_mode, min_material_mult, max_material_mult, + ) + + def generate_instance_for_pack( + self, + pack_seed: int, + task_seed: int, + task_idx: int = 0, + num_tasks: int = 8, + min_num_tiers: int = 3, + max_num_tiers: int = 4, + min_base_elements: int = 10, + max_base_elements: int = 15, + tier_shrink_min: float = 0.4, + tier_shrink_max: float = 0.6, + min_recipes_per_element: int = 2, + max_recipes_per_element: int = 3, + max_tier_gap: int = 2, + min_cross_tier_prob: float = 0.2, + max_cross_tier_prob: float = 0.35, + min_noise_node_ratio: float = 0.15, + max_noise_node_ratio: float = 0.25, + noise_chain_depth: int = 3, + material_mode: str = "unlimited", + min_material_mult: float = 1.5, + max_material_mult: float = 2.0, + ) -> GameInstance: + """Generate a game instance for per-pack scope. + + All tasks in a pack share the same graph. Each task gets a different target + selected via weighted sampling (higher tiers more likely, no duplicates). + """ + pack_rng = random.Random(pack_seed) + + num_tiers = pack_rng.randint(min_num_tiers, max_num_tiers) + base_elements = pack_rng.randint(min_base_elements, max_base_elements) + tier_shrink = pack_rng.uniform(tier_shrink_min, tier_shrink_max) + cross_tier_prob = pack_rng.uniform(min_cross_tier_prob, max_cross_tier_prob) + noise_node_ratio = pack_rng.uniform(min_noise_node_ratio, max_noise_node_ratio) + + tier_elements, recipes, local_tiers, noise_elements = self.generate_graph( + pack_rng, num_tiers, base_elements, tier_shrink, + min_recipes_per_element, max_recipes_per_element, + max_tier_gap, cross_tier_prob, noise_node_ratio, noise_chain_depth, + ) + + # Group valid (non-noise) elements by tier for target selection + tier_groups: Dict[int, List[str]] = {} + for t in range(1, num_tiers): + candidates = [e for e in tier_elements[t] if e not in noise_elements] + if candidates: + tier_groups[t] = candidates + + # Tier-weighted target selection, shuffled across tasks + ordered_targets = self._select_targets(pack_rng, tier_groups, num_tasks) + + target = ordered_targets[task_idx % len(ordered_targets)] + + task_rng = random.Random(task_seed) + return self._build_instance( + task_rng, target, tier_elements, recipes, local_tiers, + material_mode, min_material_mult, max_material_mult, + ) + + def _select_targets( + self, + rng: random.Random, + tier_groups: Dict[int, List[str]], + num_tasks: int, + ) -> List[str]: + """Select targets via tier-weighted sampling without replacement. + + First selects a tier (weight = tier + 1), then picks a random element + from that tier. Skips tiers that are exhausted. + """ + # Deep copy to avoid mutating + available = {t: list(elems) for t, elems in tier_groups.items()} + ordered: List[str] = [] + + for _ in range(num_tasks): + # Filter exhausted tiers + active = {t: elems for t, elems in available.items() if elems} + if not active: + break + + # Weighted tier selection + tiers = list(active.keys()) + weights = [t + 1 for t in tiers] + chosen_tier = rng.choices(tiers, weights=weights, k=1)[0] + + # Pick random element from chosen tier, remove it + elem = rng.choice(active[chosen_tier]) + available[chosen_tier].remove(elem) + ordered.append(elem) + + rng.shuffle(ordered) + return ordered + + def _build_instance( + self, + rng: random.Random, + target: str, + tier_elements: Dict[int, List[str]], + recipes: Dict[RecipeKey, str], + local_tiers: Dict[str, int], + material_mode: str, + min_material_mult: float, + max_material_mult: float, + ) -> GameInstance: + """Build a GameInstance for a given target.""" + all_elements = set() + for elems in tier_elements.values(): + all_elements.update(elems) + + # Trace one solution path + solution_path = self._trace_solution(target, recipes) + + # Starting inventory: all tier 0 elements, same quantity each + if material_mode == "unlimited": + starting_inventory = {e: 1 for e in tier_elements[0]} + else: + # Uniform count based on solution path needs + base_counts = self._compute_base_counts(solution_path) + max_needed = max(base_counts.values()) if base_counts else 1 + material_mult = rng.uniform(min_material_mult, max_material_mult) + uniform_count = max(max_needed, int(max_needed * material_mult + 0.5)) + starting_inventory = {e: uniform_count for e in tier_elements[0]} + + target_tier = local_tiers.get(target, 0) + + return GameInstance( + target=target, + recipes=recipes, + starting_inventory=starting_inventory, + all_elements=all_elements, + solution_path=solution_path, + depth=target_tier, + local_tiers=local_tiers, + ) + + @staticmethod + def _trace_solution( + target: str, + recipes: Dict[RecipeKey, str], + ) -> List[Tuple[str, str, str]]: + """Trace one solution path from base elements to target.""" + reverse: Dict[str, List[RecipeKey]] = defaultdict(list) + for key, result in recipes.items(): + reverse[result].append(key) + + path: List[Tuple[str, str, str]] = [] + visited: Set[str] = set() + + def _expand(elem: str): + if elem in visited: + return + visited.add(elem) + available = reverse.get(elem, []) + if not available: + return + chosen = available[0] + e1, e2 = chosen + _expand(e1) + _expand(e2) + path.append((e1, e2, elem)) + + _expand(target) + return path + + @staticmethod + def _compute_base_counts( + solution_path: List[Tuple[str, str, str]], + ) -> Dict[str, int]: + """Compute minimum base material counts for the solution path.""" + counts: Dict[str, int] = defaultdict(int) + available: Dict[str, int] = defaultdict(int) + + for e1, e2, result in solution_path: + for elem in [e1, e2]: + if available.get(elem, 0) > 0: + available[elem] -= 1 + else: + counts[elem] += 1 + available[result] = available.get(result, 0) + 1 + + return dict(counts) diff --git a/trinity/common/workflows/connect_the_dots/alchemy/workflow_random.py b/trinity/common/workflows/connect_the_dots/alchemy/workflow_random.py new file mode 100644 index 00000000000..c3be26b2cae --- /dev/null +++ b/trinity/common/workflows/connect_the_dots/alchemy/workflow_random.py @@ -0,0 +1,313 @@ +# -*- coding: utf-8 -*- +""" +CoD workflow for Random Alchemy crafting game. + +Uses procedurally generated recipe graphs instead of Little Alchemy data. +Element names are already random, so no obfuscation needed. +""" + +from typing import List, Optional, Tuple + +from trinity.common.experience import Experience +from trinity.common.models.model import ModelWrapper +from trinity.common.workflows.connect_the_dots.alchemy.env import AlchemyEnv +from trinity.common.workflows.connect_the_dots.alchemy.prompts import ( + load_system_prompt, + load_user_prompt, +) +from trinity.common.workflows.connect_the_dots.alchemy.random_graph import ( + RandomGraphGenerator, +) +from trinity.common.workflows.connect_the_dots.base_workflow import ( + AsyncCoDMultiStepWorkflow, +) +from trinity.common.workflows.connect_the_dots.utils import extract_content_between_keys +from trinity.common.workflows.workflow import Task + + +def parse_action(response: str) -> Optional[Tuple[str, str]]: + """Parse action from model response.""" + key_start, key_end = "", "" + content, success = extract_content_between_keys(response, key_start, key_end) + if not success: + return None + action_str = content.strip() + parts = action_str.split("+") + if len(parts) != 2: + return None + elem1 = parts[0].strip() + elem2 = parts[1].strip() + if not elem1 or not elem2: + return None + return (elem1, elem2) + + +class CoDRandomAlchemyWorkflow(AsyncCoDMultiStepWorkflow): + """CoD Random Alchemy workflow using procedurally generated graphs.""" + + is_async: bool = True + + def __init__( + self, + model: ModelWrapper, + task: Task, + auxiliary_models: Optional[List] = None, + use_openai_client: bool = False, + ): + super().__init__( + task=task, + model=model, + auxiliary_models=auxiliary_models, + use_openai_client=use_openai_client, + ) + + workflow_args = task.workflow_args if hasattr(task, "workflow_args") else {} + self.max_rounds = workflow_args.get("max_rounds", 15) + self.max_response_tokens_restraint = workflow_args.get( + "max_response_tokens_restraint", None + ) + self.show_recipes = workflow_args.get("show_recipes", False) + self.show_elements = workflow_args.get("show_elements", False) + self.show_tiers = workflow_args.get("show_tiers", True) + self.material_mode = workflow_args.get("material_mode", "unlimited") + self.scope = workflow_args.get("scope", "per_task") + + # Graph generation params + self.min_num_tiers = workflow_args.get("min_num_tiers", 3) + self.max_num_tiers = workflow_args.get("max_num_tiers", 4) + self.min_base_elements = workflow_args.get("min_base_elements", 10) + self.max_base_elements = workflow_args.get("max_base_elements", 15) + self.tier_shrink_min = workflow_args.get("tier_shrink_min", 0.4) + self.tier_shrink_max = workflow_args.get("tier_shrink_max", 0.6) + self.min_recipes_per_element = workflow_args.get("min_recipes_per_element", 2) + self.max_recipes_per_element = workflow_args.get("max_recipes_per_element", 3) + self.max_tier_gap = workflow_args.get("max_tier_gap", 2) + self.min_cross_tier_prob = workflow_args.get("min_cross_tier_prob", 0.2) + self.max_cross_tier_prob = workflow_args.get("max_cross_tier_prob", 0.35) + self.min_noise_node_ratio = workflow_args.get("min_noise_node_ratio", 0.15) + self.max_noise_node_ratio = workflow_args.get("max_noise_node_ratio", 0.25) + self.noise_chain_depth = workflow_args.get("noise_chain_depth", 3) + self.min_material_mult = workflow_args.get("min_material_mult", 1.5) + self.max_material_mult = workflow_args.get("max_material_mult", 2.0) + + # Raw task params + self.raw_task = task.raw_task if hasattr(task, "raw_task") else {} + self.seed = self.raw_task.get("seed", 42) + + # Generate game instance + gen = RandomGraphGenerator() + gen_kwargs = dict( + min_num_tiers=self.min_num_tiers, + max_num_tiers=self.max_num_tiers, + min_base_elements=self.min_base_elements, + max_base_elements=self.max_base_elements, + tier_shrink_min=self.tier_shrink_min, + tier_shrink_max=self.tier_shrink_max, + min_recipes_per_element=self.min_recipes_per_element, + max_recipes_per_element=self.max_recipes_per_element, + max_tier_gap=self.max_tier_gap, + min_cross_tier_prob=self.min_cross_tier_prob, + max_cross_tier_prob=self.max_cross_tier_prob, + min_noise_node_ratio=self.min_noise_node_ratio, + max_noise_node_ratio=self.max_noise_node_ratio, + noise_chain_depth=self.noise_chain_depth, + material_mode=self.material_mode, + min_material_mult=self.min_material_mult, + max_material_mult=self.max_material_mult, + ) + + if self.scope == "per_pack": + pack_seed = self.raw_task.get("pack_seed", self.seed) + task_idx = self.raw_task.get("task_idx", 0) + num_tasks = self.raw_task.get("pack_size", 8) + self.game_instance = gen.generate_instance_for_pack( + pack_seed=pack_seed, + task_seed=self.seed, + task_idx=task_idx, + num_tasks=num_tasks, + **gen_kwargs, + ) + else: + self.game_instance = gen.generate_instance( + seed=self.seed, + **gen_kwargs, + ) + + # Create environment + self.env = AlchemyEnv(self.game_instance, max_rounds=self.max_rounds, + material_mode=self.material_mode) + + # State + self.done: bool = False + self.final_reward: float = 0.0 + self.current_step: int = 0 + self.action_feedback: Optional[str] = None + self.early_termination_by_format_issue: bool = False + + def _build_system_prompt(self) -> str: + target = self.game_instance.target + + kwargs = { + "show_recipes": self.show_recipes, + "show_elements": self.show_elements, + "show_tiers": self.show_tiers, + "obfuscate_names": True, # names are random, same prompt style + "material_mode": self.material_mode, + "target": target, + } + if self.show_recipes: + kwargs["recipes"] = self.game_instance.recipes + if self.show_elements: + kwargs["all_elements"] = sorted(self.game_instance.all_elements) + if self.show_tiers: + kwargs["tier_info"] = self._build_tier_info() + + sys_prompt = load_system_prompt(**kwargs) + + sys_prompt = self._augment_system_prompt(sys_prompt) + return sys_prompt + + def _build_tier_info(self) -> str: + local_tiers = self.game_instance.local_tiers or {} + tier_groups: dict = {} + for elem, tier in sorted(local_tiers.items(), key=lambda x: x[1]): + tier_groups.setdefault(tier, []).append(elem) + + lines = [] + for tier in sorted(tier_groups.keys()): + elems = tier_groups[tier] + # label = f"tier {tier}" if tier == 0 else f"tier {tier} (needs {tier} synthesis step{'s' if tier > 1 else ''})" + label = f"tier {tier}" + lines.append(f"- {label}: {', '.join(elems)}") + target = self.game_instance.target + target_tier = local_tiers.get(target, "?") + lines.append(f"\nThe target element '{target}' is at tier {target_tier}.") + return "\n".join(lines) + + def _build_action_feedback( + self, + action: Tuple[str, str], + result: Optional[str], + invalid: bool, + ) -> str: + e1, e2 = action + if invalid: + inv = self.env.inventory + if e1 not in inv and e2 not in inv: + return f"'{e1}' and '{e2}' are not in your inventory." + elif e1 not in inv: + return f"'{e1}' is not in your inventory." + elif e2 not in inv: + return f"'{e2}' is not in your inventory." + elif e1 == e2: + return f"Not enough '{e1}' (need 2, have {inv.get(e1, 0)})." + return f"Not enough materials for '{e1}' + '{e2}'." + if result is not None: + return f"You combined {e1} + {e2} and created: {result}!" + if self.material_mode == "unlimited": + return f"{e1} + {e2} — nothing happened." + return f"{e1} + {e2} — nothing happened. Materials consumed." + + async def run_async(self) -> List[Experience]: + obs, info = self.env.reset() + self.done = False + self.final_reward = 0.0 + self.current_step = 0 + self.action_feedback = None + self.early_termination_by_format_issue = False + + self.memory.clear() + self.memory.append({"role": "system", "content": self._build_system_prompt()}) + + return await super().run_async() + + async def step_async(self, step_num: int) -> Tuple[bool, List[Experience]]: + if self.done: + return False, [] + + target = self.game_instance.target + inv = self.env.inventory + + user_content = load_user_prompt( + current_round=step_num + 1, + max_rounds=self.max_rounds, + inventory=inv, + target=target, + action_feedback=self.action_feedback, + material_mode=self.material_mode, + ) + + if self.icl_examples and step_num == 0: + user_content = f"{user_content}\n\nHere are some reference examples:\n\n{self.icl_examples}" + + self.memory.append({"role": "user", "content": user_content}) + + if self.reply_prefix: + self.memory.append({"role": "assistant", "content": self.reply_prefix}) + + experiences = await self.model.chat_async(self.memory) + response_text = experiences[0].response_text + self.memory.append({"role": "assistant", "content": response_text}) + + sys_prompt = self.memory[0]["content"] if self.memory and self.memory[0]["role"] == "system" else "" + for exp in experiences: + exp.info["sys_prompt"] = sys_prompt + exp.info["user_prompt"] = user_content + + action = parse_action(response_text) + + if action is None: + self.action_feedback = ( + "Invalid format: could not parse combination. " + "Expected format: {your reasoning process here}Element1 + Element2. Game over." + ) + self.early_termination_by_format_issue = True + self.done = True + self.current_step = step_num + 1 + return False, experiences + + e1, e2 = action + + # Case-insensitive match to inventory + inv_lower = {x.lower(): x for x in self.env.inventory} + e1_canonical = inv_lower.get(e1.lower(), e1) + e2_canonical = inv_lower.get(e2.lower(), e2) + + obs, reward, terminated, truncated, info = self.env.step( + (e1_canonical, e2_canonical) + ) + + self.action_feedback = self._build_action_feedback( + (e1_canonical, e2_canonical), + info.get("last_result"), + invalid=info.get("invalid_action", False), + ) + + self.done = terminated or truncated + self.current_step = step_num + 1 + + if terminated and reward > 0: + self.final_reward = reward + + return not self.done, experiences + + def _get_feedback(self) -> str: + target = self.game_instance.target + + if target in self.env.inventory: + return f"Success! Synthesized {target}." + if self.early_termination_by_format_issue: + return self.action_feedback + if len(self.env.get_feasible_actions()) == 0: + return ( + f"Failed: Ran out of materials. Could not synthesize {target}. " + f"Used {self.current_step} rounds." + ) + return ( + f"Failed: Ran out of rounds ({self.current_step}/{self.max_rounds}). " + f"Could not synthesize {target}." + ) + + @property + def max_step_num(self) -> int: + return self.max_rounds diff --git a/trinity/common/workflows/connect_the_dots/base_workflow.py b/trinity/common/workflows/connect_the_dots/base_workflow.py new file mode 100644 index 00000000000..6efefb7e7d1 --- /dev/null +++ b/trinity/common/workflows/connect_the_dots/base_workflow.py @@ -0,0 +1,304 @@ +# -*- coding: utf-8 -*- +""" +Base workflow class for CoD (Connect-the-Dots) multi-step workflows. + +This workflow is designed for use with CoDWorkflow, which handles metrics +post-processing via _post_process_task_solving_exp. Therefore, this base class +does NOT broadcast metrics to every experience (avoiding metric explosion like +mean@101, mean@102, etc. in eval logs). +""" + +import re +from dataclasses import asdict +from typing import Any, List, Optional, Tuple + +import openai + +from trinity.common.experience import Experience +from trinity.common.models.model import ModelWrapper +from trinity.common.workflows.connect_the_dots.agentscope_utils import ( + agentscope_msgs_to_text, +) +from trinity.common.workflows.connect_the_dots.utils import extract_content_between_keys +from trinity.common.workflows.workflow import Task, Workflow + + +class AsyncCoDMultiStepWorkflow(Workflow): + """ + Async base class for CoD multi-step workflows. + + This is the recommended base class for CoD multi-step workflows like + CoDFrozenLakeWorkflow, CoDAlfworldWorkflow, CoDTextWorldWorkflow, etc. + + Unlike RewardPropagationWorkflow, this class does NOT set metrics on every + experience. Metrics should be handled by CoDWorkflow._post_process_task_solving_exp. + + Subclasses should: + - Initialize self.memory as a list of message dicts in __init__ or reset + - The memory should follow the format: [system, user, assistant, user, assistant, ...] + """ + + is_async: bool = True + + def __init__( + self, *, task: Task, model: ModelWrapper, auxiliary_models=None, use_openai_client=True + ): + super().__init__(task=task, model=model, auxiliary_models=auxiliary_models) + self.client: Optional[openai.OpenAI] = None + if use_openai_client: + self.client = model.get_openai_client() + + # Memory for conversation history (subclasses should populate this) + self.memory: List[dict] = [] + + # State variables (TODO: merge with those in children classes) + self.final_reward: float = 0.0 # outcome reward + self.early_termination_by_format_issue: bool = False # terminate upon format error + self.reply_prefix = task.format_args.reply_prefix if task.format_args else None + self.hint = None + self.icl_examples = None + self.max_response_tokens_restraint = None + + def reset(self, task: Task): + """Set task-derived attrs. Subclasses override and call super().""" + self.task = task + self.format_args = task.format_args + self.raw_task = task.raw_task or {} + self.task_desc = task.task_desc + self.reply_prefix = task.format_args.reply_prefix if task.format_args else None + self.hint = None + self.icl_examples = None + self.max_response_tokens_restraint = None + + @property + def rollout_args(self): + return asdict(self.task.rollout_args) + + def _compress_assistant_response(self, response: str) -> str: + """Various methods of compressing assistant response to reduce context size.""" + context_compression_mode = self.task.workflow_args.get("context_compression_mode", "keep_all") + if context_compression_mode == "remove_think": + # TODO: need update, maybe return None when no match of tag + compressed = re.sub(r".*?", "", response, flags=re.DOTALL) + return compressed.strip() + elif context_compression_mode == "keep_answer": + content, ok = extract_content_between_keys(response, "", "") + if ok: + # re-wrap with tags to allow repeated compression + return f"{content.strip()}" + return "(Failed to parse answer from response)" + elif context_compression_mode == "keep_all": + return response + else: + raise ValueError(f"Invalid context_compression_mode {context_compression_mode}") + + def _compress_memory(self) -> None: + """Compress all assistant messages in memory by stripping blocks. + + Called after each step_async to prevent context bloat in multi-turn + conversations. The thinking/reasoning content is removed while + preserving answers and other non-thinking output. + + !!! TODO: change to compressing the last assistant message? !!! + """ + for msg in self.memory: + if msg.get("role") == "assistant": + msg["content"] = self._compress_assistant_response(msg["content"]) + + async def run_async(self) -> list[Experience]: + """Run the workflow asynchronously and return a list of experiences.""" + experiences = [] + step = 0 + for step in range(self.max_step_num): + # Run a single step of the agent application and get experiences directly + continue_run, exps = await self.step_async(step_num=step) + # Compress prior assistant messages in memory to avoid context bloat + self._compress_memory() + # Set the step number in each experience + for exp in exps: + exp.eid.step = step + # Store the step experiences + experiences.extend(exps) + if not continue_run: + break + + # Calculate final reward and propagate to all experiences + reward = await self.reward_async(experiences) + for exp in experiences: + exp.reward = reward + if exp.metrics is None: + exp.metrics = {} + + # Only set actual_env_steps and trajectory on the LAST experience + # (CoDWorkflow._post_process_task_solving_exp will handle per-experience metrics) + if experiences: + experiences[-1].metrics["actual_env_steps"] = step + 1 + experiences[-1].info["trajectory"] = self._build_trajectory() + + return experiences + + def _augment_system_prompt(self, sys_prompt: str) -> str: + """Augment system prompt, injecting hint and max_response_tokens_restraint.""" + if self.hint: + sys_prompt = sys_prompt + ( + "\n\n## Hints that might help\n" + "Below are some hints that might be helpful, but there is no guarantee " + "that they must be correct or applicable to the current task. You might " + "leverage them as prior knowledge, while incorporating new information " + "from your own experience in interacting with the environment.\n\n" + f"{self.hint}" + ) + if self.max_response_tokens_restraint: + sys_prompt = sys_prompt + ( + f"\n\n## Response length limit\n" + f"Please limit your response (including your thinking process) " + f"to {self.max_response_tokens_restraint} tokens." + ) + return sys_prompt + + @staticmethod + def _strip_system_prompt(sys_content: str) -> str: + """Strip hint and token limit from system prompt, keeping only task rules. + Markers correspond to those in the _augment_system_prompt method above. + """ + for marker in [ + "\n\n## Hints", + "\n\n## Response length limit", + ]: + sys_content = sys_content.split(marker)[0] + return sys_content + + @staticmethod + def _strip_icl_examples(user_content: str) -> str: + marker = "\n\nHere are some reference examples:\n\n" + if marker in user_content: + return user_content.split(marker)[0] + return user_content + + def _build_trajectory(self) -> str: + """Build trajectory string from memory for hint generation. + + Includes the system prompt (with hint/token limit stripped) so that + hint generation can see the task rules, followed by observation/action + pairs from the conversation history. + + Subclasses can override this method for custom trajectory formatting. + """ + if not self.memory: + return "" + + trajectory_parts = [] + step_num = 0 + i = 0 + while i < len(self.memory): + msg = self.memory[i] + # Include system prompt with hint/token limit stripped + if msg.get("role") == "system": + clean_sys = self._strip_system_prompt(msg.get("content", "")) + if clean_sys.strip(): + trajectory_parts.append(f"System:\n{clean_sys}") + i += 1 + continue + # Process user message + if msg.get("role") == "user": + step_num += 1 + observation = msg.get("content", "") + observation = self._strip_icl_examples(observation) + # Find the last assistant message before next user (to skip reply_prefix) + action = "(no action)" + j = i + 1 + while j < len(self.memory) and self.memory[j].get("role") == "assistant": + action = self.memory[j].get("content", "") + j += 1 + trajectory_parts.append(f"Step {step_num}:\nObservation: {observation}\nAction: {action}") + i = j + else: + i += 1 + + return "\n\n".join(trajectory_parts) + + def _build_agentscope_trajectory(self, sys_prompt: str, messages: List[Any]) -> str: + """Build CoD trajectory text from AgentScope memory messages.""" + trajectory_parts = [] + sys_prompt = self._strip_system_prompt(sys_prompt) + if sys_prompt.strip(): + trajectory_parts.append(f"System:\n{sys_prompt}") + + steps: List[Tuple[str, List[Any]]] = [] + + for msg in messages: + if msg.role == "user": + steps.append((msg.get_text_content() or "", [])) + elif steps: + steps[-1][1].append(msg) + + for step_num, (user_content, step_messages) in enumerate(steps, start=1): + observation = self._strip_icl_examples(user_content) + # !!! AgentScope stores tool calls/results as structured blocks. + action = agentscope_msgs_to_text(step_messages) + if not action: + action = "(no action)" + trajectory_parts.append( + f"Step {step_num}:\nObservation: {observation}\nAction: {action}" + ) + return "\n\n".join(trajectory_parts) + + async def step_async(self, step_num: int) -> Tuple[bool, List[Experience]]: + """Run a single step of your agent application asynchronously. + + Args: + step_num (int): The current step number. + + Returns: + Tuple[bool, List[Experience]]: A tuple of (continue_run, experiences). + - continue_run: Whether to continue running the agent application. + - experiences: List of experiences from this step. + """ + raise NotImplementedError + + def _get_feedback(self) -> str: + """Get environment feedback that will be added to exps[-1].info["feedback"]""" + raise NotImplementedError + + async def reward_async(self, exps: List[Experience]) -> float: + """Default reward function for CoD multi-step workflows.""" + if exps: + exps[-1].info["feedback"] = self._get_feedback() + if exps[-1].metrics is None: + exps[-1].metrics = {} + exps[-1].metrics["format_error_termination"] = 1.0 if self.early_termination_by_format_issue else 0.0 + exps[-1].info["early_termination_by_format_issue"] = self.early_termination_by_format_issue + + # Original outcome reward + returned_reward = self.final_reward + + # Length penalties + length_penalty_coef = self.task.workflow_args.get("length_penalty_coef", 0.0) + if exps and (length_penalty_coef > 0.0) and (self.final_reward > 0.0): + len_full_penalty = self.task.workflow_args.get("len_full_penalty", -1) + len_zero_penalty = self.task.workflow_args.get("len_zero_penalty", -1) + assert len_full_penalty > 0, "len_full_penalty should be set to a positive integer." + assert len_zero_penalty > 0, "len_zero_penalty should be set to a positive integer." + assert len_full_penalty > len_zero_penalty, "len_full_penalty should be larger than len_zero_penalty." + + # response-wise penalty + resp_len_penalties = [] + for exp in exps: + resp_len = len(exp.tokens) - exp.prompt_length + penalty = min(1.0, max(0.0, (resp_len - len_zero_penalty) / (len_full_penalty - len_zero_penalty))) + resp_len_penalties.append(penalty) + response_wise_penalty = sum(resp_len_penalties) / len(resp_len_penalties) # range [0, 1] + + # episode-wise penalty + episode_wise_penalty = len(exps) / self.max_step_num # range [0, 1] + + # combine and substract from original reward + total_penalty = length_penalty_coef * (response_wise_penalty + episode_wise_penalty) + returned_reward = max(0.0, returned_reward - total_penalty) + + return returned_reward + + @property + def max_step_num(self) -> int: + """Return the maximum number of steps in the task.""" + raise NotImplementedError diff --git a/trinity/common/workflows/connect_the_dots/cod_utils.py b/trinity/common/workflows/connect_the_dots/cod_utils.py new file mode 100644 index 00000000000..1225039fd58 --- /dev/null +++ b/trinity/common/workflows/connect_the_dots/cod_utils.py @@ -0,0 +1,190 @@ +"""Utils for CoD.""" + +from collections import defaultdict +from typing import Dict, List, Tuple + +from trinity.common.workflows.connect_the_dots.utils import extract_content_between_keys + + +def compute_part_summary(exps: List) -> dict: + """Compute summary statistics (accuracy per interaction step) for a list of experiences.""" + step_rewards: Dict[int, List[float]] = defaultdict(list) + all_rewards = [] + for exp in exps: + reward = exp.reward + all_rewards.append(reward) + step_id = exp.eid.step + if step_id is not None: + step_rewards[step_id].append(reward) + + num_samples = len(all_rewards) + overall_acc = sum(all_rewards) / num_samples if num_samples > 0 else 0.0 + + per_step_acc = {} + for sid in sorted(step_rewards.keys()): + rs = step_rewards[sid] + per_step_acc[str(sid)] = round(sum(rs) / len(rs), 4) + + return { + "overall_acc": round(overall_acc, 4), + "num_samples": num_samples, + "per_interaction_step_acc": per_step_acc, + } + + +def clone_model_with_isolated_history(model): + """Clone a model wrapper with isolated, enabled request history.""" + isolated_model = model.clone_with_isolated_history() + isolated_model.enable_history = True + return isolated_model + + +HINTS_EXAMPLE_PROMPT = """As an example, the end of a complete response might be: +I would retain the planning and verification lessons because they transfer to +future tasks, while omitting details specific to the completed task. +--- Start of updated hints --- +- Start with a list of known conditions. +- Create a high-level plan and verify each important step. +- Return the final answer only when sufficiently confident. +--- End of updated hints --- +""" + +HINT_START_MARKER = "--- Start of updated hints ---" +HINT_END_MARKER = "--- End of updated hints ---" + + +class CoDPrompts: + @staticmethod + def sys_prompt_gen_hint_iteratively(hint_example: bool = False) -> str: + """System prompt for iterative hint generation.""" + prompt = """You are a helpful assistant and a cross-task learning agent. Your responsibility is to solve **a sequence of different-but-related tasks** within the same environment, while **transferring informative hints across tasks** that can help achieve better task-solving performance. + +## Problem setting + +**In the current session, your job is to update the hints.** + +To be concrete, the user provides: +1. Previous hints that were used to guide task solving; +2. A new task-solving trajectory. + +Based on this information, you need to generate updated hints that can assist in solving future tasks within the same environment more effectively, by +- Preserving useful guidance and information from the previous hints, while discarding incorrect ones; +- Incorporating new lessons and information that can be learned from the new trajectory. + +## General principles + +- You are in an environment that might be stationary or non-stationary, and relation between tasks in this environment is initially unknown. One general principle is thus to generate **informative** hints that can serve as useful priors when solving a new task. If you are uncertain about the correctness or usefulness of certain hints, feel free to include them in your response and briefly mention uncertainty. +- There are at least two major categories of useful hints: (1) revealed information / clues about the environment that were initially unknown; (2) task-solving techniques with validated efficacy. With that said, your updated hints can certainly go beyond these two categories. + +## Response format + +First provide a brief thinking process that evaluates the previous hints and the +new trajectory. After that reasoning, finish your response with exactly this +hint block: + +--- Start of updated hints --- +- Your concise, transferable hints go here. +--- End of updated hints --- + +The two delimiter lines are literal protocol tokens. Copy both exactly, +including capitalization, spaces, and hyphens. + +Other requirements: +- Keep your thinking process concise. Avoid overthinking. +- Put the thinking process before the start delimiter, never inside the hint block. +- Do not place any text after the end delimiter. +- Keep your updated hints concise if possible, containing the most critical and helpful information. Avoid generic and uninformative statements. Avoid repeating the system prompt within the task-solving trajectory. +- Make sure that your updated hints can be potentially helpful for other related tasks, rather than specific to the task in the provided trajectory. +- Use standard Markdown format for your updated hints. Simpler structures (e.g., bullet points) are preferred. If you do need to use sections, you should start from the third level "###", while avoiding "#" and "##". +""" + if hint_example: + prompt = prompt + HINTS_EXAMPLE_PROMPT + return prompt + + @staticmethod + def user_prompt_gen_hint_iteratively( + prev_hint: str, + trajectory: str, + reward: float, + feedback: str, + ) -> str: + """User prompt for iterative hint generation. + + Args: + prev_hint: Previous hints (empty string if first iteration) + trajectory: The solution trajectory (includes task description and model responses) + reward: The reward received (e.g., 1.0 for success, 0.0 for failure) + feedback: Environment feedback (e.g., final state, error message) + """ + if prev_hint: + prev_hint_section = f"""--- Previous hints --- + +{prev_hint} + +""" + else: + prev_hint_section = """--- Previous hints --- + +(No previous hint available) + +""" + + prompt = f"""{prev_hint_section} +--- A new task-solving trajectory --- + +{trajectory} + +Reward: {reward:.3f} + +Environment feedback: {feedback} + +--- Your job --- + +Now, please think through it and return your updated hints based on the above information, following the requirements in the system prompt. +""" + # !!! For "cheating" in environments like frozenlake-obscure, add: + # "Note that there were initially some hidden clues about the environment. If you can reveal some of them based on the provided trajectory, include them in your updated hints." + + return prompt + + @staticmethod + def extract_hint(response: str) -> Tuple[str, bool]: + """Extract hints from response. + + Returns: extracted hint (str) and indicator of successful parsing (bool) + """ + content, success = extract_content_between_keys( + response, HINT_START_MARKER, HINT_END_MARKER + ) + if success and content.strip(): + return "Hints: " + content.strip(), True + return "Hints: no hint available.", False + + +# Inline warning prepended to per-trajectory ground-truth reference when +# the workflow surfaces it via the feedback string. Kept brief on purpose; +# do not enumerate transferability hints — let the system prompt drive that. +GT_LEAKAGE_WARNING = ( + "The ground-truth reference below is case-specific. " + "Updated hints will be applied to other cases, so abstract " + "transferable insights and do not copy case-specific facts." +) + + +def format_gt_block_for_feedback(gt_reference: str) -> str: + """Format a case-specific GT reference as a feedback suffix block. + + Workflows that want to surface a per-trajectory oracle signal to + the iterative hint generator should append this block to their + feedback string (via exp.info["feedback"]). + + Returns an empty string if `gt_reference` is empty, so callers can + unconditionally concatenate without dragging in a warning that + points at nothing. + """ + if not gt_reference: + return "" + return ( + f"\n\n{GT_LEAKAGE_WARNING}\n\n" + f"--- Ground truth reference (case-specific) ---\n{gt_reference}" + ) diff --git a/trinity/common/workflows/connect_the_dots/cod_workflow.py b/trinity/common/workflows/connect_the_dots/cod_workflow.py new file mode 100644 index 00000000000..ae2fa1fe802 --- /dev/null +++ b/trinity/common/workflows/connect_the_dots/cod_workflow.py @@ -0,0 +1,909 @@ +"""Connect-the-Dots (CoD) workflow for cross-task learning and generalization.""" + +import asyncio +import hashlib +import json +import os +from copy import deepcopy +from typing import Dict, List, Optional, Tuple, Union + +from trinity.common.experience import Experience +from trinity.common.models.model import ModelWrapper +from trinity.common.workflows.connect_the_dots.cod_utils import ( + clone_model_with_isolated_history, + compute_part_summary, +) +from trinity.common.workflows.connect_the_dots.update_context_workflow import ( + AsyncCoDUpdateContextWorkflow, + AsyncCoDUpdateContextAgentWorkflow, +) +from trinity.common.workflows.connect_the_dots.utils import compute_stable_pack_seed +from trinity.common.workflows.workflow import Task, Workflow + + +def compute_task_hash(task_desc: str) -> str: + """Compute hash of task description.""" + return hashlib.md5(task_desc.encode()).hexdigest()[:16] + + +def get_taskset_id(task: Task) -> int: + """Get taskset_id of a task. + + Ref: class TasksetScheduler (renamed to DatasetScheduler later), + task.index["taskset_id"] = taskset_id (type: int) + This is applicable only when there are > 1 tasksets. + """ + # return task.index["taskset_id"] + return task.index.get("taskset_id", 0) + + +def pack_tasks( + tasks: List[Task], + pack_size: int, + cod_workflow_args: dict, + pack_strategy: str, +) -> List[Task]: + """Group tasks by taskset_id, pack for each group, and merge outputs. + + For ``cod`` and ``ralph_sample``, drop a group if it has fewer than + ``pack_size`` tasks, or pad it cyclically when ``pad_tasks_to_full_pack`` + is enabled. ``ralph_full`` expands every task into a full pack and does + not drop or pad groups. + """ + if pack_size <= 0: + raise ValueError(f"pack_size should be positive, got {pack_size}.") + if pack_strategy not in ["cod", "ralph_sample", "ralph_full"]: + raise ValueError(f"Unsupported pack_strategy: {pack_strategy}") + + returned_tasks: List[Task] = [] + grouped_tasks = dict() # taskset_id -> list of original tasks + + # Group tasks by taskset_id + for task in tasks: + taskset_id = get_taskset_id(task) + if taskset_id not in grouped_tasks: + grouped_tasks[taskset_id] = [] + grouped_tasks[taskset_id].append(task) + + # Logs for debug + log_count_original = dict() + log_count_final = dict() + log_count_pack = dict() + + # Pack tasks for each group, and merge outputs + for id, task_lst in grouped_tasks.items(): + log_count_original[id] = len(task_lst) + + # ralph_full expands every real task into a pack, without drop-tail or padding. + if pack_strategy == "ralph_full": + log_count_final[id] = len(task_lst) + else: + if len(task_lst) < pack_size: + log_count_final[id] = 0 + log_count_pack[id] = 0 + continue # too few tasks, drop this group + + if cod_workflow_args.get("pad_tasks_to_full_pack", False): + remain_count = len(task_lst) % pack_size + if remain_count > 0: + pad_count = pack_size - remain_count + task_lst = task_lst + task_lst[:pad_count] # pad in a cyclic manner + log_count_final[id] = len(task_lst) + + packed_tasks = pack_tasks_from_same_taskset( + tasks=task_lst, + pack_size=pack_size, + cod_workflow_args=cod_workflow_args, + pack_strategy=pack_strategy, + ) + returned_tasks.extend(packed_tasks) + log_count_pack[id] = len(packed_tasks) + + for id in log_count_original.keys(): + print( + f"original task count {log_count_original[id]},", + f"final task count {log_count_final[id]},", + f"final pack count {log_count_pack[id]}.", + ) + + return returned_tasks + + +def pack_tasks_from_same_taskset( + tasks: List[Task], + pack_size: int, + cod_workflow_args: dict, + pack_strategy: str, +) -> List[Task]: + packed_tasks = [] + + num_packs = len(tasks) if pack_strategy == "ralph_full" else len(tasks) // pack_size + for i in range(num_packs): + if pack_strategy == "cod": + start, end = i * pack_size, (i + 1) * pack_size + sublst = tasks[start:end] + else: + # Ralph strategies repeat one anchor: every task for full, or the + # first task in each regular pack for sample. + anchor_idx = i if pack_strategy == "ralph_full" else i * pack_size + env_task_idx = i % pack_size if pack_strategy == "ralph_full" else 0 + anchor_task = tasks[anchor_idx] + sublst = [] + for _ in range(pack_size): + task = deepcopy(anchor_task) + task.index["env_task_idx"] = env_task_idx + sublst.append(task) + + new_task = deepcopy(sublst[0]) + new_task.workflow = CoDWorkflow + new_task.workflow_args = deepcopy(cod_workflow_args) + + if new_task.raw_task is None: + new_task.raw_task = dict() + new_task.raw_task["aux_tasks"] = sublst[1:] + new_task.raw_task["all_tasks"] = sublst + + packed_tasks.append(new_task) + + return packed_tasks + + +class CoDWorkflow(Workflow): + """A CoD "meta workflow" that can call an existing workflow's run method. + + Assumptions for current version: + - input task is a pack of multiple tasks in task.raw_task["all_tasks"] + - task-solving workflow is async and provides trajectory information + - reusable task-solving workflows set can_reset = True + """ + + can_reset: bool = True + can_repeat: bool = True + is_async: bool = True + + def __init__( + self, + *, + task: Task, + model: ModelWrapper, + auxiliary_models: Optional[List[ModelWrapper]] = None, + ): + super().__init__( + task=task, + model=model, + auxiliary_models=auxiliary_models, + ) + self.all_tasks: List[Task] = [] + self.all_workflow_instances: List[ + Union[Workflow, None] + ] = [] # one workflow instance for each task + self.update_context_workflow: Optional[Workflow] = None + + self.reset(task) + + def _model_for_workflow(self, workflow_cls): + """Return the model wrapper required by a child workflow.""" + # !!! CoD AgentScope workflows need isolated history to recover Experiences. + if getattr(workflow_cls, "requires_isolated_model_history", False): + return clone_model_with_isolated_history(self.model) + return self.model + + def reset(self, cod_task: Task): + """Reset current CoD workflow, as well as all task-solving workflow instances. + + (Doesn't matter much if we always use max_repeat_times_per_runner = 1) + """ + self.task = cod_task + + self.all_tasks = [] + if isinstance(cod_task.raw_task, dict) and "all_tasks" in cod_task.raw_task.keys(): + self.all_tasks = cod_task.raw_task["all_tasks"] # List[Task] + + # Optionally inject pack-level seed into sub-tasks (for per_pack mapping mode in FrozenLakeObscure, etc.) + # Enable via workflow_args: inject_pack_seed: true + if cod_task.workflow_args.get("inject_pack_seed", False): + if cod_task.workflow_args.get("stable_pack_seed", False): + # Stable dataset identities keep eval environments fixed across checkpoints. + pack_seed = compute_stable_pack_seed( + (get_taskset_id(task), task.index["index"]) + for task in self.all_tasks + ) + else: + pack_seed = hash((cod_task.batch_id, cod_task.task_id)) % (2**32) + for i, task in enumerate(self.all_tasks): + if task.raw_task is None: + task.raw_task = {} + task.raw_task["pack_seed"] = pack_seed + # Ralph packs repeat one real task; keep env task_idx fixed for per-pack mappings. + # especially for alchemy_random + task.raw_task["task_idx"] = task.index.get("env_task_idx", i) + task.raw_task["pack_size"] = len(self.all_tasks) + + # Init update-context workflow + taskset_workflow_args = ( + self.all_tasks[0].workflow_args if self.all_tasks else {} + ) + update_context_impl = taskset_workflow_args.get( + "update_context_impl", + cod_task.workflow_args.get("update_context_impl", "update_context"), + ) + if update_context_impl == "pde_update_context": + from .pde_discovery.update_context_workflow import ( + PDEUpdateContextWorkflow, + ) + + update_context_cls = PDEUpdateContextWorkflow + else: + update_context_cls = { + "update_context": AsyncCoDUpdateContextWorkflow, + "update_context_agent": AsyncCoDUpdateContextAgentWorkflow, + }[update_context_impl] + if ( + self.update_context_workflow is None + or self.update_context_workflow.__class__ != update_context_cls + ): + self.update_context_workflow = update_context_cls( + task=cod_task, + model=self._model_for_workflow(update_context_cls), + auxiliary_models=self.auxiliary_model_wrappers, + ) + else: + self.update_context_workflow.reset(cod_task) + + reuse_workflow_instance = cod_task.workflow_args["reuse_workflow_instance"] + if reuse_workflow_instance is False: + self.all_workflow_instances = [None for _ in range(len(self.all_tasks))] + return + + if len(self.all_workflow_instances) not in (0, len(self.all_tasks)): + self.all_workflow_instances = [] + + if len(self.all_workflow_instances) == 0: + for task in self.all_tasks: + workflow_instance = task.to_workflow( + self._model_for_workflow(task.workflow), + self.auxiliary_model_wrappers, + ) + self.all_workflow_instances.append(workflow_instance) + else: + for i, task in enumerate(self.all_tasks): + workflow_instance = self.all_workflow_instances[i] + # 类型一致且可以 reset 才复用 + if workflow_instance.__class__ == task.workflow and workflow_instance.can_reset: + workflow_instance.reset(task) + else: + self.all_workflow_instances[i] = task.to_workflow( + self._model_for_workflow(task.workflow), + self.auxiliary_model_wrappers, + ) + + def set_repeat_times(self, repeat_times, run_id_base): + assert ( + repeat_times == 1 + ), "Current CoD implementation requires repeat_times to be 1 here." + + self.repeat_times = repeat_times + self.task.rollout_args.n = repeat_times + self.run_id_base = run_id_base + + # Each solve-task workflow runs once per CoD trajectory. + for task in self.all_tasks: + task.rollout_args.n = 1 + for workflow_instance in self.all_workflow_instances: + if workflow_instance is not None and workflow_instance.can_repeat: + workflow_instance.set_repeat_times(1, run_id_base) + + # --- CoD utils --- + + @staticmethod + def _get_trajectory(task_exps: List[Experience]) -> str: + """Get trajectory string from experiences. + + Requires workflow to provide trajectory in exp.info["trajectory"]. + Raises KeyError if trajectory is not provided. + """ + if not task_exps: + raise ValueError("task_exps is empty, cannot get trajectory") + + last_exp = task_exps[-1] + if "trajectory" not in last_exp.info: + raise KeyError( + "Workflow must provide trajectory in exp.info['trajectory']. " + "For multi-step workflows, inherit from AsyncCoDMultiStepWorkflow. " + "For single-turn workflows, set trajectory in run_async() with user + assistant format." + ) + return last_exp.info["trajectory"] + + @staticmethod + def _get_feedback(task_exps: List[Experience], task_reward: Optional[float]) -> str: + """Get feedback string from experiences.""" + last_exp = task_exps[-1] if task_exps else None + if last_exp and "feedback" in last_exp.info: + feedback = last_exp.info["feedback"] + else: + # Fallback for workflows without feedback + assert task_reward is not None + feedback = f"Reward: {task_reward} (max: 1.0)" + return feedback + + def _post_process_hint_generation_exp( + self, + exp: Experience, + subtid: int, + part_id: int, + part_name: str, + repeat_id: int, + step_id: int, + ) -> None: + if exp.metrics is None: + exp.metrics = {} + hint = exp.info["hint"] + key_hint_len_with_suffix = f"hint_len_in_char_{part_name}" + exp.metrics.update( + { + "reward_gen_hint": exp.reward, + key_hint_len_with_suffix: len(hint), + } + ) + + exp.eid.task = "_".join( + [ + "batch", + str(self.task.batch_id), + "main", + str(self.task.task_id), + "sub", + str(subtid), + ] + ) + exp.eid.run = ( + self.run_id_base * 100 + self.repeat_times * part_id + repeat_id + ) # TODO: eliminate magic number 100 + exp.eid.step = step_id + + def _post_process_task_solving_exp( + self, + exp: Experience, + subtid: int, + part_id: int, + part_name: str, + repeat_id: int, + append_part_name_to_task_id: bool = False, + ) -> None: + # Only update metrics if exp.metrics is non-empty (i.e., the last exp in a trajectory) + # This avoids metric explosion (mean@101, mean@102, etc.) when using multi-step workflows + # !!! TODO: this prevents setting metrics in workflow, better change this if-condition to something more robust + # !!! TODO: another issue here is that response_len_in_char is only record for last exp of an episode. + if exp.metrics: + key_reward_with_suffix = f"reward_{part_name}" + key_response_len_with_suffix = f"response_len_in_char_{part_name}" + exp.metrics.update( + { + "reward": exp.reward, + key_reward_with_suffix: exp.reward, + key_response_len_with_suffix: len(exp.response_text), + } + ) + + # task / run / step id + task_id_pieces = [ + "batch", + str(self.task.batch_id), + "main", + str(self.task.task_id), + "sub", + str(subtid), + ] + if append_part_name_to_task_id: + # useful for contolling which exps to group in grpo-like advantage calculation + task_id_pieces.extend(["part", part_name]) + exp.eid.task = "_".join(task_id_pieces) + exp.eid.run = ( + self.run_id_base * 100 + self.repeat_times * part_id + repeat_id + ) # TODO: eliminate magic number 100 + # step_id is set by workflow internally (e.g., AsyncCoDMultiStepWorkflow.run_async) + + def _create_sublogs( + self, + task: Task, + exps: List[Experience], + ) -> Dict[str, Union[str, None, dict, list]]: + sublogs: Dict[str, Union[str, None, dict, list]] = {} + # Only include task_desc / truth when they are not None + if task.task_desc is not None: + sublogs["task_desc"] = task.task_desc + if task.truth is not None: + sublogs["truth"] = task.truth + if exps: + last_info = exps[-1].info + if "trajectory" in last_info: + sublogs["trajectory"] = last_info["trajectory"] + if "feedback" in last_info: + sublogs["feedback"] = last_info["feedback"] + + # Deduplicate sys_prompts and hints + sys_prompt_list: List[str] = [] + sys_prompt_map: Dict[str, int] = {} # content -> index + hint_list: List[str] = [] + hint_map: Dict[str, int] = {} # content -> index + + steps: List[dict] = [] + for exp in exps: + taskid = exp.eid.task + runid = exp.eid.run + stepid = exp.eid.step + sys_prompt = exp.info.get("sys_prompt", "(not logged)") + user_prompt = exp.info.get("user_prompt", "(not logged)") + hint = exp.info.get("hint", "(not available)") + response_text = exp.response_text + reward = exp.reward + + # Assign deduplicated index for sys_prompt + if sys_prompt not in sys_prompt_map: + sys_prompt_map[sys_prompt] = len(sys_prompt_list) + sys_prompt_list.append(sys_prompt) + # Assign deduplicated index for hint + if hint not in hint_map: + hint_map[hint] = len(hint_list) + hint_list.append(hint) + + # Simplified step key: strip batch/main/part info (already in dir/file name) + # Original taskid format: batch_X_main_Y_sub_Z[_part_NAME] + # Extract sub_Z portion + sub_part = "" + parts = taskid.split("_") + try: + sub_idx = parts.index("sub") + # Take sub and its value, stop before "part" if present + if "part" in parts: + part_idx = parts.index("part") + sub_part = "_".join(parts[sub_idx:part_idx]) + else: + sub_part = "_".join(parts[sub_idx:]) + except ValueError: + sub_part = taskid # fallback + step_key = f"{sub_part}_run_{runid}_step_{stepid}" + + step_entry = { + "step_key": step_key, + "sys_prompt_idx": sys_prompt_map[sys_prompt], + "hint_idx": hint_map[hint], + "user_prompt": user_prompt, + "response_text": response_text, + "reward": reward, + } + # Additional fields + if "exp_type" in exp.info: + step_entry["exp_type"] = exp.info["exp_type"] + if "incurred_mean_reward" in exp.info: + step_entry["incurred_mean_reward"] = exp.info["incurred_mean_reward"] + if "total_reward" in exp.info: + step_entry["total_reward"] = exp.info["total_reward"] + steps.append(step_entry) + + sublogs["sys_prompts"] = sys_prompt_list + sublogs["hints"] = hint_list + sublogs["steps"] = steps + + # Group steps into episodes by sys_prompt_idx for easier reading + episodes: Dict[int, list] = {} + for step_entry in steps: + ep_id = step_entry["sys_prompt_idx"] + if ep_id not in episodes: + episodes[ep_id] = [] + episodes[ep_id].append(step_entry) + sublogs["episodes"] = episodes + + return sublogs + + # --- CoD methods --- + + async def gen_hint_iteratively( + self, + prev_hint: str, + trajectory: str, + reward: float, + feedback: str, + ) -> Tuple[Experience, str]: + """Run one context-update episode and return (experience, new_hint).""" + self.update_context_workflow.set_context_inputs( + prev_context=prev_hint, + trajectory=trajectory, + reward=reward, + feedback=feedback, + ) + exps = await self.update_context_workflow.run_async() + exp = exps[0] + return exp, exp.info["hint"] + + def _get_or_create_workflow(self, task: Task, idx: int) -> Workflow: + """Get or create workflow instance for a task, updating self.all_workflow_instances. + + Args: + task: The task to create workflow for + idx: Index in self.all_workflow_instances + + Returns: + The workflow instance (either reused or newly created) + """ + if not self.task.workflow_args.get("reuse_workflow_instance", False): + workflow_instance = task.to_workflow( + self._model_for_workflow(task.workflow), + self.auxiliary_model_wrappers, + ) + self.all_workflow_instances[idx] = workflow_instance + return workflow_instance + + workflow_instance = self.all_workflow_instances[idx] + if workflow_instance is None: + workflow_instance = task.to_workflow( + self._model_for_workflow(task.workflow), + self.auxiliary_model_wrappers, + ) + self.all_workflow_instances[idx] = workflow_instance + elif workflow_instance.__class__ == task.workflow and workflow_instance.can_reset: + workflow_instance.reset(task) + else: + workflow_instance = task.to_workflow( + self._model_for_workflow(task.workflow), + self.auxiliary_model_wrappers, + ) + self.all_workflow_instances[idx] = workflow_instance + return workflow_instance + + async def solve_task( + self, + workflow_instance: Workflow, + hint: Optional[str] = None, + icl_examples: Optional[str] = None, + ) -> List[Experience]: + """Solve a task and return all exps (supports multi-step workflows). + + Args: + workflow_instance: Workflow instance to use (must be provided, should be already reset) + hint: Optional hint to guide the model + icl_examples: Optional ICL examples to prepend + + Returns: + List of experiences from solving the task + """ + if hint is not None: + workflow_instance.set_hint(hint=hint) + if icl_examples is not None: + workflow_instance.set_icl_examples(icl_examples) + max_tokens = self.task.workflow_args.get("max_response_tokens_restraint") + if max_tokens: + workflow_instance.set_max_response_tokens_restraint(max_tokens) + exps = await workflow_instance.run_async() + for exp in exps: + exp.info["hint"] = hint if hint else "" + return exps + + async def _run_one_iterative_hint_trajectory_e2e( + self, + part_id: int, + part_name: str, + repeat_id: int, + trajectory_id: str, + ) -> Tuple[List[Experience], float]: + """Run one iterative hint trajectory across all tasks in the pack, for end-to-end meta-training. + + Process: + z0 = "Hints: null" # we might support customizing initial hint later + -> traj T1 for task x1 + -> update z1 + -> traj T2 for task x2 + -> update z2 + -> traj T3 for task x3 + + All samples (including solve_task and gen_hint) will share the same advantage value. + + Assumption: + exp.info should contain fields "trajectory" and "feedback", as input for hint generation. + + Returns: + Tuple of (all_exps, total_reward) + - all_exps: List of all experiences (solve_task exps + gen_hint exps) + - total_reward: Sum of all task rewards (r1 + r2 + r3 + r4) + """ + all_exps: List[Experience] = [] + # gen_hint_exps: List[Experience] = [] + task_rewards: List[float] = [] + current_hint = "Hints: null" + + num_tasks = len(self.all_tasks) + + for task_idx in range(num_tasks): + task = self.all_tasks[task_idx] + taskset_id = get_taskset_id(task) + + workflow_instance = self._get_or_create_workflow(task, task_idx) + + # (1) Solve task + task_exps = await self.solve_task( + workflow_instance=workflow_instance, + hint=current_hint, + ) + + # Post-process solve_task exps + for exp in task_exps: + self._post_process_task_solving_exp( + exp=exp, + subtid=0, # !!! ensure all exps in the trajectory have the same eid.task and eid.run + part_id=part_id, + part_name=part_name, + repeat_id=repeat_id, + ) + exp.info["task_desc"] = task.task_desc + exp.info["trajectory_id"] = trajectory_id + exp.info["task_idx"] = task_idx + exp.info["taskset_id"] = taskset_id + exp.info["exp_type"] = "solve_task" + + # Add position-specific reward and format error metrics (on last exp only) + if task_exps and task_exps[-1].metrics: + task_exps[-1].metrics[f"reward_{part_name}_taskset_{taskset_id}_pos_{task_idx}"] = task_exps[-1].reward + if "format_error_termination" in task_exps[-1].metrics: + task_exps[-1].metrics[f"format_error_{part_name}_taskset_{taskset_id}_pos_{task_idx}"] = task_exps[-1].metrics["format_error_termination"] + + all_exps.extend(task_exps) + + # Get reward and feedback from last exp + task_reward = task_exps[-1].reward if task_exps else 0.0 + task_rewards.append(task_reward) + + # Build trajectory string from all responses + trajectory = self._get_trajectory(task_exps) + + # Get structured feedback from exp.info (set by workflow) + feedback = self._get_feedback(task_exps, task_reward) + + # (2) Generate updated hint (for all tasks except the last one) + if task_idx == num_tasks - 1: + continue + + exp_gen_hint, new_hint = await self.gen_hint_iteratively( + prev_hint=current_hint, + trajectory=trajectory, + reward=task_reward, + feedback=feedback, + ) + + # Post-process gen_hint exp + self._post_process_hint_generation_exp( + exp=exp_gen_hint, + subtid=0, # !!! ensure all exps in the trajectory have the same eid.task and eid.run + part_id=part_id, + part_name=part_name, + repeat_id=repeat_id, + step_id=task_idx, # use task_idx as step_id for hint. TODO: confirm that step_id make no real effect + ) + exp_gen_hint.info["trajectory_id"] = trajectory_id + exp_gen_hint.info["hint_idx"] = task_idx + exp_gen_hint.info["taskset_id"] = taskset_id + exp_gen_hint.info["exp_type"] = "gen_hint" + # gen_hint_exps.append(exp_gen_hint) + all_exps.append(exp_gen_hint) + + current_hint = new_hint + + # Calculate total reward (optionally weighted by hint_reward_decay) + # weight_i = decay^(n-1-i), later tasks weighted more + # e.g., with decay=0.8 and 4 tasks: [0.8^3, 0.8^2, 0.8^1, 0.8^0] + # default decay=1.0 means equal weights (equivalent to plain sum) + task_reward_decay = self.task.workflow_args.get("task_reward_decay", 1.0) + n = len(task_rewards) + total_reward = sum( + r * (task_reward_decay ** (n - 1 - i)) + for i, r in enumerate(task_rewards) + ) + + # Set total reward in exp.info for all exps (for advantage calculation) + for exp in all_exps: + exp.info["total_reward"] = total_reward + exp.info["task_rewards"] = task_rewards + + # Any flagged exp contaminates the shared task_rewards → mask the whole trajectory + # and strip reward-shaped metrics so wandb doesn't log judge-breakage as "wrong answer". + if any(e.info.get("judge_format_error", False) for e in all_exps): + for exp in all_exps: + exp.info["judge_format_error_in_trajectory"] = True + if exp.metrics: + for k in list(exp.metrics.keys()): + if k.startswith("reward") or k.startswith("format_error"): + exp.metrics.pop(k, None) + + return all_exps, total_reward + + async def run_iterative_hint_e2e( + self, + part_id: int, + part_name: str, + subtid: int, + ) -> List[Experience]: + """Run iterative hint for end-to-end meta-training. + + Only runs when subtid == 0 (handles all tasks in one call). + """ + if subtid != 0: + return [] + + all_exps: List[Experience] = [] + + for i in range(self.repeat_times): + trajectory_id = f"batch_{self.task.batch_id}_task_{self.task.task_id}_runidbase_{self.run_id_base}_repeat_{i}" + + exps, total_reward = await self._run_one_iterative_hint_trajectory_e2e( + part_id=part_id, + part_name=part_name, + repeat_id=i, + trajectory_id=trajectory_id, + ) + + all_exps.extend(exps) + + return all_exps + + + def _get_cod_method_mapping(self): + mapping = { + "iterative_hint_e2e": self.run_iterative_hint_e2e, + } + return mapping + + # --- main cod workflow run method --- + + async def run_parts_sequential(self) -> List[Experience]: + """Run diverse CoD methods, record metrics, save logs to json files, return experiences. + + For example, + Part 0: solve directly for each task (baseline method) + Part 1: for each task, first generate hint with aux tasks, then generate solution with hint + Part 2: for each task, first generate hint with aux items, then generate solution with hint + + Things to keep in mind: + + Set task id and step id for each exp carefully, used later by multi-step grpo advantage + + Separate reward metric names that specify rollout methods + + One workflow run -> save logs to one folder, containing one or multiple json files + """ + + cod_method_keys = self.task.workflow_args["activated_cod_methods"] + cod_method_mapping = self._get_cod_method_mapping() + + assert len(cod_method_keys) >= 1 + + returned_exps: List[Experience] = [] + + # --- prepare for logging --- + cod_log_interval = self.task.workflow_args.get("cod_log_interval", 20) + if self.task.is_eval: + explore_step_num = ( + int(str(self.task.batch_id).split("/")[0]) + 1 + ) # (eval) step num starts from 0 + log_prefix = "eval_" + str(self.task.batch_id).split("/")[1] + else: + explore_step_num = int(self.task.batch_id) # (explore) step num starts from 1 + log_prefix = "rollout" + self.save_cod_logs: bool = (explore_step_num - 1) % cod_log_interval == 0 + + save_path = None + if self.save_cod_logs: + log_dir = self.task.workflow_args["log_dir"] + exp_name = self.task.workflow_args["exp_name"] + # batch_id = str(self.task.batch_id) + task_id = str(self.task.task_id) + run_id_base = str(self.run_id_base) + subfolder = f"{log_prefix}_step_{explore_step_num}_task_{task_id}_runidbase_{run_id_base}" + save_path = os.path.join(log_dir, exp_name, subfolder) + # For eval, skip saving if subfolder already exists + if self.task.is_eval and os.path.exists(save_path): + self.save_cod_logs = False + else: + os.makedirs(save_path, exist_ok=True) + + # to be used in later parts that require auxiliary trajectories; aligned with self.all_tasks + self.all_aux_exps: List[Experience] = [] + + # Accumulate per-part summary stats for summary.json + summary_parts: Dict[str, dict] = {} + + # --- one part per activated cod method --- + for part_id, part_name in enumerate(cod_method_keys): + cod_fn = cod_method_mapping[part_name] + exps_part = [] + logs_part = dict() + + async def helper_run_one_task_part(subtid: int): + task = self.all_tasks[subtid] + exps = await cod_fn( + part_id=part_id, + part_name=part_name, + subtid=subtid, + ) + sublogs = self._create_sublogs(task=task, exps=exps) + return exps, sublogs + + # solve multiple tasks in parallel + results_part = await asyncio.gather( + *[helper_run_one_task_part(subtid) for subtid in range(len(self.all_tasks))] + ) + for subtid, rst in enumerate(results_part): + exps, sublogs = rst[0], rst[1] + exps_part.extend(exps) + logs_part[f"subtid_{subtid}"] = sublogs + + returned_exps.extend(exps_part) + if self.save_cod_logs and save_path: + file_name = f"{part_id}_{part_name}.json" + file_path = os.path.join(save_path, file_name) + with open(file_path, "w", encoding="utf-8") as f: + json.dump(logs_part, f, indent=4, ensure_ascii=False) + summary_parts[part_name] = compute_part_summary(exps_part) + + # --- Write summary.json --- + if self.save_cod_logs and save_path and summary_parts: + summary = { + "rollout_step": explore_step_num, + "parts": summary_parts, + } + summary_path = os.path.join(save_path, "summary.json") + with open(summary_path, "w", encoding="utf-8") as f: + json.dump(summary, f, indent=2, ensure_ascii=False) + + return returned_exps + + async def run_parts_parallel(self) -> List[Experience]: + raise NotImplementedError + + async def _maybe_compute_teacher_logprobs(self, exps: List[Experience]) -> None: + """Compute taskset-specific teacher logprobs for CoD training experiences. + + ``auxiliary_model_wrappers`` and training tasksets are matched by their + configuration order. Evaluation is intentionally skipped because its + experiences are not consumed by the trainer. + """ + if ( + not self.task.workflow_args.get("enable_teacher_logprobs", False) + or self.task.is_eval + ): + return + + assert self.auxiliary_model_wrappers is not None + temperature = self.task.rollout_args.temperature + temperature = 1.0 if temperature is None else temperature + + async def score_experience(exp: Experience): + taskset_id = exp.info["taskset_id"] + teacher = self.auxiliary_model_wrappers[taskset_id] + teacher_logprobs = await teacher.logprobs_async( + tokens=exp.tokens.tolist(), + temperature=temperature, + ) + response_start = exp.prompt_length - 1 + teacher_response_logprobs = teacher_logprobs[response_start:] + assert len(teacher_response_logprobs) == len(exp.logprobs), ( + "Teacher/student logprob length mismatch: " + f"teacher={len(teacher_response_logprobs)}, student={len(exp.logprobs)}" + ) + + exp.teacher_logprobs = teacher_response_logprobs + kl_per_token = (exp.logprobs - teacher_response_logprobs) * exp.action_mask + exp.metrics["kl_divergence"] = kl_per_token.sum().item() + + await asyncio.gather(*(score_experience(exp) for exp in exps)) + + async def run_async(self) -> List[Experience]: + """TODO: implement run_parts_parallel if it can bring sufficient gains in efficiency. + + reuse_workflow_instance = self.task.workflow_args["reuse_workflow_instance"] + + if reuse_workflow_instance: + # workflow instance is stateful and resettable, hence need to execute workflow runs sequentially + exps = await self.run_parts_sequential() + else: + # a new workflow instance is created for each workflow run, hence multiple runs can be executed in parallel + exps = await self.run_parts_parallel() + """ + exps = await self.run_parts_sequential() + await self._maybe_compute_teacher_logprobs(exps) + return exps diff --git a/trinity/common/workflows/connect_the_dots/frozen_lake/__init__.py b/trinity/common/workflows/connect_the_dots/frozen_lake/__init__.py new file mode 100644 index 00000000000..07d609659f5 --- /dev/null +++ b/trinity/common/workflows/connect_the_dots/frozen_lake/__init__.py @@ -0,0 +1,30 @@ +# -*- coding: utf-8 -*- +"""CoD (Connect-the-Dots) workflows for FrozenLake environment.""" + +from trinity.common.workflows.connect_the_dots.frozen_lake.workflow import ( + CoDFrozenLakeWorkflow, +) +from trinity.common.workflows.connect_the_dots.frozen_lake.workflow_obscure import ( + CoDFrozenLakeObscureWorkflow, + ALL_PERMUTATIONS, + get_mapping_by_index, + get_random_mapping, +) +from trinity.common.workflows.connect_the_dots.frozen_lake.workflow_obscure_react import ( + CoDFrozenLakeObscureReActWorkflow, +) +from trinity.common.workflows.connect_the_dots.frozen_lake.prompts import ( + load_system_prompt, + load_user_prompt, +) + +__all__ = [ + "CoDFrozenLakeWorkflow", + "CoDFrozenLakeObscureWorkflow", + "CoDFrozenLakeObscureReActWorkflow", + "ALL_PERMUTATIONS", + "get_mapping_by_index", + "get_random_mapping", + "load_system_prompt", + "load_user_prompt", +] diff --git a/trinity/common/workflows/connect_the_dots/frozen_lake/prompts/__init__.py b/trinity/common/workflows/connect_the_dots/frozen_lake/prompts/__init__.py new file mode 100644 index 00000000000..c37eea2daf6 --- /dev/null +++ b/trinity/common/workflows/connect_the_dots/frozen_lake/prompts/__init__.py @@ -0,0 +1,84 @@ +# -*- coding: utf-8 -*- +"""Prompt management for CoD FrozenLake workflow using Jinja2 templates.""" + +from pathlib import Path +from typing import Optional + +from jinja2 import Environment, FileSystemLoader + +# Get the directory where prompt templates are stored +PROMPTS_DIR = Path(__file__).parent + + +def get_jinja_env() -> Environment: + """Get Jinja2 environment with template loader.""" + return Environment( + loader=FileSystemLoader(PROMPTS_DIR), + trim_blocks=True, + lstrip_blocks=True, + ) + + +def load_system_prompt(**kwargs) -> str: + """Load and render system prompt template. + + Args: + **kwargs: Variables to pass to the template. + + Returns: + Rendered system prompt string. + """ + env = get_jinja_env() + template = env.get_template("system.jinja2") + return template.render(**kwargs) + + +def load_system_prompt_obscure(**kwargs) -> str: + """Load and render obscure system prompt template (numeric actions). + + Args: + **kwargs: Variables to pass to the template. + + Returns: + Rendered system prompt string. + """ + env = get_jinja_env() + template = env.get_template("system_obscure.jinja2") + return template.render(**kwargs) + + +def load_user_prompt( + current_step: int, + max_steps: int, + observation: str, + goal_row: int, + goal_col: int, + is_success: bool = False, + **kwargs, +) -> str: + """Load and render user prompt template. + + Args: + current_step: Current step number (1-indexed for display). + max_steps: Maximum number of steps. + observation: Current map observation. + is_success: Whether player has reached the goal. + **kwargs: Additional variables to pass to the template. + + Returns: + Rendered user prompt string. + """ + env = get_jinja_env() + template = env.get_template("user.jinja2") + return template.render( + current_step=current_step, + max_steps=max_steps, + observation=observation, + goal_row=goal_row, + goal_col=goal_col, + is_success=is_success, + **kwargs, + ) + + +__all__ = ["load_system_prompt", "load_system_prompt_obscure", "load_user_prompt", "PROMPTS_DIR"] diff --git a/trinity/common/workflows/connect_the_dots/frozen_lake/prompts/system.jinja2 b/trinity/common/workflows/connect_the_dots/frozen_lake/prompts/system.jinja2 new file mode 100644 index 00000000000..c97f2518067 --- /dev/null +++ b/trinity/common/workflows/connect_the_dots/frozen_lake/prompts/system.jinja2 @@ -0,0 +1,28 @@ +You are playing a FrozenLake game. Your goal is to reach the goal (G) from your current position (P). + +## Symbols +- P: Player (your current position) +- G: Goal (destination) +- _: Frozen tile (safe to walk) +- O: Hole (fall in and lose) + +## Rules +1. Avoid falling into holes (O). +2. Reach the goal (G) to win. +3. The map is a grid with boundaries at its edges. If you try to move beyond the edge, your position will not change and the map will remain the same. + +## Valid Actions +Up | Down | Left | Right + +## Rewards +- Fall into hole: 0 +- Reach goal: +1.0 + +## Response Format +First, wrap your reasoning in ... tags, then provide your final action using ... tags. + +The action must be exactly one of: Up, Down, Left, Right. + +Example: +The goal is to my right, so I should move right. +Right \ No newline at end of file diff --git a/trinity/common/workflows/connect_the_dots/frozen_lake/prompts/system_obscure.jinja2 b/trinity/common/workflows/connect_the_dots/frozen_lake/prompts/system_obscure.jinja2 new file mode 100644 index 00000000000..980cd92e11e --- /dev/null +++ b/trinity/common/workflows/connect_the_dots/frozen_lake/prompts/system_obscure.jinja2 @@ -0,0 +1,36 @@ +You are playing a FrozenLake-Obscure game, where the action-to-direction mapping is unknown. Your task is to reach the goal (G) from your current position (P). + +## Symbols +- P: Player (your current position) +- G: Goal (destination) +- _: Frozen tile (safe to walk) +- O: Hole (fall in and lose) + +## Rules +1. Avoid falling into holes (O). +2. Reach the goal (G) to win. +3. The map is a grid with boundaries at its edges. If you try to move beyond the edge, your position will not change and the map will remain the same. + +## Valid Actions +You can execute one of the following actions at each step: Direction 1, Direction 2, Direction 3, Direction 4. + +## Hidden Clues About the Environment +Each valid action corresponds to moving along one of the up/down/left/right directions, but the exact action-to-direction mapping is unknown. Do not make any prior assumption about the mapping. You may try to figure it out through trial-and-error, e.g., by observing how your position (P) changes on the map after executing an action. You may also utilize additional hints provided in the context. + +## Rewards +- Fall into hole (O): 0 +- Reach goal (G): 1 + +## Response Format +First think about which action to take at the current step, then provide your final answer. Format your complete response as follows: +``` +[THINKING] +Direction X +``` +where [THINKING] should be replaced with your thinking process, and X is one of 1/2/3/4. + +Other requirements: +- Be concise and avoid overthinking in your thinking process. +- Take one and only one action for each step. +- You should reach the goal (G) within a limited number of steps. Try to find a short path towards the goal. +- Balance exploration and exploitation: as you strive to reach the goal, you may acquire information about the environment along the way, which can be potentially helpful for solving other similar tasks later on. diff --git a/trinity/common/workflows/connect_the_dots/frozen_lake/prompts/user.jinja2 b/trinity/common/workflows/connect_the_dots/frozen_lake/prompts/user.jinja2 new file mode 100644 index 00000000000..d562a3e7fd1 --- /dev/null +++ b/trinity/common/workflows/connect_the_dots/frozen_lake/prompts/user.jinja2 @@ -0,0 +1,11 @@ +Step {{ current_step }}/{{ max_steps }} +{% if action_feedback %} + +{{ action_feedback }} +{% endif %} + +Current map: +{{ observation }} +{% if not is_success %} +You have not reached the goal (row={{ goal_row }}, col={{ goal_col }}) yet. Please decide your next action. +{% endif %} diff --git a/trinity/common/workflows/connect_the_dots/frozen_lake/workflow.py b/trinity/common/workflows/connect_the_dots/frozen_lake/workflow.py new file mode 100644 index 00000000000..d4571b9a918 --- /dev/null +++ b/trinity/common/workflows/connect_the_dots/frozen_lake/workflow.py @@ -0,0 +1,371 @@ +# -*- coding: utf-8 -*- +""" +CoD (Connect-the-Dots) workflow for FrozenLake environment using generic multi-turn. + +Each step produces an independent Experience with step number. Prior assistant +responses are compressed to tags only by the base class to avoid +context bloat in multi-turn conversations. +""" + +from __future__ import annotations + +import copy +from typing import TYPE_CHECKING, List, Optional, Tuple + +import numpy as np + +from trinity.common.experience import Experience +from trinity.common.workflows.connect_the_dots.base_workflow import AsyncCoDMultiStepWorkflow +from trinity.common.workflows.connect_the_dots.utils import extract_content_between_keys +from trinity.common.workflows.workflow import Task +from trinity.common.workflows.envs.frozen_lake.utils import ( + GRID_LOOKUP, + MAP_LOOKUP, + generate_random_map, + get_goal_position, +) +from trinity.common.workflows.connect_the_dots.frozen_lake.prompts import ( + load_system_prompt, + load_user_prompt, +) + +if TYPE_CHECKING: + from trinity.common.models.model import ModelWrapper + + +def parse_action(response: str) -> Optional[str]: + """Parse action from model response. + + Args: + response: Model response text. + + Returns: + Action string if found, None otherwise. + """ + VALID_ACTIONS = {"up", "down", "left", "right"} + + key_start, key_end = "", "" + content, success = extract_content_between_keys(response, key_start, key_end) + if not success: + return None + + action = content.strip().lower() + if action in VALID_ACTIONS: + return action.capitalize() + return None + + +class CoDFrozenLakeWorkflow(AsyncCoDMultiStepWorkflow): + """ + CoD FrozenLake workflow using generic multi-turn mechanism. + + This workflow: + - Produces independent Experience per step (with exp.eid.step) + - Prior assistant responses are compressed to only (via base class) + """ + + is_async: bool = True + + def __init__( + self, + model: ModelWrapper, + task: Task, + auxiliary_models: Optional[List] = None, + use_openai_client: bool = False, + ): + super().__init__( + model=model, + task=task, + auxiliary_models=auxiliary_models, + use_openai_client=use_openai_client, + ) + + # Import gymnasium here to avoid import error if not installed + try: + import gymnasium as gym + from gymnasium.envs.toy_text.frozen_lake import FrozenLakeEnv as GymFrozenLakeEnv + except ImportError as e: + raise ImportError( + f"Gymnasium is not installed. Please install gymnasium first. Error: {e}" + ) + + self._gym = gym + self._GymFrozenLakeEnv = GymFrozenLakeEnv + + # Action mapping: our action -> gym action + self.action_map = { + "Left": 0, + "Down": 1, + "Right": 2, + "Up": 3, + } + + self.reset(task) + + def reset(self, task: Task): + """Reset task-specific config, environment, and episode state.""" + super().reset(task) + + # Extract workflow-specific arguments + workflow_args = task.workflow_args if hasattr(task, "workflow_args") else {} + self.env_max_steps = workflow_args.get("env_max_steps", 8) + self.agent_max_steps = workflow_args.get("agent_max_steps", 10) + self.desc = workflow_args.get("desc", None) + self.is_slippery = workflow_args.get("is_slippery", False) + self.max_response_tokens_restraint = workflow_args.get("max_response_tokens_restraint", None) + + # Extract task-specific arguments + self.raw_task = task.raw_task if hasattr(task, "raw_task") else {} + if self.raw_task is None: + self.raw_task = {} + self.size = self.raw_task.get("size", 4) + self.p = self.raw_task.get("p", 0.8) + self.seed = self.raw_task.get("seed", 42) + + # Generate or use provided map + if self.desc is None: + random_map, goal_position = generate_random_map( + size=self.size, p=self.p, seed=self.seed, max_steps=self.env_max_steps + ) + else: + random_map = np.asarray(copy.deepcopy(self.desc), dtype="c") + goal_position = get_goal_position(random_map) + + self.goal_position = goal_position + + # Create gym environment + self.gym_env = self._GymFrozenLakeEnv( + desc=random_map[:], + is_slippery=self.is_slippery, + ) + self.action_space = self._gym.spaces.Discrete(4, start=1) + + # State variables + self.observation: Optional[str] = None + self.done: bool = False + self.final_reward: float = 0.0 + self.early_termination_by_format_issue = False + self.memory: List[dict] = [] + self.current_step: int = 0 + self.action_feedback = None + + def _get_player_position(self): + """Get current player position as (row, col).""" + return ( + self.gym_env.s // self.gym_env.ncol, + self.gym_env.s % self.gym_env.ncol, + ) + + def _is_success(self) -> bool: + """Check if player reached the goal.""" + player_pos = self._get_player_position() + return self.gym_env.desc[player_pos] == b"G" + + def render(self) -> str: + """Render current state as text.""" + room_state = copy.deepcopy(self.gym_env.desc) + + # Replace start 'S' with frozen 'F' + position_S = np.where(room_state == b"S") + room_state[position_S] = b"F" + + # Mark player position + position_P = self._get_player_position() + room_state[position_P] = b"P" + + # Convert to state array + state_array = np.vectorize(lambda x: MAP_LOOKUP[x])(room_state) + + # Handle player on hole or goal + if self.gym_env.desc[position_P] == b"H": + state_array[position_P] = 4 # player fell into hole + elif self.gym_env.desc[position_P] == b"G": + state_array[position_P] = 5 # player on goal + + # Render as text + result = "\n".join( + "".join(GRID_LOOKUP.get(cell, "?") for cell in row) + for row in state_array.tolist() + ) + return result + + def _build_action_feedback( + self, action_str: Optional[str], prev_pos: Tuple[int, int], + cur_pos: Tuple[int, int], action_effective: bool, + ) -> str: + """Build position change feedback string after an action. + + Args: + action_str: The action taken (e.g. "Up" or None if invalid). + prev_pos: (row, col) before action. + cur_pos: (row, col) after action. + action_effective: Whether the position changed. + + Returns: + Feedback string describing position change. + """ + if action_str is None: + return "Your action was invalid. Position unchanged." + if action_effective: + return ( + f"You executed action {action_str}. " + f"Your position changed from (row={prev_pos[0]}, col={prev_pos[1]}) " + f"to (row={cur_pos[0]}, col={cur_pos[1]})." + ) + else: + return ( + f"You executed action {action_str}. " + f"Your position did not change (stayed at row={cur_pos[0]}, col={cur_pos[1]}). " + f"This may be because you tried to move beyond the edge of the map." + ) + + def env_step(self, action: Optional[str]): + """Execute action in environment. + + Args: + action: Action string (Up/Down/Left/Right) or None for invalid. + + Returns: + Tuple of (observation, reward, done, info). + """ + if self._is_success(): + return self.render(), 1.0, True, {"action_is_effective": False} + + if action is None or action not in self.action_map: + return self.render(), 0.0, False, {"action_is_effective": False} + + prev_pos = int(self.gym_env.s) + + # Execute action + _, reward, done, _, _ = self.gym_env.step(self.action_map[action]) + + obs = self.render() + action_effective = prev_pos != int(self.gym_env.s) + + return obs, reward, done, {"action_is_effective": action_effective} + + def _build_system_prompt(self) -> str: + """Build system prompt with hint and token limit.""" + sys_prompt = load_system_prompt() + sys_prompt = self._augment_system_prompt(sys_prompt) + return sys_prompt + + async def run_async(self) -> List[Experience]: + """Run the workflow.""" + # Reset environment + self.gym_env.reset(seed=self.seed) + self.observation = self.render() + self.done = False + self.final_reward = 0.0 + self.current_step = 0 + self.action_feedback = None # Position change feedback for next step + + # Initialize memory with system prompt (includes hint and token limit) + self.memory.clear() + self.memory.append({"role": "system", "content": self._build_system_prompt()}) + + return await super().run_async() + + async def step_async(self, step_num: int) -> Tuple[bool, List[Experience]]: + """Execute one step of the workflow. + + Args: + step_num: Current step number. + + Returns: + Tuple of (continue_run, experiences): + - continue_run: True to continue, False to stop. + - experiences: List of experiences from this step. + """ + if self.done: + return False, [] + + # Format observation as user message using jinja template + user_content = load_user_prompt( + current_step=step_num + 1, + max_steps=self.agent_max_steps, + observation=self.observation, + goal_row=self.goal_position[0], + goal_col=self.goal_position[1], + is_success=self._is_success(), + action_feedback=self.action_feedback, + ) + + # ========== CoD-specific: add icl_examples and reply_prefix ========== + # icl_examples: only added in the first turn to avoid token waste + if self.icl_examples and step_num == 0: + user_content = f"{user_content}\n\nHere are some reference examples:\n\n{self.icl_examples}" + + self.memory.append({"role": "user", "content": user_content}) + + # reply_prefix: prefix for model's reply, used for guided generation + if self.reply_prefix: + self.memory.append({"role": "assistant", "content": self.reply_prefix}) + # ==================================================================== + + # Get model response - chat_async returns List[Experience] directly + experiences = await self.model.chat_async(self.memory) + response_text = experiences[0].response_text + self.memory.append({"role": "assistant", "content": response_text}) + + # Store prompt information in exp.info for logging + sys_prompt = self.memory[0]["content"] if self.memory and self.memory[0]["role"] == "system" else "" + for exp in experiences: + exp.info["sys_prompt"] = sys_prompt + exp.info["user_prompt"] = user_content + + # Parse action + action = parse_action(response_text) + + # Track position before action + prev_pos = self._get_player_position() + + # Execute action + observation, reward, done, info = self.env_step(action) + + # Build action feedback for next step + cur_pos = self._get_player_position() + self.action_feedback = self._build_action_feedback( + action_str=action, prev_pos=prev_pos, cur_pos=cur_pos, + action_effective=info.get("action_is_effective", False), + ) + + # Update state + self.observation = observation + self.done = done + self.current_step = step_num + 1 + + if done and reward > 0: + self.final_reward = reward + + return not self.done, experiences + + def _get_feedback(self) -> str: + """Generate structured feedback about task completion. + + Returns feedback relevant to reward calculation: + - Success: reached the goal + - Failure: fell into hole or ran out of steps + """ + if self._is_success(): + return "Success! Reached the goal (G)." + + player_pos = self._get_player_position() + if self.gym_env.desc[player_pos] == b"H": + return "Failed: Fell into a hole (O). Avoid holes by planning a safe path." + + if self.early_termination_by_format_issue: + return self.action_feedback + + # Ran out of steps + return f"Failed: Ran out of steps ({self.current_step}/{self.agent_max_steps}) without reaching the goal. Try to find a shorter path." + + @property + def max_step_num(self) -> int: + """Maximum number of steps.""" + return self.agent_max_steps + + def __del__(self): + """Cleanup environment.""" + if hasattr(self, "gym_env"): + self.gym_env.close() diff --git a/trinity/common/workflows/connect_the_dots/frozen_lake/workflow_obscure.py b/trinity/common/workflows/connect_the_dots/frozen_lake/workflow_obscure.py new file mode 100644 index 00000000000..055f7598d81 --- /dev/null +++ b/trinity/common/workflows/connect_the_dots/frozen_lake/workflow_obscure.py @@ -0,0 +1,222 @@ +# -*- coding: utf-8 -*- +""" +CoD FrozenLake Obscure Workflow - uses numeric actions (1,2,3,4) instead of directions. + +Mapping modes: +- "global": All episodes use the same fixed mapping +- "per_pack": Tasks in the same CoD pack share one mapping +- "per_task": Each task uses its own mapping (via raw_task) +""" + +from __future__ import annotations + +import itertools +import random +import re +from typing import Dict, List, Optional, Tuple, TYPE_CHECKING + +from trinity.common.experience import Experience +from trinity.common.workflows.connect_the_dots.frozen_lake.workflow import ( + CoDFrozenLakeWorkflow, +) +from trinity.common.workflows.connect_the_dots.utils import extract_content_between_keys + +if TYPE_CHECKING: + from trinity.common.models.model import ModelWrapper + from trinity.common.workflows.workflow import Task + + +# All 24 permutations of directions +ALL_PERMUTATIONS = list(itertools.permutations(["Down", "Right", "Up", "Left"])) + +# Default mapping: 1=Down, 2=Right, 3=Up, 4=Left +DEFAULT_ACTION_MAPPING = {1: "Down", 2: "Right", 3: "Up", 4: "Left"} + + +def get_mapping_by_index(index: int) -> Dict[int, str]: + """Get action mapping by permutation index (0-23).""" + perm = ALL_PERMUTATIONS[index % len(ALL_PERMUTATIONS)] + return {i + 1: perm[i] for i in range(4)} + + +def get_random_mapping(seed: Optional[int] = None) -> Dict[int, str]: + """Get a random action mapping.""" + rng = random.Random(seed) + directions = ["Down", "Right", "Up", "Left"] + rng.shuffle(directions) + return {i + 1: directions[i] for i in range(4)} + + +def parse_numeric_action_number(response: str) -> Optional[int]: + """Parse 'Direction X' action from response, return the integer or None.""" + key_start, key_end = "", "" + content, success = extract_content_between_keys(response, key_start, key_end) + if not success: + return None + action_str = content.strip() + dir_match = re.match(r"[Dd]irection\s+(\d+)", action_str) + if dir_match: + return int(dir_match.group(1)) + return None + + +class CoDFrozenLakeObscureWorkflow(CoDFrozenLakeWorkflow): + """ + FrozenLake with numeric actions (1,2,3,4) instead of Up/Down/Left/Right. + + Configuration (workflow_args): + - mapping_mode: "global", "per_pack", or "per_task" + - "global": All tasks use the same fixed mapping + - "per_pack": All tasks in one CoD pack share one mapping, + different packs get different mappings + - "per_task": Each task uses its own mapping + - global_mapping: {1: "Down", ...} for global mode + - mapping_index: 0-23 for predefined permutations (global mode) + + Per-task (raw_task, for per_task mode): + - action_mapping: explicit mapping + - mapping_index: permutation index + - mapping_seed: seed for random mapping + """ + + def __init__( + self, + model: ModelWrapper, + task: Task, + auxiliary_models: Optional[List] = None, + use_openai_client: bool = False, + ): + super().__init__( + model=model, + task=task, + auxiliary_models=auxiliary_models, + use_openai_client=use_openai_client, + ) + + def reset(self, task: Task): + """Reset FrozenLake state and recompute numeric action mapping.""" + workflow_args = task.workflow_args if hasattr(task, "workflow_args") else {} + self.mapping_mode = workflow_args.get("mapping_mode", "global") + self.global_mapping = workflow_args.get("global_mapping", None) + self.global_mapping_index = workflow_args.get("mapping_index", None) + + super().reset(task) + self.action_mapping = self._determine_action_mapping() + + def _determine_action_mapping(self) -> Dict[int, str]: + """Determine action mapping based on mode. + + Modes: + - "global": All tasks use the same mapping (from workflow_args) + - "per_pack": All tasks in the same CoD pack share one mapping, + different packs use different mappings + - "per_task": Each task uses its own mapping (from raw_task) + """ + if self.mapping_mode == "per_pack": + # Use pack-level seed injected by CoDWorkflow.reset() + if "pack_seed" in self.raw_task: + return get_random_mapping(self.raw_task["pack_seed"]) + else: + # Fallback: not running under CoDWorkflow, use task seed + return get_random_mapping(self.raw_task.get("seed", 42)) + elif self.mapping_mode == "per_task": + if "action_mapping" in self.raw_task: + return self.raw_task["action_mapping"] + elif "mapping_index" in self.raw_task: + return get_mapping_by_index(self.raw_task["mapping_index"]) + elif "mapping_seed" in self.raw_task: + return get_random_mapping(self.raw_task["mapping_seed"]) + else: + return get_random_mapping(self.raw_task.get("seed", 42)) + else: + # "global" mode + if self.global_mapping is not None: + return self.global_mapping + elif self.global_mapping_index is not None: + return get_mapping_by_index(self.global_mapping_index) + else: + return DEFAULT_ACTION_MAPPING + + def _build_system_prompt(self) -> str: + """Build system prompt with numeric actions using jinja template.""" + from trinity.common.workflows.connect_the_dots.frozen_lake.prompts import load_system_prompt_obscure + + sys_prompt = load_system_prompt_obscure() + sys_prompt = self._augment_system_prompt(sys_prompt) + return sys_prompt + + async def step_async(self, step_num: int) -> Tuple[bool, List[Experience]]: + """Execute one step - override to use numeric action parsing.""" + if self.done: + return False, [] + + from trinity.common.workflows.connect_the_dots.frozen_lake.prompts import load_user_prompt + + user_content = load_user_prompt( + current_step=step_num + 1, + max_steps=self.agent_max_steps, + observation=self.observation, + goal_row=self.goal_position[0], + goal_col=self.goal_position[1], + is_success=self._is_success(), + action_feedback=self.action_feedback, + ) + + if self.icl_examples and step_num == 0: + user_content = f"{user_content}\n\nHere are some reference examples:\n\n{self.icl_examples}" + + self.memory.append({"role": "user", "content": user_content}) + + if self.reply_prefix: + self.memory.append({"role": "assistant", "content": self.reply_prefix}) + + experiences = await self.model.chat_async(self.memory) + response_text = experiences[0].response_text + self.memory.append({"role": "assistant", "content": response_text}) + + sys_prompt = self.memory[0]["content"] if self.memory and self.memory[0]["role"] == "system" else "" + for exp in experiences: + exp.info["sys_prompt"] = sys_prompt + exp.info["user_prompt"] = user_content + exp.info["action_mapping"] = self.action_mapping + + # Check for malformed response: missing tag + numeric_action = parse_numeric_action_number(response_text) + if numeric_action is None: + # Early termination: reward 0, give feedback about format + self.action_feedback = "Invalid format: could not parse action. Expected format: {your reasoning process here}Direction X, where X is 1, 2, 3, or 4. Game over." + self.done = True + self.current_step = step_num + 1 + self.early_termination_by_format_issue = True + return False, experiences + + direction = self.action_mapping.get(numeric_action) + if direction is None: + # Valid format but invalid action number (e.g. Direction 5) + self.action_feedback = f"Invalid format: Direction {numeric_action} is not a valid action. Expected: Direction 1, 2, 3, or 4. Game over." + self.done = True + self.current_step = step_num + 1 + self.early_termination_by_format_issue = True + return False, experiences + + # Track position before action + prev_pos = self._get_player_position() + + observation, reward, done, info = self.env_step(direction) + + # Build action feedback with "Direction X" format + cur_pos = self._get_player_position() + self.action_feedback = self._build_action_feedback( + action_str=f"Direction {numeric_action}", + prev_pos=prev_pos, cur_pos=cur_pos, + action_effective=info.get("action_is_effective", False), + ) + + self.observation = observation + self.done = done + self.current_step = step_num + 1 + + if done and reward > 0: + self.final_reward = reward + + return not self.done, experiences diff --git a/trinity/common/workflows/connect_the_dots/frozen_lake/workflow_obscure_react.py b/trinity/common/workflows/connect_the_dots/frozen_lake/workflow_obscure_react.py new file mode 100644 index 00000000000..a31c1ad8110 --- /dev/null +++ b/trinity/common/workflows/connect_the_dots/frozen_lake/workflow_obscure_react.py @@ -0,0 +1,150 @@ +# -*- coding: utf-8 -*- +"""ReAct-agent version of the CoD FrozenLake-Obscure solve workflow. + +An AgentScope ReActAgent emits one numeric action per environment turn; the +workflow steps the environment and feeds the next observation back. Each model +call inside a turn is captured via an isolated-history model clone. +""" + +from __future__ import annotations + +from typing import List, Tuple + +from trinity.common.experience import Experience +from trinity.common.workflows.connect_the_dots.agentscope_utils import ( + build_agentscope_react_agent, + run_agentscope_agent_step, +) +from trinity.common.workflows.connect_the_dots.frozen_lake.prompts import load_user_prompt +from trinity.common.workflows.connect_the_dots.frozen_lake.workflow_obscure import ( + CoDFrozenLakeObscureWorkflow, + parse_numeric_action_number, +) + + +class CoDFrozenLakeObscureReActWorkflow(CoDFrozenLakeObscureWorkflow): + """FrozenLake-Obscure solved by a ReActAgent that acts once per turn. + + Env, mapping, prompts, feedback and reward are inherited from the parent; + this class replaces only the generation path. + """ + + is_async: bool = True + can_reset: bool = True + requires_isolated_model_history: bool = True + + async def run_async(self) -> List[Experience]: + # Reset env and per-episode state. + self.gym_env.reset(seed=self.seed) + self.observation = self.render() + self.done = False + self.final_reward = 0.0 + self.current_step = 0 + self.action_feedback = None + self.early_termination_by_format_issue = False + + # reply_prefix cannot be injected through the agent loop; fail loudly + # rather than diverge silently. + if self.reply_prefix: + raise NotImplementedError( + "reply_prefix is not supported by the ReAct solve workflow." + ) + + # Rebuild the agent each run since the hint in the system prompt changes. + self.model.history.clear() + sys_prompt = self._build_system_prompt() + agent = await build_agentscope_react_agent( + name="cod_frozenlake_obscure", + model=self.model, + system_prompt=sys_prompt, + compress_assistant_fn=self._compress_assistant_response, + max_iters=self.task.workflow_args.get("react_max_iters", 1), + ) + + history_spans: List[Tuple[int, int, str]] = [] + for step_num in range(self.agent_max_steps): + if self.done: + break + user_content = load_user_prompt( + current_step=step_num + 1, + max_steps=self.agent_max_steps, + observation=self.observation, + goal_row=self.goal_position[0], + goal_col=self.goal_position[1], + is_success=self._is_success(), + action_feedback=self.action_feedback, + ) + if self.icl_examples and step_num == 0: + user_content = ( + f"{user_content}\n\nHere are some reference examples:\n\n" + f"{self.icl_examples}" + ) + + history_start = len(self.model.history) + response_text = await run_agentscope_agent_step(agent, user_content) + history_spans.append( + (history_start, len(self.model.history), user_content) + ) + + numeric_action = parse_numeric_action_number(response_text) + if numeric_action is None: + self.action_feedback = ( + "Invalid format: could not parse action. Expected format: " + "{your reasoning process here}Direction X, " + "where X is 1, 2, 3, or 4. Game over." + ) + self.done = True + self.early_termination_by_format_issue = True + self.current_step = step_num + 1 + break + + direction = self.action_mapping.get(numeric_action) + if direction is None: + self.action_feedback = ( + f"Invalid format: Direction {numeric_action} is not a valid " + "action. Expected: Direction 1, 2, 3, or 4. Game over." + ) + self.done = True + self.early_termination_by_format_issue = True + self.current_step = step_num + 1 + break + + prev_pos = self._get_player_position() + observation, reward, done, info = self.env_step(direction) + cur_pos = self._get_player_position() + self.action_feedback = self._build_action_feedback( + action_str=f"Direction {numeric_action}", + prev_pos=prev_pos, + cur_pos=cur_pos, + action_effective=info.get("action_is_effective", False), + ) + self.observation = observation + self.done = done + self.current_step = step_num + 1 + if done and reward > 0: + self.final_reward = reward + + experiences = self.model.extract_experience_from_history() + exp_step = 0 + for history_start, history_end, user_content in history_spans: + for exp in experiences[history_start:history_end]: + exp.eid.step = exp_step + exp.info["sys_prompt"] = sys_prompt + exp.info["user_prompt"] = user_content + exp.info["action_mapping"] = self.action_mapping + exp_step += 1 + + trajectory = self._build_agentscope_trajectory( + sys_prompt, + await agent.memory.get_memory(prepend_summary=False), + ) + + reward = await self.reward_async(experiences) + for exp in experiences: + exp.reward = reward + if exp.metrics is None: + exp.metrics = {} + if experiences: + experiences[-1].metrics["actual_env_steps"] = self.current_step + experiences[-1].info["trajectory"] = trajectory + return experiences diff --git a/trinity/common/workflows/connect_the_dots/grid_navigation/__init__.py b/trinity/common/workflows/connect_the_dots/grid_navigation/__init__.py new file mode 100644 index 00000000000..108fb8a00d5 --- /dev/null +++ b/trinity/common/workflows/connect_the_dots/grid_navigation/__init__.py @@ -0,0 +1,17 @@ +"""CoD grid-navigation environment and workflow.""" + +from trinity.common.workflows.connect_the_dots.grid_navigation.env import ( + GridNavigationEnv, + GridNavigationPackState, + GridNavigationTask, +) +from trinity.common.workflows.connect_the_dots.grid_navigation.workflow import ( + CoDGridNavigationWorkflow, +) + +__all__ = [ + "CoDGridNavigationWorkflow", + "GridNavigationEnv", + "GridNavigationPackState", + "GridNavigationTask", +] diff --git a/trinity/common/workflows/connect_the_dots/grid_navigation/env.py b/trinity/common/workflows/connect_the_dots/grid_navigation/env.py new file mode 100644 index 00000000000..a9444e87a0c --- /dev/null +++ b/trinity/common/workflows/connect_the_dots/grid_navigation/env.py @@ -0,0 +1,267 @@ +"""Grid-navigation environment used by the CoD research workflow.""" + +from __future__ import annotations + +import weakref +from dataclasses import dataclass +from typing import Optional, Tuple + +import numpy as np + +Position = Tuple[int, int] + + +def _generate_landscape_costs( + *, + rng: np.random.Generator, + size: int, + num_components: int, + min_scale: float, + max_scale: float, +) -> np.ndarray: + """Generate a spatially correlated integer cost landscape.""" + if num_components < 2: + raise ValueError("landscape_num_components must be at least 2.") + if not 0.0 < min_scale <= max_scale: + raise ValueError("landscape scales must satisfy 0 < min_scale <= max_scale.") + + coordinates = np.linspace(0.0, 1.0, size) + row_grid, col_grid = np.meshgrid(coordinates, coordinates, indexing="ij") + landscape = np.zeros((size, size), dtype=np.float64) + + for _ in range(num_components): + center_row, center_col = rng.uniform(0.0, 1.0, size=2) + long_scale = float(rng.uniform(min_scale, max_scale)) + short_scale = float(rng.uniform(min_scale, long_scale)) + angle = float(rng.uniform(0.0, np.pi)) + amplitude = float(rng.uniform(0.8, 1.2)) + + row_delta = row_grid - center_row + col_delta = col_grid - center_col + cos_angle = np.cos(angle) + sin_angle = np.sin(angle) + long_delta = cos_angle * row_delta + sin_angle * col_delta + short_delta = -sin_angle * row_delta + cos_angle * col_delta + landscape += amplitude * np.exp( + -0.5 * ((long_delta / long_scale) ** 2 + (short_delta / short_scale) ** 2) + ) + + # Positive Gaussian components form costly hills over a connected low-cost + # background. Squaring the normalized height makes routes through valleys cheap + # without flattening the spatial differences through a rank transformation. + landscape_min = float(landscape.min()) + landscape_span = float(landscape.max() - landscape_min) + if landscape_span == 0.0: + return np.zeros_like(landscape, dtype=np.int64) + normalized = (landscape - landscape_min) / landscape_span + costs = np.rint(99.0 * np.square(normalized)) + return costs.astype(np.int64) + + +@dataclass +class GridNavigationPackState: + """The cost map and observations shared by all tasks in one CoD pack.""" + + pack_seed: int + costs: np.ndarray + revealed: np.ndarray + + @classmethod + def generate( + cls, + pack_seed: int, + grid_min_size: int, + grid_max_size: int, + landscape_num_components: int = 6, + landscape_min_scale: float = 0.10, + landscape_max_scale: float = 0.35, + ) -> "GridNavigationPackState": + """Generate a deterministic square cost map for one pack.""" + rng = np.random.default_rng(pack_seed) + size = int(rng.integers(grid_min_size, grid_max_size + 1)) + costs = _generate_landscape_costs( + rng=rng, + size=size, + num_components=landscape_num_components, + min_scale=landscape_min_scale, + max_scale=landscape_max_scale, + ) + return cls( + pack_seed=pack_seed, + costs=costs, + revealed=np.zeros_like(costs, dtype=bool), + ) + + @property + def size(self) -> int: + """Return the side length of the square grid.""" + return int(self.costs.shape[0]) + + def reveal(self, path: list[Position], radius: int) -> None: + """Reveal the Chebyshev neighborhood of every entered cell.""" + for row, col in path: + row_start = max(0, row - radius) + row_end = min(self.size, row + radius + 1) + col_start = max(0, col - radius) + col_end = min(self.size, col + radius + 1) + self.revealed[row_start:row_end, col_start:col_end] = True + + +_PACK_STATES: weakref.WeakValueDictionary = weakref.WeakValueDictionary() + + +def get_pack_state( + *, + pack_seed: int, + task_idx: int, + grid_min_size: int, + grid_max_size: int, + landscape_num_components: int = 6, + landscape_min_scale: float = 0.10, + landscape_max_scale: float = 0.35, +) -> GridNavigationPackState: + """Create the state for task 0, and reuse it for later tasks in the pack.""" + key = pack_seed + if task_idx == 0: + state = GridNavigationPackState.generate( + pack_seed=pack_seed, + grid_min_size=grid_min_size, + grid_max_size=grid_max_size, + landscape_num_components=landscape_num_components, + landscape_min_scale=landscape_min_scale, + landscape_max_scale=landscape_max_scale, + ) + _PACK_STATES[key] = state + return state + return _PACK_STATES[key] + + +@dataclass(frozen=True) +class GridNavigationTask: + """Task-specific start, goal, and exact round count.""" + + start: Position + goal: Position + num_rounds: int + + +def _sample_other_coordinate(rng: np.random.Generator, current: int, size: int) -> int: + value = int(rng.integers(0, size - 1)) + return value + 1 if value >= current else value + + +def generate_task( + *, + task_seed: int, + grid_size: int, + min_rounds: int, + max_rounds: int, +) -> GridNavigationTask: + """Generate a task together with an implicit exact-length witness path.""" + rng = np.random.default_rng(task_seed) + num_rounds = int(rng.integers(min_rounds, max_rounds + 1)) + + while True: + start = ( + int(rng.integers(0, grid_size)), + int(rng.integers(0, grid_size)), + ) + current = start + for _ in range(num_rounds): + row, col = current + if int(rng.integers(0, 2)) == 0: + current = (row, _sample_other_coordinate(rng, col, grid_size)) + else: + current = (_sample_other_coordinate(rng, row, grid_size), col) + if current != start: + return GridNavigationTask(start=start, goal=current, num_rounds=num_rounds) + + +class GridNavigationEnv: + """One fixed-round navigation task over a shared cost map.""" + + def __init__( + self, + *, + pack_state: GridNavigationPackState, + task: GridNavigationTask, + reveal_radius: int, + ): + """Initialize one task over an existing pack state.""" + self.pack_state = pack_state + self.task = task + self.reveal_radius = reveal_radius + self.reset_task() + + def reset_task(self) -> None: + """Reset task-local state while retaining pack observations.""" + self.current_position = self.task.start + self.rounds_taken = 0 + self.entered_costs: list[int] = [] + self.last_path: list[Position] = [] + + @property + def normalized_loss(self) -> float: + """Return the visit-weighted mean cell cost normalized to [0, 1].""" + if not self.entered_costs: + return 0.0 + return float(sum(self.entered_costs) / (100.0 * len(self.entered_costs))) + + def validate_destination(self, destination: Position) -> Optional[str]: + """Return an error for an illegal rook move, otherwise None.""" + row, col = destination + cur_row, cur_col = self.current_position + if not (0 <= row < self.pack_state.size and 0 <= col < self.pack_state.size): + return f"Destination ({row},{col}) is outside the grid." + if destination == self.current_position: + return "The destination must differ from the current position." + if row != cur_row and col != cur_col: + return "The destination must be in the same row or the same column." + return None + + def path_to(self, destination: Position) -> list[Position]: + """List entered cells, excluding the origin and including the destination.""" + cur_row, cur_col = self.current_position + dest_row, dest_col = destination + if cur_row == dest_row: + step = 1 if dest_col > cur_col else -1 + return [(cur_row, col) for col in range(cur_col + step, dest_col + step, step)] + step = 1 if dest_row > cur_row else -1 + return [(row, cur_col) for row in range(cur_row + step, dest_row + step, step)] + + def step(self, destination: Position) -> dict: + """Execute one previously validated move.""" + path = self.path_to(destination) + step_costs = [int(self.pack_state.costs[position]) for position in path] + self.entered_costs.extend(step_costs) + self.pack_state.reveal(path, self.reveal_radius) + self.current_position = destination + self.rounds_taken += 1 + self.last_path = path + + done = self.rounds_taken == self.task.num_rounds + success = done and self.current_position == self.task.goal + reward = 1.0 - self.normalized_loss if success else 0.0 + return { + "path": path, + "normalized_loss": self.normalized_loss, + "done": done, + "success": success, + "reward": reward, + } + + def render(self) -> str: + """Render row and column labels with known costs and question marks.""" + header = " " + "".join(f"{col:>3}" for col in range(self.pack_state.size)) + rows = [header] + for row in range(self.pack_state.size): + cells = [] + for col in range(self.pack_state.size): + value = ( + str(int(self.pack_state.costs[row, col])) + if self.pack_state.revealed[row, col] + else "?" + ) + cells.append(f"{value:>3}") + rows.append(f"{row:>3} " + "".join(cells)) + return "\n".join(rows) diff --git a/trinity/common/workflows/connect_the_dots/grid_navigation/prompts/__init__.py b/trinity/common/workflows/connect_the_dots/grid_navigation/prompts/__init__.py new file mode 100644 index 00000000000..4612b8e3721 --- /dev/null +++ b/trinity/common/workflows/connect_the_dots/grid_navigation/prompts/__init__.py @@ -0,0 +1,54 @@ +"""Prompt loading helpers for the CoD grid-navigation workflow.""" + +from pathlib import Path +from typing import Optional, Tuple + +from jinja2 import Environment, FileSystemLoader + +PROMPTS_DIR = Path(__file__).parent + + +def get_jinja_env() -> Environment: + return Environment( + loader=FileSystemLoader(PROMPTS_DIR), + trim_blocks=True, + lstrip_blocks=True, + ) + + +def load_system_prompt(*, reveal_radius: int) -> str: + """Render the system prompt with the configured observation radius.""" + return ( + get_jinja_env() + .get_template("system.jinja2") + .render( + reveal_radius=reveal_radius, + ) + ) + + +def load_user_prompt( + *, + current_round: int, + max_rounds: int, + current_position: Tuple[int, int], + goal_position: Tuple[int, int], + observation: str, + action_feedback: Optional[str], +) -> str: + """Render one round's grid observation and task state.""" + return ( + get_jinja_env() + .get_template("user.jinja2") + .render( + current_round=current_round, + max_rounds=max_rounds, + current_position=current_position, + goal_position=goal_position, + observation=observation, + action_feedback=action_feedback, + ) + ) + + +__all__ = ["load_system_prompt", "load_user_prompt", "PROMPTS_DIR"] diff --git a/trinity/common/workflows/connect_the_dots/grid_navigation/prompts/system.jinja2 b/trinity/common/workflows/connect_the_dots/grid_navigation/prompts/system.jinja2 new file mode 100644 index 00000000000..1deaf7c15b1 --- /dev/null +++ b/trinity/common/workflows/connect_the_dots/grid_navigation/prompts/system.jinja2 @@ -0,0 +1,35 @@ +You are navigating a two-dimensional cost grid. Your goal is to finish the final round exactly at the target position while keeping the normalized navigation loss low. + +## Grid and observations +- Rows and columns use zero-based indices. +- Every cell has a fixed integer cost from 0 to 99. +- Unknown costs are shown as `?`. Observed costs are shown as integers. +- The current task belongs to a stream of tasks within the same environment. Every task in this stream shares the same fixed cost grid. +- After each move, the environment reveals every cell whose Chebyshev distance from an entered cell is at most {{ reveal_radius }}. These observations are retained automatically in later rounds and later tasks in the stream. + +## Movement rules +- The task has a fixed number of rounds, and you must take exactly one move per round. +- A move selects any other cell in the same row or the same column as your current position. +- You may cross any positive number of cells up to the grid boundary. +- The entered cells include every crossed cell and the destination, but not the starting cell of the move. +- Reaching the target before the final round does not end the task. You must continue moving and be at the target when the final round ends. + +## Loss and reward +Let V be the list of all entered cells over all rounds. Repeated visits appear in V repeatedly. The mean entered-cell cost and normalized loss are: + +mean_cost = sum(cost(v) for v in V) / |V| +loss = mean_cost / 100 = sum(cost(v) for v in V) / (100 * |V|) + +The loss depends on the mean cost, not the total cost. There is no separate penalty for path length or the number of entered cells, so entering fewer cells is not inherently better. A longer route can lower the final loss when the additional cells are cheaper than the running mean, while an expensive detour raises it. Do not default to a conventional shortest-path or minimum-total-cost objective. + +- If your position after the final round is exactly the target, reward = 1 - loss. +- Otherwise, reward = 0. +- An invalidly formatted or illegal move ends the task with reward 0. + +Explore broadly across the full grid over the stream of tasks instead of observing only a short corridor to the current target. As the map becomes more complete, exploit the structure of the cost grid to improve both the current route and later routes. + +## Response format +Think concisely about the move, then give exactly one destination. Format your response as follows: + +(your thinking process goes here) +row,col diff --git a/trinity/common/workflows/connect_the_dots/grid_navigation/prompts/user.jinja2 b/trinity/common/workflows/connect_the_dots/grid_navigation/prompts/user.jinja2 new file mode 100644 index 00000000000..a1fed2165e5 --- /dev/null +++ b/trinity/common/workflows/connect_the_dots/grid_navigation/prompts/user.jinja2 @@ -0,0 +1,14 @@ +Round {{ current_round }}/{{ max_rounds }} +{% if action_feedback %} + +{{ action_feedback }} +{% endif %} + +Current position: ({{ current_position[0] }},{{ current_position[1] }}) +Target position: ({{ goal_position[0] }},{{ goal_position[1] }}) + +Observed cost grid: +{{ observation }} + +Choose one legal destination for this round. +Remember to follow every requirement in the `## Movement rules` section. diff --git a/trinity/common/workflows/connect_the_dots/grid_navigation/workflow.py b/trinity/common/workflows/connect_the_dots/grid_navigation/workflow.py new file mode 100644 index 00000000000..240a11d2f90 --- /dev/null +++ b/trinity/common/workflows/connect_the_dots/grid_navigation/workflow.py @@ -0,0 +1,218 @@ +"""CoD multi-step workflow for the persistent-cost grid-navigation task.""" + +from __future__ import annotations + +import re +from typing import TYPE_CHECKING, List, Optional, Tuple + +from trinity.common.experience import Experience +from trinity.common.workflows.connect_the_dots.base_workflow import AsyncCoDMultiStepWorkflow +from trinity.common.workflows.connect_the_dots.grid_navigation.env import ( + GridNavigationEnv, + generate_task, + get_pack_state, +) +from trinity.common.workflows.connect_the_dots.grid_navigation.prompts import ( + load_system_prompt, + load_user_prompt, +) +from trinity.common.workflows.connect_the_dots.utils import extract_content_between_keys +from trinity.common.workflows.workflow import Task + +if TYPE_CHECKING: + from trinity.common.models.model import ModelWrapper + + +def parse_action(response: str) -> Optional[Tuple[int, int]]: + """Parse one row,column destination from the answer tags.""" + content, success = extract_content_between_keys(response, "", "") + if not success: + return None + match = re.fullmatch(r"\s*(-?\d+)\s*,\s*(-?\d+)\s*", content) + if match is None: + return None + return int(match.group(1)), int(match.group(2)) + + +class CoDGridNavigationWorkflow(AsyncCoDMultiStepWorkflow): + """Navigate a hidden persistent cost grid for an exact number of rounds.""" + + is_async: bool = True + can_reset: bool = True + + def __init__( + self, + model: ModelWrapper, + task: Task, + auxiliary_models: Optional[List] = None, + use_openai_client: bool = False, + ): + """Initialize the workflow and bind it to its pack state.""" + super().__init__( + model=model, + task=task, + auxiliary_models=auxiliary_models, + use_openai_client=use_openai_client, + ) + self.reset(task) + + def reset(self, task: Task): + """Reset task configuration and acquire the shared pack environment.""" + super().reset(task) + args = task.workflow_args + self.grid_min_size = int(args.get("grid_min_size", 12)) + self.grid_max_size = int(args.get("grid_max_size", 16)) + self.landscape_num_components = int(args.get("landscape_num_components", 6)) + self.landscape_min_scale = float(args.get("landscape_min_scale", 0.10)) + self.landscape_max_scale = float(args.get("landscape_max_scale", 0.35)) + self.min_rounds = int(args.get("min_rounds", 6)) + self.max_rounds = int(args.get("max_rounds", 10)) + self.reveal_radius = int(args.get("reveal_radius", 1)) + + self.seed = int(self.raw_task.get("seed", 42)) + self.pack_seed = int(self.raw_task.get("pack_seed", self.seed)) + self.task_idx = int(self.raw_task.get("task_idx", 0)) + + pack_state = get_pack_state( + pack_seed=self.pack_seed, + task_idx=self.task_idx, + grid_min_size=self.grid_min_size, + grid_max_size=self.grid_max_size, + landscape_num_components=self.landscape_num_components, + landscape_min_scale=self.landscape_min_scale, + landscape_max_scale=self.landscape_max_scale, + ) + navigation_task = generate_task( + task_seed=self.seed, + grid_size=pack_state.size, + min_rounds=self.min_rounds, + max_rounds=self.max_rounds, + ) + self.env = GridNavigationEnv( + pack_state=pack_state, + task=navigation_task, + reveal_radius=self.reveal_radius, + ) + + self.done = False + self.final_reward = 0.0 + self.current_step = 0 + self.action_feedback: Optional[str] = None + self.early_termination_by_format_issue = False + self.memory: List[dict] = [] + + def _build_system_prompt(self) -> str: + return self._augment_system_prompt(load_system_prompt(reveal_radius=self.reveal_radius)) + + async def run_async(self) -> List[Experience]: + """Reset task-local episode state and run all navigation rounds.""" + self.env.reset_task() + self.done = False + self.final_reward = 0.0 + self.current_step = 0 + self.action_feedback = None + self.early_termination_by_format_issue = False + + self.memory.clear() + self.memory.append({"role": "system", "content": self._build_system_prompt()}) + return await super().run_async() + + def _terminate_invalid_action(self, message: str) -> None: + self.action_feedback = f"Invalid action: {message} Game over." + self.done = True + self.early_termination_by_format_issue = True + + async def step_async(self, step_num: int) -> Tuple[bool, List[Experience]]: + """Prompt for, parse, and execute one grid move.""" + if self.done: + return False, [] + + user_content = load_user_prompt( + current_round=step_num + 1, + max_rounds=self.env.task.num_rounds, + current_position=self.env.current_position, + goal_position=self.env.task.goal, + observation=self.env.render(), + action_feedback=self.action_feedback, + ) + if self.icl_examples and step_num == 0: + user_content = ( + f"{user_content}\n\nHere are some reference examples:\n\n{self.icl_examples}" + ) + + self.memory.append({"role": "user", "content": user_content}) + if self.reply_prefix: + self.memory.append({"role": "assistant", "content": self.reply_prefix}) + + experiences = await self.model.chat_async(self.memory) + response_text = experiences[0].response_text + self.memory.append({"role": "assistant", "content": response_text}) + + sys_prompt = self.memory[0]["content"] + for exp in experiences: + exp.info["sys_prompt"] = sys_prompt + exp.info["user_prompt"] = user_content + + destination = parse_action(response_text) + if destination is None: + self._terminate_invalid_action( + "expected exactly one row,col destination." + ) + self.current_step = step_num + 1 + return False, experiences + + validation_error = self.env.validate_destination(destination) + if validation_error is not None: + self._terminate_invalid_action(validation_error) + self.current_step = step_num + 1 + return False, experiences + + result = self.env.step(destination) + path_text = " -> ".join(f"({row},{col})" for row, col in result["path"]) + self.action_feedback = ( + f"You moved to ({destination[0]},{destination[1]}) through {len(result['path'])} " + f"entered cells: {path_text}. " + f"Cumulative normalized loss: {result['normalized_loss']:.4f}." + ) + self.current_step = step_num + 1 + self.done = bool(result["done"]) + if self.done: + self.final_reward = float(result["reward"]) + + for exp in experiences: + exp.info["destination"] = destination + exp.info["normalized_loss"] = result["normalized_loss"] + + return not self.done, experiences + + def _build_trajectory(self) -> str: + trajectory = super()._build_trajectory() + final_state = ( + "Final environment state:\n" + f"Position: {self.env.current_position}\n" + f"Target: {self.env.task.goal}\n" + f"Normalized loss: {self.env.normalized_loss:.4f}\n" + f"Observed cost grid:\n{self.env.render()}" + ) + return f"{trajectory}\n\n{final_state}" if trajectory else final_state + + def _get_feedback(self) -> str: + if self.early_termination_by_format_issue: + outcome = self.action_feedback or "Invalid action." + elif self.env.current_position == self.env.task.goal: + outcome = ( + f"Success: finished round {self.env.task.num_rounds} at the target. " + f"Normalized loss: {self.env.normalized_loss:.4f}; " + f"reward: {self.final_reward:.4f}." + ) + else: + outcome = ( + f"Failed: after {self.env.task.num_rounds} rounds, position " + f"{self.env.current_position} did not equal target {self.env.task.goal}. Reward: 0." + ) + return f"{outcome}\n\nFinal observed cost grid:\n{self.env.render()}" + + @property + def max_step_num(self) -> int: + """Return the task's exact number of rounds.""" + return self.env.task.num_rounds diff --git a/trinity/common/workflows/connect_the_dots/learn2ask/__init__.py b/trinity/common/workflows/connect_the_dots/learn2ask/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/trinity/common/workflows/connect_the_dots/learn2ask/prompts/__init__.py b/trinity/common/workflows/connect_the_dots/learn2ask/prompts/__init__.py new file mode 100644 index 00000000000..fd27f28b7af --- /dev/null +++ b/trinity/common/workflows/connect_the_dots/learn2ask/prompts/__init__.py @@ -0,0 +1,67 @@ +# -*- coding: utf-8 -*- +"""Prompt management for CoD Learn2Ask workflow using Jinja2 templates. + +Two templates live here: + +* system.jinja2 — rollout system prompt. A single template covers + both training modes: the `` guideline is emitted when + train_mode != "Ra", matching the original rollout_prompt_med / + rollout_prompt_med_Ra split in + examples/learn_to_ask/workflow/prompt_learn2ask.py byte-for-byte. + +* reward_judge.jinja2 — judge system prompt fed to the auxiliary + model. Holds a single placeholder, info_truth, that Trinity fills + per-cid. Renders byte-identically to + reward_prompt_med.format(info_truth). +""" + +from pathlib import Path + +from jinja2 import Environment, FileSystemLoader + + +PROMPTS_DIR = Path(__file__).parent + + +def get_jinja_env() -> Environment: + """Get Jinja2 environment with template loader.""" + return Environment( + loader=FileSystemLoader(PROMPTS_DIR), + trim_blocks=True, + lstrip_blocks=True, + ) + + +def load_system_prompt(train_mode: str = "Ra+Rs", **kwargs) -> str: + """Render the rollout system prompt. + + Args: + train_mode: "Ra+Rs" / "Rs" / "Ra" — "Ra" drops the + instruction because the decision signal is not part of Ra's + reward. All other values emit it. + **kwargs: extra variables forwarded to the template. + """ + env = get_jinja_env() + template = env.get_template("system.jinja2") + return template.render(train_mode=train_mode, **kwargs) + + +def load_reward_judge_prompt(info_truth: str, **kwargs) -> str: + """Render the judge system prompt for the current cid. + + Args: + info_truth: comma-separated structured symptom points extracted + by data prep step 1/2 (same string the upstream Learn2Ask + reward_fn passes to reward_prompt_med.format). + **kwargs: extra variables forwarded to the template. + """ + env = get_jinja_env() + template = env.get_template("reward_judge.jinja2") + return template.render(info_truth=info_truth, **kwargs) + + +__all__ = [ + "load_system_prompt", + "load_reward_judge_prompt", + "PROMPTS_DIR", +] diff --git a/trinity/common/workflows/connect_the_dots/learn2ask/prompts/reward_judge.jinja2 b/trinity/common/workflows/connect_the_dots/learn2ask/prompts/reward_judge.jinja2 new file mode 100644 index 00000000000..0d25f9b750b --- /dev/null +++ b/trinity/common/workflows/connect_the_dots/learn2ask/prompts/reward_judge.jinja2 @@ -0,0 +1,29 @@ + +# Task +You are an evaluation assistant. The user will provide a dialogue history between a doctor and a patient. You must analyze the dialogue and evaluate the doctor's last message. + +# Grading Policy +## Format Score +- 1.0: The doctor's last message contains exactly **one question**. +- 0.5: The doctor's last message contains **two questions**. +- 0.0: The doctor's last message contains **three or more questions**. + +## Content Score +- 1.0: The question(s) **directly ask about** any item in the Reference Information. +- 0.5: The question(s) are **highly relevant** to, but not directly asking about, any item in the Reference Information. +- 0.0: The question(s) are **irrelevant** to all items in the Reference Information. + +# Reference Information +{{ info_truth }} + +# Output Format + +After reasoning, respond with exactly three bracketed blocks, in this order: + +[format_score]Insert only the format score as a float, one of 1.0, 0.5, 0.0.[/format_score] +[content_score]Insert only the content score as a float, one of 1.0, 0.5, 0.0.[/content_score] +[feedback]Insert one short abstract sentence on the response's strengths or gaps, grounded in the grading policy above. Do NOT include any specific term from the Reference Information. Keep on a single line.[/feedback] + +Important: +- The three bracketed blocks above must appear in your final answer, in order. +- Scores must be based only on the doctor's last message and the provided Reference Information. diff --git a/trinity/common/workflows/connect_the_dots/learn2ask/prompts/system.jinja2 b/trinity/common/workflows/connect_the_dots/learn2ask/prompts/system.jinja2 new file mode 100644 index 00000000000..42e9f1fa5d1 --- /dev/null +++ b/trinity/common/workflows/connect_the_dots/learn2ask/prompts/system.jinja2 @@ -0,0 +1,11 @@ + +# Task +You are a medical assistant. Your task is to understand the ongoing conversation and continue the medical inquiry in English. + +## Guidelines +- Each response must contain exactly one clear and concise medical question with 2 to 3 answer choices. +- Do not repeat any previous question. +- Your response must be a single sentence. +{% if train_mode != "Ra" %} +- If enough information has been gathered to make a medication suggestion, output only: +{% endif %} diff --git a/trinity/common/workflows/connect_the_dots/learn2ask/workflow.py b/trinity/common/workflows/connect_the_dots/learn2ask/workflow.py new file mode 100644 index 00000000000..063459afec9 --- /dev/null +++ b/trinity/common/workflows/connect_the_dots/learn2ask/workflow.py @@ -0,0 +1,369 @@ +# -*- coding: utf-8 -*- +""" +CoD (Connect-the-Dots) workflow for the Learn2Ask medical inquiry task. + +Single-turn workflow adapted from examples/learn_to_ask/workflow/ +workflow_learn2ask.py. Uses an auxiliary model as LLM judge (Qwen judge +from the Learn2Ask paper) for content/format scoring; action score is a +hard rule based on detection. On top of the original workflow +this module surfaces trajectory / feedback into exp.info so the CoD +meta-workflow can consume them for cross-task hint learning. +""" + +from __future__ import annotations + +import asyncio +import re +from typing import TYPE_CHECKING, Dict, List, Optional, Tuple + +from trinity.common.experience import Experience +from trinity.common.workflows.connect_the_dots.base_workflow import ( + AsyncCoDMultiStepWorkflow, +) +from trinity.common.workflows.connect_the_dots.learn2ask.prompts import ( + load_reward_judge_prompt, + load_system_prompt, +) +from trinity.common.workflows.workflow import ( + Task, + log_sys_user_prompts_in_exp, +) + +if TYPE_CHECKING: + import openai + + from trinity.common.models.model import ModelWrapper + + +# --- Tag parsing for judge output --- + +_TAG_PATTERN = re.compile(r"\[(\w+)\](.*?)\[/\1\]", re.DOTALL) +_THINK_PATTERN = re.compile(r".*?", re.DOTALL) + + +def _parse_tag_string(text: str) -> Dict[str, str]: + """Parse [tag]...[/tag] pairs, ignoring blocks.""" + return { + tag: value.strip() + for tag, value in _TAG_PATTERN.findall(_THINK_PATTERN.sub("", text)) + } + + +# --- Dialogue formatting helpers --- + +def _merge_dialogue(msg_list: List[dict]) -> str: + """Render a list of {role, content} messages as 'patient:' / 'doctor:' lines.""" + lines = [] + for msg in msg_list: + role = msg.get("role") + content = msg.get("content", "") + if role == "user": + lines.append(f"patient: {content}") + elif role == "assistant": + lines.append(f"doctor: {content}") + return "\n".join(lines) + + +def _merge_consecutive_same_role(messages: List[dict]) -> List[dict]: + """Collapse consecutive same-role messages by joining contents with newlines.""" + if not messages: + return messages + merged = [dict(messages[0])] + for msg in messages[1:]: + if msg.get("role") == merged[-1].get("role"): + merged[-1]["content"] = f"{merged[-1]['content']}\n{msg['content']}" + else: + merged.append(dict(msg)) + return merged + + +class CoDLearn2AskWorkflow(AsyncCoDMultiStepWorkflow): + """Learn2Ask single-turn workflow wired into the CoD meta-workflow. + + Reward stays aligned with the original Learn2Ask paper (LLM judge for + content/format, hard-rule action gating, fusion formula). The main + CoD-side additions are: + + * `set_hint / set_icl_examples / set_max_response_tokens_restraint` + to receive the CoD context between tasks. + * `exp.info["trajectory"]` and `exp.info["feedback"]` populated for + the iterative hint generator to consume. + """ + + @property + def max_step_num(self) -> int: + return 1 + + def __init__( + self, + *, + task: Task, + model: "ModelWrapper", + auxiliary_models: Optional[List["ModelWrapper"]] = None, + use_openai_client: bool = False, + ): + assert ( + auxiliary_models is not None and len(auxiliary_models) == 1 + ), "CoDLearn2AskWorkflow expects exactly one auxiliary model (judge)." + super().__init__( + task=task, + model=model, + auxiliary_models=auxiliary_models, + use_openai_client=use_openai_client, + ) + self.reset(task) + + # ------------------------------------------------------------------ + # Lifecycle: reset / CoD setters + # ------------------------------------------------------------------ + + def reset(self, task: Task): + """Reset per-task state.""" + super().reset(task) + + workflow_args = task.workflow_args or {} + self.train_mode: str = workflow_args.get("train_mode", "Ra+Rs") + self.fusion_mode: str = workflow_args.get("fusion_mode", "default") + + if isinstance(self.task_desc, list): + self.task_desc = _merge_consecutive_same_role(self.task_desc) + + # Ground truth (from predata) + self.action_truth: str = self.raw_task.get("decision_truth", "continue") + self.info_truth: str = self.raw_task.get("info_truth", "") + self.session_id = str(self.raw_task.get("session_id", "")) + self.diagn: str = self.raw_task.get("diagn", "") + + self.system_prompt: str = load_system_prompt(train_mode=self.train_mode) + + # ------------------------------------------------------------------ + # Prompt assembly + # ------------------------------------------------------------------ + + def format_messages(self) -> List[dict]: + sys_prompt = self._augment_system_prompt(self.system_prompt) + if isinstance(self.task_desc, list): + messages = [{"role": "system", "content": sys_prompt}] + self.task_desc + elif isinstance(self.task_desc, str): + messages = [ + {"role": "system", "content": sys_prompt}, + {"role": "user", "content": self.task_desc}, + ] + else: + raise ValueError( + f"task.task_desc must be a list of messages or a str, got {type(self.task_desc)}" + ) + if self.reply_prefix: + messages.append({"role": "assistant", "content": self.reply_prefix}) + return messages + + # ------------------------------------------------------------------ + # Main execution + # ------------------------------------------------------------------ + + async def run_async(self) -> List[Experience]: + messages = self.format_messages() + responses = await self.model.chat_async(messages, **self.rollout_args) + + # Score every response in parallel so concurrent judge calls overlap. + score_results = await asyncio.gather( + *(self._compute_reward(r.response_text or "") for r in responses) + ) + + history_text = self._history_as_dialogue() + task_desc_str = history_text # single-turn: the prompt dialogue is the task + + for response, (reward, component_scores, judge_feedback, judge_format_error) in zip( + responses, score_results + ): + response.reward = reward + if response.metrics is None: + response.metrics = {} + # Skip per-component scores when judge failed — they're 0/0/0 placeholders, not real signal. + if not judge_format_error: + response.metrics.update(component_scores) + + resp_text = response.response_text or "" + response.info["task_desc"] = task_desc_str + response.info["trajectory"] = ( + f"# Task instructions given to the assistant\n{self.system_prompt}\n\n" + f"# Resulting conversation\n{history_text}\ndoctor: {resp_text}" + if history_text else + f"# Task instructions given to the assistant\n{self.system_prompt}\n\n" + f"# Resulting conversation\ndoctor: {resp_text}" + ) + response.info["feedback"] = self._build_feedback( + component_scores, resp_text, judge_feedback + ) + if judge_format_error: + response.info["judge_format_error"] = True + + log_sys_user_prompts_in_exp(messages, responses) + return responses + + # ------------------------------------------------------------------ + # Reward computation (mirrors examples/learn_to_ask reward_fn) + # ------------------------------------------------------------------ + + async def _compute_reward( + self, response: str + ) -> Tuple[float, Dict[str, float], str, bool]: + action_response = "stop" if "" in response else "continue" + judge_feedback = "" + judge_format_error = False + + if self.action_truth != action_response: + action_score = format_score = content_score = 0.0 + else: + action_score = 1.0 + if self.action_truth == "continue": + score_dict = await self._llm_judge(response) + try: + format_score = float(score_dict.get("format_score")) + content_score = float(score_dict.get("content_score")) + except (TypeError, ValueError): + judge_format_error = True + format_score, content_score = 0.0, 0.0 + judge_feedback = (score_dict.get("feedback") or "").strip() + else: + content_score = 1.0 + format_score = 1.0 if response.strip() == "" else 0.0 + + final_reward = self._fuse_scores(action_score, content_score, format_score) + metrics = { + "action_score": action_score, + "content_score": content_score, + "format_score": format_score, + } + return final_reward, metrics, judge_feedback, judge_format_error + + def _fuse_scores( + self, action_score: float, content_score: float, format_score: float + ) -> float: + if self.train_mode == "Ra+Rs": + if self.fusion_mode == "sum": + return action_score + content_score + format_score + return action_score * (1 + 2 * content_score) + format_score + if self.train_mode == "Ra": + return 2 * content_score + format_score + # "Rs" + return action_score * 3 + format_score + + async def _llm_judge( + self, response: str, max_retries: int = 5 + ) -> Dict[str, str]: + """Call the auxiliary judge model with the dialogue + candidate reply. + + Matches examples/learn_to_ask/workflow/llm_reward's retry loop: + one initial attempt plus `max_retries` retries (total of + max_retries + 1 = 6 tries), sleeping attempt-count seconds + between tries. + """ + client: "openai.AsyncOpenAI" = self.auxiliary_models[0] + history = self._history_as_dialogue() + judge_user = ( + f"{history}\ndoctor: {response}\n" if history else f"doctor: {response}\n" + ) + judge_messages = [ + {"role": "system", "content": load_reward_judge_prompt(info_truth=self.info_truth)}, + {"role": "user", "content": judge_user}, + ] + + total_tries = max_retries + 1 + for attempt in range(total_tries): + try: + completion = await client.chat.completions.create( + model=client.model_path, + messages=judge_messages, + stream=False, + temperature=1.0, + top_p=0.95, + presence_penalty=0.0, + extra_body={ + "chat_template_kwargs": {"enable_thinking": True}, + "top_k": 20, + "min_p": 0.0, + "repetition_penalty": 1.0, + }, + ) + content = completion.choices[0].message.content or "" + return _parse_tag_string(content) + except Exception: + if attempt >= total_tries - 1: + return {} + await asyncio.sleep(1.0 * (attempt + 1)) + return {} + + # ------------------------------------------------------------------ + # Auxiliary artefacts for CoD hint generation + # ------------------------------------------------------------------ + + def _history_as_dialogue(self) -> str: + """Render task_desc (observed context) as patient/doctor dialogue.""" + if isinstance(self.task_desc, list): + return _merge_dialogue(self.task_desc) + if isinstance(self.task_desc, str): + return f"patient: {self.task_desc}" + return "" + + def _build_feedback( + self, component_scores: Dict[str, float], response: str, judge_feedback: str = "" + ) -> str: + """Human-readable feedback summarising the judge + action outcome.""" + action_s = component_scores["action_score"] + content_s = component_scores["content_score"] + format_s = component_scores["format_score"] + got = "stop" if "" in response else "continue" + total = self._fuse_scores(action_s, content_s, format_s) + max_reward = self._fuse_scores(1.0, 1.0, 1.0) + + parts = [ + f"Reward {total:.1f}/{max_reward:.1f}. Subscores: action={action_s:.0f}, " + f"content={content_s:.1f}, format={format_s:.1f}." + ] + + if action_s == 0: + if self.action_truth == "continue": + parts.append( + "Action: chose 'stop' but expected 'continue'. This is premature " + "termination. The patient still has unfilled symptom dimensions " + "worth probing; a focused follow-up question targeting one of them " + f"would have scored up to {max_reward:.1f}. action=0 zeros all " + "subscores. The big " + "penalty for stopping early signals that information sufficiency is " + "the primary stop criterion." + ) + else: + parts.append( + "Action: chose 'continue' and asked a question but expected 'stop'. " + "This is a redundant question. Critical symptom info has already been " + "gathered, so the right action was to emit '' for a " + f"{max_reward:.1f} reward. " + "action=0 zeros all subscores. The penalty signals that over-asking " + "after enough info is collected is as costly as a wrong question." + ) + return " ".join(parts) + + parts.append(f"Action: '{got}' matched expected. action=1.0.") + + if self.action_truth == "stop": + if format_s > 0: + parts.append("Format=1.0: response was clean ''. Content auto-1.0.") + else: + parts.append( + "Format=0.0: response had extra characters around '', " + "losing 1.0. Content auto-1.0." + ) + else: + parts.append( + f"Format={format_s:.1f}. Judge penalizes multi-question turns: " + f"1 question scores 1.0, 2 scores 0.5, 3 or more scores 0.0." + ) + parts.append( + f"Content={content_s:.1f}. Judge rates how directly the question targets " + f"an unfilled symptom dimension: direct=1.0, relevant=0.5, irrelevant=0.0." + ) + if judge_feedback: + parts.append(f"Judge note: {judge_feedback}") + + return " ".join(parts) diff --git a/trinity/common/workflows/connect_the_dots/optimalcontrol/__init__.py b/trinity/common/workflows/connect_the_dots/optimalcontrol/__init__.py new file mode 100644 index 00000000000..b209d0bca74 --- /dev/null +++ b/trinity/common/workflows/connect_the_dots/optimalcontrol/__init__.py @@ -0,0 +1,22 @@ +# -*- coding: utf-8 -*- +"""CoD (Connect-the-Dots) workflow for the optimal-control deployment environment.""" + +from trinity.common.workflows.connect_the_dots.optimalcontrol.env import ( + ActionFn, + OptimalControlEnv, +) +from trinity.common.workflows.connect_the_dots.optimalcontrol.prompts import ( + load_system_prompt, + load_user_prompt, +) +from trinity.common.workflows.connect_the_dots.optimalcontrol.workflow import ( + CoDOptimalControlWorkflow, +) + +__all__ = [ + "CoDOptimalControlWorkflow", + "OptimalControlEnv", + "ActionFn", + "load_system_prompt", + "load_user_prompt", +] diff --git a/trinity/common/workflows/connect_the_dots/optimalcontrol/env.py b/trinity/common/workflows/connect_the_dots/optimalcontrol/env.py new file mode 100644 index 00000000000..06d601a5dde --- /dev/null +++ b/trinity/common/workflows/connect_the_dots/optimalcontrol/env.py @@ -0,0 +1,516 @@ +# -*- coding: utf-8 -*- +"""Optimal-control environment with hidden linear dynamics. + +State: s_t = (x_t, v_t) +Dynamics: + v_{t+1} = a_env * v_t + b_env * u_t + epsilon_t + x_{t+1} = x_t + v_{t+1} +Control constraint: u_t in [-1, 1] +""" + +import ast +import math +import operator +import random +from typing import Any, Callable, Dict, List, Tuple, TypedDict +from xml.sax.saxutils import escape + +ActionFn = Callable[[float, float, int, float, float], float] + + +class ResolvedOptimalControlTask(TypedDict): + a_env: float + b_env: float + min_abs_b_env: float #!!! + x0: float + v0: float + x_target: float + v_target: float + horizon: int + control_penalty_coef: float + enable_process_noise: bool #!!! + process_noise_std: float #!!! + task_seed: int #!!! + + +def _clip(value: float, lower: float, upper: float) -> float: + """Clip value to [lower, upper].""" + return max(lower, min(value, upper)) + + +def format_reward(value: float) -> str: + """Keep three decimals unless a positive reward would appear to be zero.""" + rounded = f"{value:.3f}" + if value > 0.0 and rounded == "0.000": + return f"{value:.3e} (nonzero; below three-decimal precision)" + return rounded + + +_VARIABLE_NAMES = { + "x", + "v", + "t", + "x_target", + "v_target", + "horizon", + "remaining_steps", +} +_BINARY_OPERATORS = { + ast.Add: operator.add, + ast.Sub: operator.sub, + ast.Mult: operator.mul, + ast.Div: operator.truediv, + ast.Pow: operator.pow, +} +_UNARY_OPERATORS = { + ast.UAdd: operator.pos, + ast.USub: operator.neg, + ast.Not: operator.not_, +} +_COMPARISON_OPERATORS = { + ast.Eq: operator.eq, + ast.NotEq: operator.ne, + ast.Lt: operator.lt, + ast.LtE: operator.le, + ast.Gt: operator.gt, + ast.GtE: operator.ge, +} +_FUNCTION_COMPARISONS = { + "lt": operator.lt, + "le": operator.le, + "gt": operator.gt, + "ge": operator.ge, + "eq": operator.eq, + "ne": operator.ne, +} +_FUNCTION_ARITY = { + "abs": (1, 1), + "min": (1, None), + "max": (1, None), + "pow": (2, 2), + "clip": (3, 3), + "ifelse": (3, 3), + "lt": (2, 2), + "le": (2, 2), + "gt": (2, 2), + "ge": (2, 2), + "eq": (2, 2), + "ne": (2, 2), +} + + +def _error_detail(error: Exception) -> str: + """Return a concise error description suitable for model feedback.""" + return f"{type(error).__name__}: {error}"[:500] + + +def _validate_expression_node(node: ast.AST) -> None: + """Validate the supported scalar expression language.""" + if isinstance(node, ast.Expression): + _validate_expression_node(node.body) + return + if isinstance(node, ast.Constant): + if isinstance(node.value, bool) or not isinstance(node.value, (int, float)): + raise ValueError("only numeric constants are allowed") + try: + finite = math.isfinite(float(node.value)) + except OverflowError as error: + raise ValueError("numeric constants must be finite") from error + if not finite: + raise ValueError("numeric constants must be finite") + return + if isinstance(node, ast.Name): + if node.id not in _VARIABLE_NAMES: + raise ValueError(f"unknown variable: {node.id}") + return + if isinstance(node, ast.BinOp): + if type(node.op) not in _BINARY_OPERATORS: + raise ValueError("unsupported binary operator") + _validate_expression_node(node.left) + _validate_expression_node(node.right) + return + if isinstance(node, ast.UnaryOp): + if type(node.op) not in _UNARY_OPERATORS: + raise ValueError("unsupported unary operator") + _validate_expression_node(node.operand) + return + if isinstance(node, ast.Call): + if not isinstance(node.func, ast.Name) or node.func.id not in _FUNCTION_ARITY: + raise ValueError("unsupported function call") + if node.keywords: + raise ValueError("keyword arguments are not allowed") + minimum, maximum = _FUNCTION_ARITY[node.func.id] + if len(node.args) < minimum or (maximum is not None and len(node.args) > maximum): + raise ValueError(f"invalid argument count for {node.func.id}") + for argument in node.args: + _validate_expression_node(argument) + return + if isinstance(node, ast.IfExp): + _validate_expression_node(node.test) + _validate_expression_node(node.body) + _validate_expression_node(node.orelse) + return + if isinstance(node, ast.Compare): + if any(type(comparator) not in _COMPARISON_OPERATORS for comparator in node.ops): + raise ValueError("unsupported comparison operator") + _validate_expression_node(node.left) + for comparator in node.comparators: + _validate_expression_node(comparator) + return + if isinstance(node, ast.BoolOp): + if not isinstance(node.op, (ast.And, ast.Or)): + raise ValueError("unsupported boolean operator") + for value in node.values: + _validate_expression_node(value) + return + raise ValueError(f"unsupported expression syntax: {type(node).__name__}") + + +def validate_action_expression(expression: object) -> str: + """Parse and validate one controller expression.""" + if not isinstance(expression, str): + raise ValueError("action_expression must be text") + expression = expression.strip() + if not expression: + raise ValueError("action_expression cannot be empty") + try: + tree = ast.parse(expression, mode="eval") + _validate_expression_node(tree) + except SyntaxError as error: + raise ValueError(f"invalid action_expression syntax: {error.msg}") from error + except RecursionError as error: + raise ValueError("action_expression is too deeply nested") from error + return expression + + +def _evaluate_expression(node: ast.AST, variables: Dict[str, float]) -> Any: + """Evaluate one previously validated expression node.""" + if isinstance(node, ast.Constant): + return float(node.value) + if isinstance(node, ast.Name): + return variables[node.id] + if isinstance(node, ast.BinOp): + return _BINARY_OPERATORS[type(node.op)]( + _evaluate_expression(node.left, variables), + _evaluate_expression(node.right, variables), + ) + if isinstance(node, ast.UnaryOp): + return _UNARY_OPERATORS[type(node.op)](_evaluate_expression(node.operand, variables)) + if isinstance(node, ast.Call): + function_name = node.func.id + if function_name == "ifelse": + condition = _evaluate_expression(node.args[0], variables) + branch = node.args[1] if condition else node.args[2] + return _evaluate_expression(branch, variables) + values = [_evaluate_expression(argument, variables) for argument in node.args] + if function_name == "abs": + return abs(values[0]) + if function_name == "min": + return min(values) + if function_name == "max": + return max(values) + if function_name == "pow": + return pow(values[0], values[1]) + if function_name == "clip": + return _clip(values[0], values[1], values[2]) + return _FUNCTION_COMPARISONS[function_name](values[0], values[1]) + if isinstance(node, ast.IfExp): + branch = node.body if _evaluate_expression(node.test, variables) else node.orelse + return _evaluate_expression(branch, variables) + if isinstance(node, ast.Compare): + left = _evaluate_expression(node.left, variables) + for operation, comparator in zip(node.ops, node.comparators): + right = _evaluate_expression(comparator, variables) + if not _COMPARISON_OPERATORS[type(operation)](left, right): + return False + left = right + return True + if isinstance(node, ast.BoolOp): + if isinstance(node.op, ast.And): + result = _evaluate_expression(node.values[0], variables) + for value in node.values[1:]: + if not result: + return result + result = _evaluate_expression(value, variables) + return result + result = _evaluate_expression(node.values[0], variables) + for value in node.values[1:]: + if result: + return result + result = _evaluate_expression(value, variables) + return result + raise TypeError(f"unsupported expression node: {type(node).__name__}") + + +def compile_action_expression(expression: str, horizon: int) -> Tuple[ActionFn | None, str | None]: + """Compile one validated scalar expression into a feedback action.""" + try: + expression = validate_action_expression(expression) + tree = ast.parse(expression, mode="eval") + except ValueError as error: + return None, str(error) + + def action(x, v, t, x_target, v_target): + variables = { + "x": float(x), + "v": float(v), + "t": float(t), + "x_target": float(x_target), + "v_target": float(v_target), + "horizon": float(horizon), + "remaining_steps": float(max(1, horizon - t)), + } + return _evaluate_expression(tree.body, variables) + + return action, None + + +def serialize_action_expression_xml(expression: str) -> str: + """Serialize a validated expression as canonical controller XML.""" + expression = validate_action_expression(expression) + return ( + "\n" + " \n" + f" {escape(expression)}\n" + " \n" + "" + ) + + +class OptimalControlEnv: + """One-dimensional optimal-control environment used by CoD-Deploy demos.""" + + @staticmethod #!!! + def _sample_b_env( #!!! + rng: random.Random, lower: float, upper: float, min_abs: float #!!! + ) -> float: #!!! + """Sample uniformly while excluding (-min_abs, min_abs).""" #!!! + if lower > upper: #!!! + raise ValueError("b_env_range lower bound must not exceed upper bound") #!!! + if min_abs < 0.0: #!!! + raise ValueError("min_abs_b_env must be non-negative") #!!! + if min_abs == 0.0: #!!! + return rng.uniform(lower, upper) #!!! + left_upper = min(upper, -min_abs) #!!! + right_lower = max(lower, min_abs) #!!! + left_length = max(0.0, left_upper - lower) #!!! + right_length = max(0.0, upper - right_lower) #!!! + total_length = left_length + right_length #!!! + if total_length == 0.0: #!!! + if lower <= -min_abs <= upper: #!!! + return -min_abs #!!! + if lower <= min_abs <= upper: #!!! + return min_abs #!!! + raise ValueError("b_env_range contains no value satisfying min_abs_b_env") #!!! + offset = rng.random() * total_length #!!! + if offset < left_length: #!!! + return lower + offset #!!! + return right_lower + (offset - left_length) #!!! + + @staticmethod + def resolve_task( + raw_task: Dict[str, Any], + workflow_args: Dict[str, Any], + ) -> ResolvedOptimalControlTask: + """Resolve one task against its runtime CoD pack environment.""" + environment_seed = int(raw_task.get("pack_seed", raw_task["seed"])) + rng = random.Random(environment_seed) + a_env = rng.uniform(*workflow_args["a_env_range"]) + b_lower, b_upper = map(float, workflow_args["b_env_range"]) #!!! + min_abs_b_env = float(workflow_args.get("min_abs_b_env", 0.0)) #!!! + b_env = OptimalControlEnv._sample_b_env( #!!! + rng, b_lower, b_upper, min_abs_b_env #!!! + ) #!!! + x0 = float(raw_task["x0"]) + v0 = float(raw_task["v0"]) + x_target = float(raw_task["x_target"]) + v_target = float(raw_task["v_target"]) + horizon_value = raw_task.get("horizon") #!!! + if horizon_value is None: #!!! + horizon_value = raw_task["max_horizon"] #!!! + horizon = int(horizon_value) #!!! + control_penalty_coef = float(raw_task["control_penalty_coef"]) + enable_process_noise = bool( #!!! + workflow_args.get("enable_process_noise", False) #!!! + ) #!!! + process_noise_std = float(workflow_args.get("process_noise_std", 0.0)) #!!! + if process_noise_std < 0.0: #!!! + raise ValueError("process_noise_std must be non-negative") #!!! + task_seed = int(raw_task["seed"]) #!!! + + return { + "a_env": a_env, + "b_env": b_env, + "min_abs_b_env": min_abs_b_env, #!!! + "x0": x0, + "v0": v0, + "x_target": x_target, + "v_target": v_target, + "horizon": horizon, + "control_penalty_coef": control_penalty_coef, + "enable_process_noise": enable_process_noise, #!!! + "process_noise_std": process_noise_std, #!!! + "task_seed": task_seed, #!!! + } + + def __init__( + self, + a_env: float, + b_env: float, + x0: float, + v0: float, + x_target: float, + v_target: float, + horizon: int = 8, + control_penalty_coef: float = 0.03, + enable_process_noise: bool = False, #!!! + process_noise_std: float = 0.0, #!!! + task_seed: int = 0, #!!! + ): + self.a_env = float(a_env) + self.b_env = float(b_env) + self.x0 = float(x0) + self.v0 = float(v0) + self.x_target = float(x_target) + self.v_target = float(v_target) + self.horizon = int(horizon) + self.control_penalty_coef = float(control_penalty_coef) + self.enable_process_noise = bool(enable_process_noise) #!!! + self.process_noise_std = float(process_noise_std) #!!! + if self.process_noise_std < 0.0: #!!! + raise ValueError("process_noise_std must be non-negative") #!!! + self.task_seed = int(task_seed) #!!! + + def reset(self) -> Tuple[float, float]: + """Reset to initial state.""" + return self.x0, self.v0 + + def step( #!!! + self, x: float, v: float, u: float, process_noise: float = 0.0 #!!! + ) -> Tuple[float, float]: #!!! + """Apply one control and return the next state.""" + u = float(_clip(u, -1.0, 1.0)) + v_next = self.a_env * v + self.b_env * u + float(process_noise) #!!! + x_next = x + v_next + return x_next, v_next + + def rollout(self, action_fn: ActionFn) -> Dict[str, Any]: + """Roll out the policy for the full horizon. + + Args: + action_fn: callable action(x, v, t, x_target, v_target) -> u_t. + + Returns: + dict with rollout states, controls, terminal metrics, and error status. #!!! + """ + x, v = self.reset() + xs = [x] + vs = [v] + us: List[float] = [] + process_noises: List[float] = [] #!!! + noise_rng = random.Random(self.task_seed ^ 0x5DEECE66D) #!!! + + for t in range(self.horizon): + try: + raw_u = action_fn(x, v, t, self.x_target, self.v_target) #!!! + if isinstance(raw_u, bool) or not isinstance(raw_u, (int, float)): #!!! + raise TypeError("action must return an int or float scalar") #!!! + u = float(raw_u) #!!! + if not math.isfinite(u): + raise ValueError("action must return a finite scalar") + u = _clip(u, -1.0, 1.0) + except Exception as error: + raise RuntimeError(f"action failed at t={t}: {_error_detail(error)}") from error + us.append(u) + process_noise = ( #!!! + noise_rng.gauss(0.0, self.process_noise_std) #!!! + if self.enable_process_noise and self.process_noise_std > 0.0 #!!! + else 0.0 #!!! + ) #!!! + process_noises.append(process_noise) #!!! + x, v = self.step(x, v, u, process_noise=process_noise) #!!! + xs.append(x) + vs.append(v) + + loss = self.compute_loss(xs[-1], vs[-1], us) #!!! + reward = self.compute_reward(loss) #!!! + process_noise_rms = ( #!!! + math.sqrt( #!!! + sum(noise * noise for noise in process_noises) / len(process_noises) #!!! + ) #!!! + if process_noises #!!! + else 0.0 #!!! + ) #!!! + + return { + "xs": xs, + "vs": vs, + "us": us, + "process_noises": process_noises, #!!! + "process_noise_rms": process_noise_rms, #!!! + "x_final": xs[-1], + "v_final": vs[-1], + "loss": loss, + "reward": reward, + } + + def rollout_action_expression( + self, + expression: str, + ) -> Tuple[Dict[str, Any] | None, str, str | None]: + """Compile and fully evaluate one feedback expression.""" + action_fn, validation_error = compile_action_expression(expression, self.horizon) + if action_fn is None: + return None, "format_error", validation_error + try: + return self.rollout(action_fn), "ok", None + except Exception as error: + return None, "action_error", _error_detail(error) + + def compute_loss( #!!! + self, x_final: float, v_final: float, us: List[float] #!!! + ) -> float: #!!! + """Compute terminal-state and control-effort loss.""" #!!! + pos_error = x_final - self.x_target #!!! + vel_error = v_final - self.v_target #!!! + control_cost = self.control_penalty_coef * sum(u * u for u in us) #!!! + return pos_error * pos_error + 2.0 * vel_error * vel_error + control_cost #!!! + + def compute_reward(self, loss: float) -> float: + """Convert loss to reward.""" + return 1.0 / (1.0 + loss) + + def render_trajectory(self, result: Dict[str, Any]) -> str: + """Render rollout result as a table aligned by time step. + + The control column u_t shows the control applied *at* step t while the + system is in state (x_t, v_t). At the terminal step T no further control + is available, so it is marked as ``--``. + """ + xs = result["xs"] + vs = result["vs"] + us = result["us"] + lines = [ + f"{'t':>3} {'x_t':>10} {'v_t':>10} {'u_t':>10}", + f"{0:>3} {xs[0]:>10.3f} {vs[0]:>10.3f} {us[0]:>10.3f}", + ] + for t in range(1, self.horizon): + u_str = f"{us[t]:>10.3f}" + lines.append(f"{t:>3} {xs[t]:>10.3f} {vs[t]:>10.3f} {u_str}") + lines.append( + f"{self.horizon:>3} {xs[self.horizon]:>10.3f} {vs[self.horizon]:>10.3f} {'--':>10}" + ) + return "\n".join(lines) + + def render_summary(self, result: Dict[str, Any]) -> str: + """Render the terminal state, target, loss, and reward.""" + reward_text = format_reward(float(result["reward"])) + lines = [ + "Terminal state:", + f"x_{self.horizon} = {result['x_final']:.3f}, v_{self.horizon} = {result['v_final']:.3f}.", + "Target:", + f"x_target = {self.x_target:.3f}, v_target = {self.v_target:.3f}.", + "Loss and reward:", + f"L = {result['loss']:.3f}, R = {reward_text}.", + ] + return "\n".join(lines) diff --git a/trinity/common/workflows/connect_the_dots/optimalcontrol/prompts/__init__.py b/trinity/common/workflows/connect_the_dots/optimalcontrol/prompts/__init__.py new file mode 100644 index 00000000000..1ecae727285 --- /dev/null +++ b/trinity/common/workflows/connect_the_dots/optimalcontrol/prompts/__init__.py @@ -0,0 +1,82 @@ +# -*- coding: utf-8 -*- +"""Prompt management for CoD Optimal Control workflow using Jinja2 templates.""" + +from pathlib import Path + +from jinja2 import Environment, FileSystemLoader + +PROMPTS_DIR = Path(__file__).parent + + +def get_jinja_env() -> Environment: + """Get Jinja2 environment with template loader.""" + return Environment( + loader=FileSystemLoader(PROMPTS_DIR), + trim_blocks=True, + lstrip_blocks=True, + ) + + +def load_system_prompt( + control_penalty_coef: float = 0.03, + enable_process_noise: bool = False, #!!! + process_noise_std: float = 0.0, #!!! + known_b_sign=None, + **kwargs, +) -> str: + """Load and render the system prompt template. + + Args: + control_penalty_coef: Coefficient used in the control-effort term of + the loss function shown to the agent. + known_b_sign: Disclosed control direction, or ``None`` when hidden. + **kwargs: Additional template variables. + + Returns: + Rendered system prompt string. + """ + env = get_jinja_env() + template = env.get_template("system.jinja2") + return template.render( + control_penalty_coef=control_penalty_coef, + enable_process_noise=enable_process_noise, #!!! + process_noise_std=process_noise_std, #!!! + known_b_sign=known_b_sign, + **kwargs, + ) + + +def load_user_prompt( + x0: float, + v0: float, + x_target: float, + v_target: float, + horizon: int, + **kwargs, +) -> str: + """Load and render the user prompt template for one optimal-control task. + + Args: + x0: Initial position. + v0: Initial velocity. + x_target: Target position. + v_target: Target velocity. + horizon: Rollout horizon T. + **kwargs: Additional template variables. + + Returns: + Rendered user prompt string. + """ + env = get_jinja_env() + template = env.get_template("user.jinja2") + return template.render( + x0=x0, + v0=v0, + x_target=x_target, + v_target=v_target, + horizon=horizon, + **kwargs, + ) + + +__all__ = ["load_system_prompt", "load_user_prompt", "PROMPTS_DIR"] diff --git a/trinity/common/workflows/connect_the_dots/optimalcontrol/prompts/system.jinja2 b/trinity/common/workflows/connect_the_dots/optimalcontrol/prompts/system.jinja2 new file mode 100644 index 00000000000..0e91dc10c97 --- /dev/null +++ b/trinity/common/workflows/connect_the_dots/optimalcontrol/prompts/system.jinja2 @@ -0,0 +1,102 @@ +You are solving a one-dimensional optimal-control problem by designing a scalar feedback expression. + +## Environment Dynamics + +The state is `(x_t, v_t)`. At step `t`, the environment evaluates your expression using the current state and applies the resulting scalar control `u_t`. Controls are applied for `t = 0, 1, ..., T - 1`, and the terminal state `(x_T, v_T)` is evaluated after the final transition. + +```text +{% if enable_process_noise %} +v_(t+1) = a_env * v_t + b_env * u_t + epsilon_t +{% else %} +v_(t+1) = a_env * v_t + b_env * u_t +{% endif %} +x_(t+1) = x_t + v_(t+1) +``` + +The numerical values of `a_env` and `b_env` are unknown and hidden from you. +{% if known_b_sign == "positive" %} +You are told that `b_env` is strictly positive: positive `u_t` increases the control contribution to the next velocity. Its magnitude remains unknown. +{% elif known_b_sign == "negative" %} +You are told that `b_env` is strictly negative: positive `u_t` decreases the control contribution to the next velocity. Its magnitude remains unknown. +{% else %} +The control coefficient `b_env` may be positive or negative, so infer its sign from supplied hints or earlier-controller evidence rather than assuming a direction. +{% endif %} +The environment clips every evaluated control to `[-1, 1]`. +{% if enable_process_noise %} + +The velocity transition also includes an unpredictable disturbance: + +```text +epsilon_t ~ Normal(0, {{ process_noise_std }}^2) +``` + +The noise seed and future values are unavailable. Use current-state feedback to correct disturbances that have already affected the trajectory. +{% endif %} + +## Shared Environment Evidence + +Related tasks in the same pack share the hidden `a_env` and `b_env`. Supplied hints and earlier-controller evidence may therefore help infer the shared dynamics. Earlier tasks can have different initial states, targets, and horizons, so reuse their control law only after adapting it to the current task. + +## Objective + +```text +L = (x_T - x_target)^2 + + 2 * (v_T - v_target)^2 + + {{ control_penalty_coef }} * sum_t u_t^2 +R = 1 / (1 + L) +``` + +Maximize `R` by reaching the target position with the requested terminal velocity while avoiding unnecessary control effort. + +## Action Expression + +Submit one scalar expression. The environment reevaluates it at every step with these variables: + +- `x`: current position `x_t` +- `v`: current velocity `v_t` +- `t`: current integer step index +- `x_target`: target position +- `v_target`: target velocity +- `horizon`: rollout horizon `T` +- `remaining_steps`: `max(1, horizon - t)` + +Supported syntax: + +- finite numeric constants, including decimal and scientific notation +- arithmetic operators: `+`, `-`, `*`, `/`, `**` +- comparisons: `<`, `<=`, `>`, `>=`, `==`, `!=` +- Boolean operators: `and`, `or`, `not` +- conditional expression: `value_if_true if condition else value_if_false` +- numeric functions: `abs(value)`, `min(values...)`, `max(values...)`, `pow(base, exponent)`, and `clip(value, lower, upper)` +- XML-safe conditional functions: `ifelse(condition, value_if_true, value_if_false)` with `lt(a,b)`, `le(a,b)`, `gt(a,b)`, `ge(a,b)`, `eq(a,b)`, or `ne(a,b)` + +Normal arithmetic precedence applies. The final result must be one finite numeric scalar. The environment performs the final `[-1, 1]` clipping, so an outer `clip` is optional. Assignment and local variables are unsupported: write `0.3*(x_target-x)` directly, not `dx = x_target-x`. Statements, function definitions, imports, loops, comprehensions, containers, indexing, attributes, and all other function calls are also unsupported. + +A fixed signed-PD expression can be written as: + +```text +0.25 * (x_target - x) + 0.32 * (v_target - v) +``` + +A time-varying expression can be written without XML escaping by using the XML-safe conditional functions: + +```text +0.25 * (x_target - x) + ifelse(lt(t, horizon - 2), 0.5, 1.0) * (v_target - v) +``` + +Use `ifelse` with `lt/le/gt/ge/eq/ne` in the final XML whenever possible. Raw `<` and `<=` are also valid expression operators, but XML requires them to be escaped as `<` and `<=`. PDE equations normally need no escaping because they do not contain comparison operators. + + + + 0.25 * (x_target - x) + ifelse(lt(t, horizon - 2), 0.5, 1.0) * (v_target - v) + + + +## Response Contract + +Reason briefly about the inferred dynamics and feedback law, then end the response with exactly one `...` submission. + +- The final `` block must be unique and must end the response. +- `` must contain only expression text and no nested XML elements. +- Do not output Python functions, assignments, JSON, Markdown code fences, multiple candidate expressions, additional XML actions, or text after ``. +- Do not expose or assume numerical values for hidden dynamics without evidence. Use the disclosed sign of `b_env` when one is provided. diff --git a/trinity/common/workflows/connect_the_dots/optimalcontrol/prompts/user.jinja2 b/trinity/common/workflows/connect_the_dots/optimalcontrol/prompts/user.jinja2 new file mode 100644 index 00000000000..4bdbd9b0397 --- /dev/null +++ b/trinity/common/workflows/connect_the_dots/optimalcontrol/prompts/user.jinja2 @@ -0,0 +1,5 @@ +Initial state: x_0 = {{ "%.3f"|format(x0) }}, v_0 = {{ "%.3f"|format(v0) }}. +Target state: x_target = {{ "%.3f"|format(x_target) }}, v_target = {{ "%.3f"|format(v_target) }}. +The rollout has {{ horizon }} steps (t = 0, 1, ..., {{ horizon - 1 }}). + +Think briefly about the feedback law, then end with exactly one `...` submission. diff --git a/trinity/common/workflows/connect_the_dots/optimalcontrol/workflow.py b/trinity/common/workflows/connect_the_dots/optimalcontrol/workflow.py new file mode 100644 index 00000000000..243fc37d11c --- /dev/null +++ b/trinity/common/workflows/connect_the_dots/optimalcontrol/workflow.py @@ -0,0 +1,673 @@ +# -*- coding: utf-8 -*- +"""CoD workflow for the one-dimensional optimal-control deployment demo. + +Each task is a single-turn controller-design problem: the model submits one +scalar action expression, which is evaluated as closed-loop feedback under +hidden linear dynamics for a fixed horizon. + +The workflow inherits from ``AsyncCoDMultiStepWorkflow`` to reuse standard CoD +utilities (system-prompt augmentation/trajectory stripping, ICL-example +stripping, and trajectory formatting). +""" + +from __future__ import annotations + +import asyncio +import datetime +import json +import os +from typing import TYPE_CHECKING, List, Optional, Tuple + +import openai +import torch + +from trinity.common.experience import Experience +from trinity.common.workflows.connect_the_dots.base_workflow import ( + AsyncCoDMultiStepWorkflow, +) +from trinity.common.workflows.connect_the_dots.optimalcontrol.env import ( + OptimalControlEnv, + format_reward, + serialize_action_expression_xml, + validate_action_expression, +) +from trinity.common.workflows.connect_the_dots.optimalcontrol.prompts import ( + load_system_prompt, + load_user_prompt, +) +from trinity.common.workflows.connect_the_dots.utils import parse_xml_answer +from trinity.common.workflows.workflow import Task + +if TYPE_CHECKING: + from trinity.common.models.model import ModelWrapper + + +_PACK_CONTROL_POLICIES: dict[tuple[str, int], dict[str, object]] = {} #!!! +PARSE_FAILURE_REWARD = -0.1 + + +def parse_action_expression_submission(response: str) -> tuple[Optional[str], str]: + """Parse one terminal XML action expression.""" + if "```" in response: + return None, "python_or_markdown_code_is_not_allowed" + payload, parse_error = parse_xml_answer(response) + if payload is None: + return None, parse_error + answer_end = response.rfind("") + len("") + if response[answer_end:].strip(): + return None, "answer_must_end_the_response" + if payload.get("action") != "action_expression": + return None, "expected_action_expression" + try: + return validate_action_expression(payload.get("args")), "" + except ValueError as error: + return None, str(error) + + +class CoDOptimalControlWorkflow(AsyncCoDMultiStepWorkflow): + """Single-turn action-expression workflow for optimal control. + + The model emits one scalar feedback expression; the workflow validates it, + evaluates it under hidden dynamics, and assigns the rollout reward. + """ + + is_async: bool = True + can_reset: bool = True + can_repeat: bool = False + + def __init__( + self, + model: ModelWrapper, + task: Task, + auxiliary_models: Optional[List] = None, + ): + # The sync OpenAI client from the base class is not needed; in + # external mode an async client is created lazily in `_chat`. + super().__init__( + task=task, + model=model, + auxiliary_models=auxiliary_models, + use_openai_client=False, + ) + self.reset(task) + + def reset(self, task: Task): + """Reset task-specific configuration and environment state.""" + super().reset(task) + + self.workflow_args = task.workflow_args if hasattr(task, "workflow_args") else {} + self.trajectory_dump_dir = self.workflow_args.get("trajectory_dump_dir", None) + self.task_idx = int(self.raw_task["task_idx"]) + self.include_previous_control_policy = bool( + self.workflow_args.get("include_previous_control_policy", False) + ) + self.previous_control_policy: Optional[str] = None + self.submitted_control_policy: Optional[str] = None + self.previous_control_policy_evaluation: Optional[dict[str, object]] = None #!!! + self.submitted_control_policy_evaluation: Optional[dict[str, object]] = None #!!! + + task_config = OptimalControlEnv.resolve_task(self.raw_task, self.workflow_args) + self.a_env = task_config["a_env"] + self.b_env = task_config["b_env"] + self.min_abs_b_env = task_config["min_abs_b_env"] #!!! + self.x0 = task_config["x0"] + self.v0 = task_config["v0"] + self.x_target = task_config["x_target"] + self.v_target = task_config["v_target"] + self.horizon = task_config["horizon"] + self.control_penalty_coef = task_config["control_penalty_coef"] + self.enable_process_noise = task_config["enable_process_noise"] #!!! + self.process_noise_std = task_config["process_noise_std"] #!!! + self.task_seed = task_config["task_seed"] #!!! + + self.env = OptimalControlEnv( + a_env=self.a_env, + b_env=self.b_env, + x0=self.x0, + v0=self.v0, + x_target=self.x_target, + v_target=self.v_target, + horizon=self.horizon, + control_penalty_coef=self.control_penalty_coef, + enable_process_noise=self.enable_process_noise, #!!! + process_noise_std=self.process_noise_std, #!!! + task_seed=self.task_seed, #!!! + ) + + # State for the current task + self.memory: List[dict] = [] + + @property + def max_step_num(self) -> int: + """This is a single-turn workflow: one environment step per task.""" + return 1 + + def _pack_policy_key(self) -> tuple[str, int]: + pack_seed = int(self.raw_task.get("pack_seed", self.raw_task.get("seed", 42))) + return str(self.task.batch_id), pack_seed + + def _build_system_prompt(self) -> str: + """Build system prompt, appending CoD hint and token-limit notes.""" + b_lower, b_upper = map(float, self.workflow_args["b_env_range"]) + known_b_sign = "positive" if b_lower > 0.0 else "negative" if b_upper < 0.0 else None + sys_prompt = load_system_prompt( + control_penalty_coef=self.control_penalty_coef, + enable_process_noise=self.enable_process_noise, #!!! + process_noise_std=self.process_noise_std, #!!! + known_b_sign=known_b_sign, + ) + return self._augment_system_prompt(sys_prompt) + + def _build_user_prompt(self) -> str: + """Build the task-specific user prompt.""" + user_prompt = load_user_prompt( + x0=self.x0, + v0=self.v0, + x_target=self.x_target, + v_target=self.v_target, + horizon=self.horizon, + ) + if self.icl_examples: + user_prompt = ( + f"{user_prompt}\n\nHere are some reference examples:\n\n{self.icl_examples}" + ) + if self.previous_control_policy: #!!! + evaluation = self.previous_control_policy_evaluation or {} + executed = bool(evaluation.get("action_execution_success", False)) + reward = float(evaluation.get("reward", 0.0)) + reward_text = format_reward(reward) + position_error = evaluation.get("signed_position_error") + velocity_error = evaluation.get("signed_velocity_error") + behavior = ( + "The expression completed the earlier rollout." + if executed + else "The expression did not complete the earlier rollout." + ) + position_text = "unknown" if position_error is None else f"{float(position_error):.3f}" + velocity_text = "unknown" if velocity_error is None else f"{float(velocity_error):.3f}" + user_prompt += ( #!!! + "\n\n## Controller evidence from an earlier task\n\n" + "The expression and measurements below come from an earlier, different " + "task in the same hidden environment. They are not rollout results for " + "the current task. Use them only as evidence about the shared hidden " + "dynamics and controller behavior.\n\n" + f"- Earlier initial state: x0={evaluation.get('x0')}, " + f"v0={evaluation.get('v0')}\n" + f"- Earlier target: x_target={evaluation.get('x_target')}, " + f"v_target={evaluation.get('v_target')}\n" + f"- Earlier horizon: {evaluation.get('horizon')}\n" + f"- Completed that earlier rollout: {'yes' if executed else 'no'}\n" + f"- Reward on that earlier task: {reward_text}\n" + f"- Signed terminal position residual x_T - x_target: {position_text}\n" + f"- Signed terminal velocity residual v_T - v_target: {velocity_text}\n" + f"- Main behavior: {behavior}\n\n" + "Action expression used on that earlier task:\n\n" + f"{serialize_action_expression_xml(self.previous_control_policy)}\n\n" #!!! + "Adapt the expression to the current initial state, target, and horizon, " + "or submit a new one if it performed poorly." + ) #!!! + return user_prompt + + def _build_feedback( + self, + result: Optional[dict], + evaluation_status: str, + error_detail: Optional[str] = None, + response_truncated: bool = False, + ) -> str: + """Build environment feedback string for this rollout.""" + if evaluation_status == "format_error" and response_truncated: + return ( + "Invalid response: generation reached the maximum response length before a " + "valid final action-expression XML submission was completed. The environment " + f"was not executed and reward was set to {PARSE_FAILURE_REWARD}." + ) + if evaluation_status == "format_error": + detail = error_detail or "invalid action-expression XML" + return ( + "Invalid controller submission: expected exactly one terminal " + "... block " + f"({detail}). The environment was not executed and reward was set to " + f"{PARSE_FAILURE_REWARD}." + ) + if evaluation_status == "action_error": + detail = error_detail or "expression evaluation failed." + return ( + f"Invalid policy: {detail}. Controller evaluation stopped when the expression " + "failed; no complete trajectory or loss was returned, and reward was forced " + "to 0." + ) + if result is None: + return "Invalid policy: controller evaluation failed. Reward was forced to 0." + + return "\n\n".join( + [ + "The environment executed the submitted policy under the fixed hidden dynamics.", + self.env.render_trajectory(result), + self.env.render_summary(result), + ] + ) + + def _build_rollout_trajectory( + self, + system_prompt: str, + user_prompt: str, + response_text: str, + action_expression: Optional[str], + ) -> str: + """Build controller evidence for CoD hint generation.""" + if action_expression is not None: + policy_label = "Submitted action expression" + policy = serialize_action_expression_xml(action_expression) + else: + policy_label = "Raw agent response; no valid action expression was extracted" + policy = response_text + return "\n\n".join( + [ + f"System:\n{self._strip_system_prompt(system_prompt)}", + f"Task:\n{self._strip_icl_examples(user_prompt)}", + f"{policy_label}:\n{policy}", + ] + ) + + def _dump_trajectory( + self, + run_id: int, + system_prompt: str, + user_prompt: str, + response_text: str, + action_expression: Optional[str], + result: Optional[dict], + format_error: bool, + action_error: bool, + response_token_count: int, + max_response_tokens: Optional[int], + response_hit_max_tokens: bool, + response_truncated: bool, + error_detail: Optional[str], + feedback: str, + ) -> None: + """Persist one task trajectory to a JSONL file for offline analysis.""" + if not self.trajectory_dump_dir: + return + + os.makedirs(self.trajectory_dump_dir, exist_ok=True) + dump_path = os.path.join( + self.trajectory_dump_dir, + "trajectories.jsonl", + ) + + # This workflow only handles the actual solve_task step; the CoD + # hint-generation step is executed by AsyncCoDUpdateContextWorkflow and + # already captured in the CoD per-part JSON logs. + exp_type = "solve_task" + rollout = None + metrics = { + "reward": ( + PARSE_FAILURE_REWARD + if format_error + else 0.0 if result is None else result["reward"] + ), + "format_error": format_error, + "expression_schema_error": format_error, + "action_error": action_error, + "response_token_count": response_token_count, + "max_response_tokens": max_response_tokens, + "response_hit_max_tokens": response_hit_max_tokens, + "response_truncated": response_truncated, + } + if result is not None: + rollout = { + "xs": result["xs"], + "vs": result["vs"], + "us": result["us"], + "process_noises": result["process_noises"], + } + metrics.update( + { + "loss": result["loss"], + "x_final": result["x_final"], + "v_final": result["v_final"], + } + ) + + record = { + "timestamp": datetime.datetime.now().isoformat(), + "run_id": run_id, + "batch_id": self.task.batch_id, + "task_id": self.task.task_id, + "task_idx": self.task_idx, + "seed": self.raw_task.get("seed", None), + "exp_type": exp_type, + "a_env": self.a_env, + "b_env": self.b_env, + "min_abs_b_env": self.min_abs_b_env, #!!! + "x0": self.x0, + "v0": self.v0, + "x_target": self.x_target, + "v_target": self.v_target, + "horizon": self.horizon, + "control_penalty_coef": self.control_penalty_coef, + "enable_process_noise": self.enable_process_noise, #!!! + "process_noise_std": self.process_noise_std, #!!! + "task_seed": self.task_seed, #!!! + "system_prompt": system_prompt, + "user_prompt": user_prompt, + "response_text": response_text, + "action_expression": action_expression, + "controller_xml": ( + None + if action_expression is None + else serialize_action_expression_xml(action_expression) + ), + "controller_error_detail": error_detail, + "hint": self.hint, + "rollout": rollout, + "metrics": metrics, + "feedback": feedback, + } + + try: + with open(dump_path, "a", encoding="utf-8") as f: + f.write(json.dumps(record, ensure_ascii=False) + "\n") + except Exception as e: + # Never crash the main workflow because of a logging failure. + print(f"[OptimalControlWorkflow] Failed to dump trajectory: {e}") + + async def _chat(self, messages: List[dict]) -> List[Experience]: + """Generate responses via the local engine or an external API. + + For local engines (vLLM/SGLang) this delegates to ``chat_async``. In + Trinity's native external mode (``model.external_model.enable: true``, + bench only), ``ModelWrapper.model`` is ``None`` and only the OpenAI + client path is available, so the external API is called directly and + each choice is wrapped into a minimal ``Experience``. + """ + if self.model.model is not None: + return await self.model.chat_async(messages, **self.rollout_args) + + # Trinity appends `/v1` when building the OpenAI client, so tolerate + # base URLs that already include the suffix (e.g. DashScope's + # `.../compatible-mode/v1`). + if ( + self.model.openai_async_client is None + and self.model.api_address + and self.model.api_address.rstrip("/").endswith("/v1") + ): + self.model.api_address = self.model.api_address.rstrip("/")[: -len("/v1")] + client = self.model.get_openai_async_client() + rollout_args = self.rollout_args + request_kwargs = { + "model": self.model.config.external_model_config.model_name, + "messages": messages, + "max_completion_tokens": rollout_args.get("max_tokens") + or self.model.config.max_response_tokens, + "n": rollout_args.get("n", 1), + } + if rollout_args.get("temperature") is not None: + request_kwargs["temperature"] = rollout_args["temperature"] + # DashScope-compatible endpoints expect a top-level `enable_thinking` + # in extra_body (not vLLM-style chat_template_kwargs). + if self.model.config.enable_thinking is not None: + request_kwargs["extra_body"] = {"enable_thinking": self.model.config.enable_thinking} + + max_retries = 5 + response = None + for attempt in range(max_retries): + try: + response = await client.chat.completions.create(**request_kwargs) + break + except (openai.RateLimitError, openai.InternalServerError): + if attempt == max_retries - 1: + raise + await asyncio.sleep(1.0 * (2**attempt)) + + usage_metrics = {} + usage = getattr(response, "usage", None) + if usage is not None: + for usage_key in ("prompt_tokens", "completion_tokens", "total_tokens"): + usage_val = getattr(usage, usage_key, None) + if isinstance(usage_val, (int, float)): + usage_metrics[f"usage/{usage_key}"] = float(usage_val) + + experiences: List[Experience] = [] + for choice in response.choices: + finish_reason = str(choice.finish_reason or "").lower() + experiences.append( + Experience( + # Minimal valid token tensor; external APIs return no token ids. + tokens=torch.tensor([0, 0], dtype=torch.int32), + logprobs=torch.tensor([0.0], dtype=torch.float32), + prompt_length=1, + response_text=choice.message.content or "", + truncate_status=("response_truncated" if finish_reason == "length" else None), + metrics=dict(usage_metrics), + info={"finish_reason": finish_reason}, + ) + ) + return experiences + + async def step_async(self, step_num: int) -> Tuple[bool, List[Experience]]: + """Execute one single-turn optimal-control task.""" + system_prompt = self._build_system_prompt() + user_prompt = self._build_user_prompt() + + self.memory = [ + {"role": "system", "content": system_prompt}, + {"role": "user", "content": user_prompt}, + ] + if self.reply_prefix: + self.memory.append({"role": "assistant", "content": self.reply_prefix}) + + responses = await self._chat(self.memory) + + experiences: List[Experience] = [] + for i, response in enumerate(responses): + response.eid.run = i + response.eid.step = 0 + response_text = response.response_text or "" + response.response_text = response_text + action_expression, expression_error = parse_action_expression_submission(response_text) + controller_xml = ( + None + if action_expression is None + else serialize_action_expression_xml(action_expression) + ) + + max_response_tokens = self.rollout_args.get("max_tokens") + if max_response_tokens is None: + max_response_tokens = self.model.config.max_response_tokens + if self.model.model is None: + response_token_count = ( + int(response.metrics.get("usage/completion_tokens", -1)) + if len(responses) == 1 + else -1 + ) + response_hit_max_tokens = response.truncate_status == "response_truncated" + else: + response_token_count = len(response.tokens) - response.prompt_length + response_hit_max_tokens = ( + max_response_tokens is not None and response_token_count >= max_response_tokens + ) + response_truncated = response_hit_max_tokens and action_expression is None + + if action_expression is None: + result = None + evaluation_status = "format_error" + error_detail = expression_error + else: + result, evaluation_status, error_detail = self.env.rollout_action_expression( + action_expression + ) + format_error = evaluation_status == "format_error" + action_error = evaluation_status == "action_error" + reward = ( + PARSE_FAILURE_REWARD + if format_error + else 0.0 if result is None else float(result["reward"]) + ) + feedback = self._build_feedback( + result, + evaluation_status=evaluation_status, + error_detail=error_detail, + response_truncated=response_truncated, + ) + if action_expression is not None: #!!! + self.submitted_control_policy = action_expression #!!! + self.submitted_control_policy_evaluation = { #!!! + "format_parse_success": True, #!!! + "action_execution_success": not action_error, #!!! + "reward": reward, #!!! + "max_score": 1.0, #!!! + "x_final": None if result is None else float(result["x_final"]), #!!! + "v_final": None if result is None else float(result["v_final"]), #!!! + "position_error": ( + None if result is None else abs(float(result["x_final"]) - self.x_target) + ), + "velocity_error": ( + None if result is None else abs(float(result["v_final"]) - self.v_target) + ), + "signed_position_error": ( + None if result is None else float(result["x_final"]) - self.x_target + ), + "signed_velocity_error": ( + None if result is None else float(result["v_final"]) - self.v_target + ), + "x0": self.x0, + "v0": self.v0, + "x_target": self.x_target, + "v_target": self.v_target, + "horizon": self.horizon, + "feedback": feedback, #!!! + } #!!! + trajectory = self._build_rollout_trajectory( + system_prompt, + user_prompt, + response_text, + action_expression, + ) + + response.reward = reward + if response.metrics is None: + response.metrics = {} + response.metrics.update( + { + "reward": reward, + "format_error": float(format_error), + "format_error_termination": float(format_error), + "expression_schema_error": float(format_error), + "action_error": float(action_error), + "response_token_count": float(response_token_count), + "response_hit_max_tokens": float(response_hit_max_tokens), + "response_truncated": float(response_truncated), + "previous_control_policy_used": float(bool(self.previous_control_policy)), + "process_noise_rms": ( #!!! + 0.0 if result is None else float(result["process_noise_rms"]) #!!! + ), #!!! + } + ) + if result is not None: + response.metrics.update( + { + "loss": float(result["loss"]), + "x_final": float(result["x_final"]), + "v_final": float(result["v_final"]), + } + ) + + if response.info is None: + response.info = {} + response.info.update( + { + "sys_prompt": system_prompt, + "user_prompt": user_prompt, + "feedback": feedback, + "trajectory": trajectory, + "a_env": self.a_env, + "b_env": self.b_env, + "min_abs_b_env": self.min_abs_b_env, #!!! + "x0": self.x0, + "v0": self.v0, + "x_target": self.x_target, + "v_target": self.v_target, + "horizon": self.horizon, + "previous_control_policy": self.previous_control_policy or "", + "action_expression": action_expression or "", + "controller_xml": controller_xml or "", + "previous_control_policy_evaluation": ( #!!! + self.previous_control_policy_evaluation or {} #!!! + ), #!!! + "enable_process_noise": self.enable_process_noise, #!!! + "process_noise_std": self.process_noise_std, #!!! + "task_seed": self.task_seed, #!!! + "early_termination_by_format_issue": format_error, + "expression_schema_error": format_error, + "action_error": action_error, + "response_token_count": response_token_count, + "max_response_tokens": max_response_tokens, + "response_hit_max_tokens": response_hit_max_tokens, + "response_truncated": response_truncated, + "controller_error_detail": error_detail or "", + } + ) + + # Persist and update conversation memory. + self._dump_trajectory( + run_id=response.eid.run, + system_prompt=system_prompt, + user_prompt=user_prompt, + response_text=response_text, + action_expression=action_expression, + result=result, + format_error=format_error, + action_error=action_error, + response_token_count=response_token_count, + max_response_tokens=max_response_tokens, + response_hit_max_tokens=response_hit_max_tokens, + response_truncated=response_truncated, + error_detail=error_detail, + feedback=feedback, + ) + + experiences.append(response) + + # Single-step workflow: do not continue after this turn. + return False, experiences + + async def run_async(self) -> List[Experience]: + """Run one optimal-control task with optional pack-level policy reuse.""" + policy_key = None + if self.include_previous_control_policy: + policy_key = self._pack_policy_key() + if self.task_idx == 0: + _PACK_CONTROL_POLICIES.pop(policy_key, None) + previous_record = _PACK_CONTROL_POLICIES.get(policy_key) #!!! + if previous_record: #!!! + self.previous_control_policy = str(previous_record["policy"]) #!!! + self.previous_control_policy_evaluation = dict( #!!! + previous_record.get("evaluation", {}) #!!! + ) #!!! + + _, experiences = await self.step_async(step_num=0) + + if policy_key is not None: + submitted_evaluation = self.submitted_control_policy_evaluation or {} + if self.submitted_control_policy and bool( + submitted_evaluation.get("action_execution_success", False) + ): + _PACK_CONTROL_POLICIES[policy_key] = { #!!! + "policy": self.submitted_control_policy, #!!! + "evaluation": submitted_evaluation, #!!! + } #!!! + pack_size = int(self.raw_task.get("pack_size", self.task_idx + 1)) + if self.task_idx + 1 >= pack_size: + _PACK_CONTROL_POLICIES.pop(policy_key, None) + + if experiences: + # Set step and env-step count only on the last experience; CoD + # post-processing will propagate metrics as needed. The trajectory + # string was already built per-response during step_async. + experiences[-1].eid.step = 0 + experiences[-1].metrics["actual_env_steps"] = 1 + + return experiences diff --git a/trinity/common/workflows/connect_the_dots/pde_discovery/__init__.py b/trinity/common/workflows/connect_the_dots/pde_discovery/__init__.py new file mode 100644 index 00000000000..870580d9a35 --- /dev/null +++ b/trinity/common/workflows/connect_the_dots/pde_discovery/__init__.py @@ -0,0 +1,8 @@ +# -*- coding: utf-8 -*- +"""CoD workflow for PDE discovery.""" + +from trinity.common.workflows.connect_the_dots.pde_discovery.workflow import ( + CoDPDEDiscoveryWorkflow, +) + +__all__ = ["CoDPDEDiscoveryWorkflow"] diff --git a/trinity/common/workflows/connect_the_dots/pde_discovery/candidate.py b/trinity/common/workflows/connect_the_dots/pde_discovery/candidate.py new file mode 100644 index 00000000000..f41b6a0690e --- /dev/null +++ b/trinity/common/workflows/connect_the_dots/pde_discovery/candidate.py @@ -0,0 +1,197 @@ +# -*- coding: utf-8 -*- +"""Dictionary and candidate-equation utilities for PDE discovery.""" + +from __future__ import annotations + +import re +from typing import Dict, List, Tuple + +import numpy as np + + +DEFAULT_DICTIONARY = [ + "u", + "u**2", + "u**3", + "u**4", + "u**5", + "u**7", + "sin(u)", + "sin(2*u)", + "exp(u)-1", + "log(1+u**2)", + "tanh(u)", + "u/(1+u**2)", + "u**2/(1+u**2)", + "u**4/(1+u**4)", + "u/(1+u+u**2)", +] + +SYMBOLIC_COEFFICIENT_RE = re.compile( + r"(? str: + if term == "u": + return r"u(?!\*\*)" + escaped = re.escape(term) + if "+" in term or "-" in term: + # format_equation wraps compound terms to preserve expression grouping. + return rf"(?:{escaped}|\({escaped}\))" + return escaped + + +def symbolic_coefficient_tokens(equation: object) -> List[str]: + return sorted(set(SYMBOLIC_COEFFICIENT_RE.findall(str(equation or "")))) + + +def basis_values(term: str, u: np.ndarray) -> np.ndarray: + if term == "u": + return u + if term == "u**2": + return u**2 + if term == "u**3": + return u**3 + if term == "u**4": + return u**4 + if term == "u**5": + return u**5 + if term == "u**7": + return u**7 + if term == "sin(u)": + return np.sin(u) + if term == "sin(2*u)": + return np.sin(2.0 * u) + if term == "exp(u)-1": + return np.expm1(u) + if term == "log(1+u**2)": + return np.log1p(u**2) + if term == "tanh(u)": + return np.tanh(u) + if term == "u/(1+u**2)": + return u / (1.0 + u**2) + if term == "u**2/(1+u**2)": + return (u**2) / (1.0 + u**2) + if term == "u**4/(1+u**4)": + return (u**4) / (1.0 + u**4) + if term == "u/(1+u+u**2)": + return u / (1.0 + u + u**2) + raise ValueError(f"Unsupported dictionary term: {term}") + + +def dictionary_matrix( + u: np.ndarray, + dictionary: List[str], +) -> Tuple[np.ndarray, List[str]]: + columns = [] + terms = [] + for term in dictionary: + try: + basis = basis_values(str(term), u) + except ValueError: + continue + if np.all(np.isfinite(basis)): + columns.append(basis.astype(float)) + terms.append(str(term)) + if not columns: + return np.zeros((len(u), 0), dtype=float), [] + return np.column_stack(columns), terms + + +def sanitize_dictionary(dictionary: object, fallback_terms: List[str]) -> List[str]: + if not isinstance(dictionary, list) or not dictionary: + return list(fallback_terms) + sanitized = [] + for term in dictionary: + term = str(term) + if term in DEFAULT_DICTIONARY and term not in sanitized: + sanitized.append(term) + return sanitized or list(fallback_terms) + + +def reaction_value(u: np.ndarray, coefficients: Dict[str, float]) -> np.ndarray: + values = np.zeros_like(u, dtype=float) + for term, coef in coefficients.items(): + try: + values = values + coef * basis_values(term, u) + except ValueError: + continue + return values + + +def format_equation(coefficients: Dict[str, float]) -> str: + if not coefficients: + return "0" + + def format_term(term: str) -> str: + if "+" in term or "-" in term: + return f"({term})" + return term + + pieces = [ + f"{coef:.3g}*{format_term(term)}" + for term, coef in coefficients.items() + ] + return pieces[0] + "".join( + f" {'+' if piece[0] != '-' else '-'} {piece.lstrip('-')}" + for piece in pieces[1:] + ) + + +def candidate_coefficients( + candidate: str, + dictionary_terms: List[str], +) -> Dict[str, float]: + text = str(candidate or "").strip() + if "=" in text: + left, right = text.split("=", 1) + # Accept only the equation form used by the environment. + if not re.fullmatch(r"\s*f\s*\(\s*u\s*\)\s*", left, re.IGNORECASE): + return {} + text = right + + coefficients: Dict[str, float] = {} + remaining = re.sub(r"\s+", "", text) + for term in sorted(dictionary_terms, key=len, reverse=True): + pattern = ( + r"([+-]?(?:\d+(?:\.\d*)?|\.\d+)(?:e[+-]?\d+)?)" + r"\*" + + term_pattern(term) + ) + while match := re.search(pattern, remaining, flags=re.IGNORECASE): + # Sum repeated terms rather than silently keeping only one. + coefficients[term] = coefficients.get(term, 0.0) + float(match.group(1)) + start, end = match.span() + remaining = remaining[:start] + (" " * (end - start)) + remaining[end:] + + # Reject unsupported terms instead of silently dropping them. + if remaining.replace(" ", "") not in {"", "+", "-"}: + return {} + return { + term: coefficient + for term, coefficient in coefficients.items() + if abs(coefficient) > 1e-12 + } + + +def candidate_support( + candidate: str, + dictionary: List[str], + dictionary_terms: List[str], +) -> List[str]: + coefficient_support = candidate_coefficients(candidate, dictionary_terms) + if coefficient_support: + return [term for term in dictionary if term in coefficient_support] + compact = candidate.replace(" ", "") + support = [] + for term in sorted(dictionary, key=len, reverse=True): + pattern = term_pattern(term) + if re.search(pattern, compact, flags=re.IGNORECASE): + support.append(term) + compact = re.sub(pattern, "", compact, flags=re.IGNORECASE) + support_set = set(support) + return [term for term in dictionary if term in support_set] diff --git a/trinity/common/workflows/connect_the_dots/pde_discovery/ground_truth.py b/trinity/common/workflows/connect_the_dots/pde_discovery/ground_truth.py new file mode 100644 index 00000000000..6d8bb0fc9f1 --- /dev/null +++ b/trinity/common/workflows/connect_the_dots/pde_discovery/ground_truth.py @@ -0,0 +1,286 @@ +# -*- coding: utf-8 -*- +"""Hidden ground-truth sampling for PDE discovery.""" + +from __future__ import annotations + +import json +import os +import re +from functools import lru_cache +from itertools import product +from typing import Dict, List, Tuple + +import numpy as np + +from .candidate import ( + DEFAULT_DICTIONARY, + reaction_value, +) +from . import pde_numeric + + +DEFAULT_GT_FAMILY_DIR = os.path.join( + os.path.dirname(__file__), + "ground_truth_families", +) +DEFAULT_HIDDEN_GT_FAMILY_FILE = "physical_general.json" +COEFFICIENT_ABS_MIN = 0.5 +COEFFICIENT_ABS_MAX = 2.0 +GT_STABILITY_MAX_ATTEMPTS = 128 +GT_STABILITY_GRID_SIZE = 65 +GT_STABILITY_TIME_STEPS = 401 + + +def _serialized_templates(templates: List[dict]) -> tuple: + return tuple( + ( + tuple(template["terms"]), + tuple( + (term, float(template.get("signs", {}).get(term, 1.0))) + for term in template["terms"] + ), + ) + for template in templates + ) + + +@lru_cache(maxsize=None) +def _calibrated_family_initial_amplitude_upper( + serialized_templates: tuple, + min_terms: int, + max_terms: int, + state_abs_limit: float, +) -> float: + """Find one conservative amplitude bound shared by a GT family.""" + nx = GT_STABILITY_GRID_SIZE + nt = GT_STABILITY_TIME_STEPS + x_grid = np.linspace(0.0, 1.0, nx) + dt = 1.0 / (nt - 1) + dx = 1.0 / (nx - 1) + ratio = dt / (dx * dx) + interior = nx - 2 + lower = -ratio * np.ones(interior - 1, dtype=float) + diag = (1.0 + 2.0 * ratio) * np.ones(interior, dtype=float) + upper = -ratio * np.ones(interior - 1, dtype=float) + + # Cover the corners of the allowed coefficient box for every eligible support. + coefficient_vertices = [] + for terms, raw_signs in serialized_templates: + if not min_terms <= len(terms) <= max_terms: + continue + signs = dict(raw_signs) + for magnitudes in product( + (COEFFICIENT_ABS_MIN, COEFFICIENT_ABS_MAX), + repeat=len(terms), + ): + coefficient_vertices.append( + { + term: (1.0 if signs[term] >= 0.0 else -1.0) * magnitude + for term, magnitude in zip(terms, magnitudes) + } + ) + if not coefficient_vertices: + raise ValueError("Ground-truth family has no templates to calibrate") + + state_limit = float(state_abs_limit) + if not np.isfinite(state_limit) or state_limit <= 0.0: + raise ValueError("state_abs_limit must be finite and positive") + # The first Dirichlet mode is the broadest profile and receives the least + # diffusion damping among the admissible sine modes. + base_profile = np.sin(np.pi * x_grid) + + def family_is_stable(amplitude: float) -> bool: + for coefficients in coefficient_vertices: + state = amplitude * base_profile + for _ in range(nt - 1): + with np.errstate(over="ignore", invalid="ignore", divide="ignore"): + rhs = state[1:-1] + dt * reaction_value( + state[1:-1], coefficients + ) + if not np.all(np.isfinite(rhs)): + return False + next_state = pde_numeric.solve_tridiagonal( + lower, diag, upper, rhs + ) + if ( + not np.all(np.isfinite(next_state)) + or np.any(np.abs(next_state) >= state_limit) + ): + return False + state[1:-1] = next_state + return True + + safe = 0.0 + unsafe = state_limit + resolution = state_limit / (GT_STABILITY_GRID_SIZE - 1) + while unsafe - safe > resolution: + midpoint = (safe + unsafe) / 2.0 + if family_is_stable(midpoint): + safe = midpoint + else: + unsafe = midpoint + if safe <= 0.0: + raise ValueError("Ground-truth family has no positive stable amplitude") + return safe + + +def calibrate_family_initial_amplitude_upper( + templates: List[dict], + min_terms: int, + max_terms: int, + state_abs_limit: float, +) -> float: + """Return a cached, ground-truth-independent bound for one family.""" + return _calibrated_family_initial_amplitude_upper( + _serialized_templates(templates), + int(min_terms), + int(max_terms), + float(state_abs_limit), + ) + + +def load_hidden_gt_templates( + dictionary_terms: List[str], + family: str = DEFAULT_HIDDEN_GT_FAMILY_FILE, +) -> List[dict]: + family_file = str(family).strip() + if family_file.endswith(".json"): + family_file = family_file[:-5] + if not re.fullmatch(r"[A-Za-z0-9_-]+", family_file): + raise ValueError(f"Invalid hidden ground-truth family: {family!r}") + path = os.path.join(DEFAULT_GT_FAMILY_DIR, f"{family_file}.json") + with open(path, "r", encoding="utf-8") as file: + payload = json.load(file) + raw_templates = payload.get("templates", payload) if isinstance(payload, dict) else payload + if not isinstance(raw_templates, list): + raise ValueError(f"Hidden GT family file has no template list: {path}") + + templates = [] + for raw_template in raw_templates: + if not isinstance(raw_template, dict): + continue + terms = [ + str(term) + for term in raw_template.get("terms", []) + if str(term) in DEFAULT_DICTIONARY and str(term) in dictionary_terms + ] + terms = list(dict.fromkeys(terms)) + if not terms: + continue + raw_signs = raw_template.get("signs", {}) + signs = {} + if isinstance(raw_signs, dict): + for term in terms: + try: + sign = float(raw_signs.get(term, 0.0)) + except (TypeError, ValueError): + sign = 0.0 + signs[term] = 1.0 if sign >= 0.0 else -1.0 + templates.append({"terms": terms, "signs": signs}) + if not templates: + raise ValueError(f"Hidden GT family file has no usable templates: {path}") + return templates + + +def sample_hidden_reaction( + rng: np.random.Generator, + templates: List[dict], + min_terms: int, + max_terms: int, + template_index: int | None = None, +) -> Tuple[List[str], Dict[str, float]]: + eligible = [ + template + for template in templates + if min_terms <= len(template["terms"]) <= max_terms + ] + if not eligible: + raise ValueError( + "Hidden GT family has no templates compatible with " + f"min_reaction_terms={min_terms}, max_reaction_terms={max_terms}." + ) + if template_index is None: + template = eligible[int(rng.integers(0, len(eligible)))] + else: + if not 0 <= template_index < len(eligible): + raise ValueError( + f"ground_truth_template_index={template_index} is outside the " + f"eligible template range [0, {len(eligible) - 1}]" + ) + template = eligible[template_index] + terms = sorted(list(template["terms"]), key=DEFAULT_DICTIONARY.index) + signs = template.get("signs", {}) + coefficients = {} + for term in terms: + sign = float(signs.get(term, 0.0)) + if abs(sign) <= 1e-12: + sign = 1.0 if rng.random() < 0.5 else -1.0 + magnitude = round( + float(rng.uniform(COEFFICIENT_ABS_MIN, COEFFICIENT_ABS_MAX)), + 2, + ) + coefficients[term] = (1.0 if sign >= 0.0 else -1.0) * magnitude + return terms, coefficients + + +def is_stable_reaction_candidate( + coefficients: Dict[str, float], + pde_grid_size: int, + pde_time_steps: int, + trajectory_count: int, + pack_seed: int, + initial_amplitude_upper: float, + initial_condition_shape: pde_numeric.InitialConditionShapeConfig, + state_abs_limit: float, +) -> bool: + nx = max(17, min(pde_grid_size, GT_STABILITY_GRID_SIZE)) + nt = max(17, min(pde_time_steps, GT_STABILITY_TIME_STEPS)) + x_grid = np.linspace(0.0, 1.0, nx) + t_grid = np.linspace(0.0, 1.0, nt) + dx = float(x_grid[1] - x_grid[0]) + dt = float(t_grid[1] - t_grid[0]) + r = dt / (dx * dx) + + interior = nx - 2 + lower = -r * np.ones(interior - 1, dtype=float) + diag = (1.0 + 2.0 * r) * np.ones(interior, dtype=float) + upper = -r * np.ones(interior - 1, dtype=float) + + clip_limit = float(state_abs_limit) + if not np.isfinite(clip_limit) or clip_limit <= 0.0: + raise ValueError("state_abs_limit must be finite and positive") + for traj_idx in range(trajectory_count): + u_grid = np.zeros((nt, nx), dtype=float) + u_grid[0] = pde_numeric.initial_condition( + x_grid, + pack_seed=pack_seed, + trajectory_index=traj_idx, + trajectory_count=trajectory_count, + amplitude_upper=initial_amplitude_upper, + shape_config=initial_condition_shape, + ) + u_grid[:, 0] = 0.0 + u_grid[:, -1] = 0.0 + if not np.all(np.isfinite(u_grid[0])): + return False + + for step in range(nt - 1): + rhs_raw = u_grid[step, 1:-1] + dt * reaction_value( + u_grid[step, 1:-1], + coefficients, + ) + if ( + not np.all(np.isfinite(rhs_raw)) + or np.any(rhs_raw <= -clip_limit) + or np.any(rhs_raw >= clip_limit) + ): + return False + u_next = pde_numeric.solve_tridiagonal(lower, diag, upper, rhs_raw) + if ( + not np.all(np.isfinite(u_next)) + or np.any(u_next <= -clip_limit) + or np.any(u_next >= clip_limit) + ): + return False + u_grid[step + 1, 1:-1] = u_next + return True diff --git a/trinity/common/workflows/connect_the_dots/pde_discovery/ground_truth_families/physical_full.json b/trinity/common/workflows/connect_the_dots/pde_discovery/ground_truth_families/physical_full.json new file mode 100644 index 00000000000..4b272744b00 --- /dev/null +++ b/trinity/common/workflows/connect_the_dots/pde_discovery/ground_truth_families/physical_full.json @@ -0,0 +1,61 @@ +{ + "name": "physical_full", + "description": "Expanded physical PDE reaction family with one- and two-term dissipative, periodic-potential, and biochemical kinetics. The agent does not see this file.", + "templates": [ + {"terms": ["u"], "signs": {"u": -1}}, + {"terms": ["u**3"], "signs": {"u**3": -1}}, + {"terms": ["sin(u)"], "signs": {"sin(u)": -1}}, + {"terms": ["tanh(u)"], "signs": {"tanh(u)": -1}}, + {"terms": ["u**2/(1+u**2)"], "signs": {"u**2/(1+u**2)": -1}}, + {"terms": ["u", "u**2"], "signs": {"u": 1, "u**2": -1}}, + {"terms": ["u", "u**3"], "signs": {"u": 1, "u**3": -1}}, + {"terms": ["u", "u**4"], "signs": {"u": 1, "u**4": -1}}, + {"terms": ["u", "u**5"], "signs": {"u": 1, "u**5": -1}}, + {"terms": ["u**2", "u**3"], "signs": {"u**2": 1, "u**3": -1}}, + {"terms": ["u**2", "u**4"], "signs": {"u**2": 1, "u**4": -1}}, + {"terms": ["sin(u)", "u**3"], "signs": {"sin(u)": 1, "u**3": -1}}, + {"terms": ["sin(u)", "u**5"], "signs": {"sin(u)": 1, "u**5": -1}}, + {"terms": ["exp(u)-1", "u**2"], "signs": {"exp(u)-1": 1, "u**2": -1}}, + {"terms": ["exp(u)-1", "u**3"], "signs": {"exp(u)-1": 1, "u**3": -1}}, + {"terms": ["log(1+u**2)", "u**2"], "signs": {"log(1+u**2)": 1, "u**2": -1}}, + {"terms": ["log(1+u**2)", "u**4"], "signs": {"log(1+u**2)": 1, "u**4": -1}}, + {"terms": ["tanh(u)", "u**2"], "signs": {"tanh(u)": 1, "u**2": -1}}, + {"terms": ["tanh(u)", "u**3"], "signs": {"tanh(u)": 1, "u**3": -1}}, + {"terms": ["tanh(u)", "u**5"], "signs": {"tanh(u)": 1, "u**5": -1}}, + {"terms": ["u/(1+u**2)", "u**2"], "signs": {"u/(1+u**2)": 1, "u**2": -1}}, + {"terms": ["u/(1+u**2)", "u**3"], "signs": {"u/(1+u**2)": 1, "u**3": -1}}, + {"terms": ["u/(1+u**2)", "u**5"], "signs": {"u/(1+u**2)": 1, "u**5": -1}}, + {"terms": ["u**2/(1+u**2)", "u"], "signs": {"u**2/(1+u**2)": 1, "u": -1}}, + {"terms": ["u**2/(1+u**2)", "u**3"], "signs": {"u**2/(1+u**2)": 1, "u**3": -1}}, + {"terms": ["u**2/(1+u**2)", "u**4"], "signs": {"u**2/(1+u**2)": 1, "u**4": -1}}, + {"family": "phase_field", "terms": ["u**3", "u**5"], "signs": {"u**3": 1, "u**5": -1}}, + {"family": "phase_field", "terms": ["u", "u**7"], "signs": {"u": 1, "u**7": -1}}, + {"family": "mean_field_relaxation", "terms": ["u", "tanh(u)"], "signs": {"u": -1, "tanh(u)": 1}}, + {"family": "periodic_potential", "terms": ["sin(u)", "sin(2*u)"], "signs": {"sin(u)": 1, "sin(2*u)": -1}}, + {"family": "periodic_potential", "terms": ["sin(u)", "sin(2*u)"], "signs": {"sin(u)": 1, "sin(2*u)": 1}}, + {"family": "hill_kinetics", "terms": ["u**2/(1+u**2)", "u**4/(1+u**4)"], "signs": {"u**2/(1+u**2)": 1, "u**4/(1+u**4)": -1}}, + {"family": "hill_kinetics", "terms": ["u**2/(1+u**2)", "u**4/(1+u**4)"], "signs": {"u**2/(1+u**2)": -1, "u**4/(1+u**4)": 1}}, + {"family": "periodic_potential", "terms": ["sin(u)", "sin(2*u)"], "signs": {"sin(u)": -1, "sin(2*u)": -1}}, + {"family": "periodic_potential", "terms": ["sin(u)", "sin(2*u)"], "signs": {"sin(u)": -1, "sin(2*u)": 1}}, + {"family": "hill_kinetics", "terms": ["u", "u**4/(1+u**4)"], "signs": {"u": -1, "u**4/(1+u**4)": 1}}, + {"family": "haldane_kinetics", "terms": ["u", "u/(1+u+u**2)"], "signs": {"u": -1, "u/(1+u+u**2)": 1}}, + {"family": "periodic_potential_with_absorption", "terms": ["sin(2*u)", "u**2"], "signs": {"sin(2*u)": 1, "u**2": -1}}, + {"family": "periodic_potential_with_absorption", "terms": ["sin(2*u)", "u**3"], "signs": {"sin(2*u)": 1, "u**3": -1}}, + {"family": "periodic_potential_with_absorption", "terms": ["sin(2*u)", "u**5"], "signs": {"sin(2*u)": 1, "u**5": -1}}, + {"family": "periodic_potential_with_absorption", "terms": ["sin(u)", "u**4"], "signs": {"sin(u)": 1, "u**4": -1}}, + {"family": "bounded_activation", "terms": ["tanh(u)", "u**4"], "signs": {"tanh(u)": 1, "u**4": -1}}, + {"family": "substrate_inhibition", "terms": ["u/(1+u**2)", "u**4"], "signs": {"u/(1+u**2)": 1, "u**4": -1}}, + {"family": "haldane_kinetics", "terms": ["u/(1+u+u**2)", "u**2"], "signs": {"u/(1+u+u**2)": 1, "u**2": -1}}, + {"family": "haldane_kinetics", "terms": ["u/(1+u+u**2)", "u**3"], "signs": {"u/(1+u+u**2)": 1, "u**3": -1}}, + {"family": "periodic_potential_with_absorption", "terms": ["sin(2*u)", "u**2"], "signs": {"sin(2*u)": -1, "u**2": -1}}, + {"family": "periodic_potential_with_absorption", "terms": ["sin(2*u)", "u**3"], "signs": {"sin(2*u)": -1, "u**3": -1}}, + {"family": "periodic_potential_with_absorption", "terms": ["sin(2*u)", "u**5"], "signs": {"sin(2*u)": -1, "u**5": -1}}, + {"family": "periodic_potential_with_absorption", "terms": ["sin(u)", "u**4"], "signs": {"sin(u)": -1, "u**4": -1}}, + {"family": "periodic_potential_with_absorption", "terms": ["sin(u)", "u**3"], "signs": {"sin(u)": -1, "u**3": -1}}, + {"family": "substrate_inhibition", "terms": ["u", "u/(1+u**2)"], "signs": {"u": -1, "u/(1+u**2)": 1}}, + {"family": "hill_kinetics", "terms": ["u**2", "u**2/(1+u**2)"], "signs": {"u**2": -1, "u**2/(1+u**2)": 1}}, + {"family": "higher_order_phase_field", "terms": ["u**3", "u**7"], "signs": {"u**3": 1, "u**7": -1}}, + {"family": "higher_order_population", "terms": ["u**2", "u**5"], "signs": {"u**2": 1, "u**5": -1}}, + {"family": "hill_kinetics", "terms": ["u**2/(1+u**2)", "u**5"], "signs": {"u**2/(1+u**2)": 1, "u**5": -1}} + ] +} diff --git a/trinity/common/workflows/connect_the_dots/pde_discovery/ground_truth_families/physical_full_eval.json b/trinity/common/workflows/connect_the_dots/pde_discovery/ground_truth_families/physical_full_eval.json new file mode 100644 index 00000000000..59671315966 --- /dev/null +++ b/trinity/common/workflows/connect_the_dots/pde_discovery/ground_truth_families/physical_full_eval.json @@ -0,0 +1,21 @@ +{ + "name": "physical_full_eval", + "description": "Support-disjoint evaluation split of the expanded physical PDE family. Its basis terms are seen during training, but its two-term compositions are held out. The agent does not see this file.", + "templates": [ + {"terms": ["u", "u**4"], "signs": {"u": 1, "u**4": -1}}, + {"terms": ["u**2", "u**3"], "signs": {"u**2": 1, "u**3": -1}}, + {"terms": ["sin(u)", "u**5"], "signs": {"sin(u)": 1, "u**5": -1}}, + {"terms": ["exp(u)-1", "u**3"], "signs": {"exp(u)-1": 1, "u**3": -1}}, + {"terms": ["log(1+u**2)", "u**4"], "signs": {"log(1+u**2)": 1, "u**4": -1}}, + {"terms": ["tanh(u)", "u**5"], "signs": {"tanh(u)": 1, "u**5": -1}}, + {"terms": ["u/(1+u**2)", "u**5"], "signs": {"u/(1+u**2)": 1, "u**5": -1}}, + {"terms": ["u**2/(1+u**2)", "u**4"], "signs": {"u**2/(1+u**2)": 1, "u**4": -1}}, + {"family": "phase_field", "terms": ["u", "u**7"], "signs": {"u": 1, "u**7": -1}}, + {"family": "periodic_potential", "terms": ["sin(u)", "sin(2*u)"], "signs": {"sin(u)": 1, "sin(2*u)": -1}}, + {"family": "periodic_potential", "terms": ["sin(u)", "sin(2*u)"], "signs": {"sin(u)": 1, "sin(2*u)": 1}}, + {"family": "hill_kinetics", "terms": ["u**2/(1+u**2)", "u**4/(1+u**4)"], "signs": {"u**2/(1+u**2)": 1, "u**4/(1+u**4)": -1}}, + {"family": "hill_kinetics", "terms": ["u**2/(1+u**2)", "u**4/(1+u**4)"], "signs": {"u**2/(1+u**2)": -1, "u**4/(1+u**4)": 1}}, + {"family": "periodic_potential", "terms": ["sin(u)", "sin(2*u)"], "signs": {"sin(u)": -1, "sin(2*u)": -1}}, + {"family": "periodic_potential", "terms": ["sin(u)", "sin(2*u)"], "signs": {"sin(u)": -1, "sin(2*u)": 1}} + ] +} diff --git a/trinity/common/workflows/connect_the_dots/pde_discovery/ground_truth_families/physical_full_eval_4000.json b/trinity/common/workflows/connect_the_dots/pde_discovery/ground_truth_families/physical_full_eval_4000.json new file mode 100644 index 00000000000..c05edee720f --- /dev/null +++ b/trinity/common/workflows/connect_the_dots/pde_discovery/ground_truth_families/physical_full_eval_4000.json @@ -0,0 +1,31 @@ +{ + "name": "physical_full_eval_4000", + "description": "Support-disjoint 4000-task evaluation family with 17 two-term and 8 three-term synthetic reaction kinetics. Every basis term is seen during training, while every complete support is held out. The added templates are physically motivated phenomenological combinations for the benchmark's calibrated non-negative, bounded state range; they are not claims of unique canonical equations. The agent does not see this file.", + "templates": [ + {"terms": ["u", "u**4"], "signs": {"u": 1, "u**4": -1}}, + {"terms": ["u**2", "u**3"], "signs": {"u**2": 1, "u**3": -1}}, + {"terms": ["sin(u)", "u**5"], "signs": {"sin(u)": 1, "u**5": -1}}, + {"terms": ["exp(u)-1", "u**3"], "signs": {"exp(u)-1": 1, "u**3": -1}}, + {"terms": ["log(1+u**2)", "u**4"], "signs": {"log(1+u**2)": 1, "u**4": -1}}, + {"terms": ["tanh(u)", "u**5"], "signs": {"tanh(u)": 1, "u**5": -1}}, + {"terms": ["u/(1+u**2)", "u**5"], "signs": {"u/(1+u**2)": 1, "u**5": -1}}, + {"terms": ["u**2/(1+u**2)", "u**4"], "signs": {"u**2/(1+u**2)": 1, "u**4": -1}}, + {"family": "phase_field", "terms": ["u", "u**7"], "signs": {"u": 1, "u**7": -1}}, + {"family": "periodic_potential", "terms": ["sin(u)", "sin(2*u)"], "signs": {"sin(u)": 1, "sin(2*u)": -1}}, + {"family": "hill_kinetics", "terms": ["u**2/(1+u**2)", "u**4/(1+u**4)"], "signs": {"u**2/(1+u**2)": 1, "u**4/(1+u**4)": -1}}, + {"family": "exponential_absorption", "terms": ["exp(u)-1", "u**5"], "signs": {"exp(u)-1": 1, "u**5": -1}}, + {"family": "exponential_periodic_competition", "terms": ["exp(u)-1", "sin(u)"], "signs": {"exp(u)-1": -1, "sin(u)": 1}}, + {"family": "higher_order_phase_field", "terms": ["u**5", "u**7"], "signs": {"u**5": 1, "u**7": -1}}, + {"family": "bounded_saturating_competition", "terms": ["tanh(u)", "u/(1+u**2)"], "signs": {"tanh(u)": -1, "u/(1+u**2)": 1}}, + {"family": "adjacent_polynomial_competition", "terms": ["u**3", "u**4"], "signs": {"u**3": 1, "u**4": -1}}, + {"family": "even_saturating_competition", "terms": ["log(1+u**2)", "u**2/(1+u**2)"], "signs": {"log(1+u**2)": -1, "u**2/(1+u**2)": 1}}, + {"family": "polynomial_logistic", "terms": ["u", "u**2", "u**3"], "signs": {"u": 1, "u**2": 1, "u**3": -1}}, + {"family": "bounded_polynomial", "terms": ["u", "tanh(u)", "u**3"], "signs": {"u": 1, "tanh(u)": 1, "u**3": -1}}, + {"family": "phase_field", "terms": ["u", "u**3", "u**5"], "signs": {"u": 1, "u**3": -1, "u**5": -1}}, + {"family": "periodic_potential_with_absorption", "terms": ["sin(u)", "sin(2*u)", "u**3"], "signs": {"sin(u)": 1, "sin(2*u)": -1, "u**3": -1}}, + {"family": "higher_order_phase_field", "terms": ["u", "u**3", "u**7"], "signs": {"u": 1, "u**3": 1, "u**7": -1}}, + {"family": "saturating_high_order_kinetics", "terms": ["u**2/(1+u**2)", "u**4/(1+u**4)", "u**5"], "signs": {"u**2/(1+u**2)": 1, "u**4/(1+u**4)": 1, "u**5": -1}}, + {"family": "exponential_bounded_high_order", "terms": ["exp(u)-1", "tanh(u)", "u**7"], "signs": {"exp(u)-1": 1, "tanh(u)": 1, "u**7": -1}}, + {"family": "periodic_saturating_high_order", "terms": ["sin(2*u)", "u**2/(1+u**2)", "u**7"], "signs": {"sin(2*u)": 1, "u**2/(1+u**2)": 1, "u**7": -1}} + ] +} diff --git a/trinity/common/workflows/connect_the_dots/pde_discovery/ground_truth_families/physical_full_eval_hard.json b/trinity/common/workflows/connect_the_dots/pde_discovery/ground_truth_families/physical_full_eval_hard.json new file mode 100644 index 00000000000..973d5c072ed --- /dev/null +++ b/trinity/common/workflows/connect_the_dots/pde_discovery/ground_truth_families/physical_full_eval_hard.json @@ -0,0 +1,24 @@ +{ + "name": "physical_full_eval_hard", + "description": "Support-disjoint evaluation split with 12 two-term and 6 three-term PDEs spanning simple polynomial, periodic, high-order, and saturating kinetics. Its basis terms are seen during training, but its compositions are held out. The agent does not see this file.", + "templates": [ + {"terms": ["u", "u**4"], "signs": {"u": 1, "u**4": -1}}, + {"terms": ["u**2", "u**3"], "signs": {"u**2": 1, "u**3": -1}}, + {"terms": ["sin(u)", "u**5"], "signs": {"sin(u)": 1, "u**5": -1}}, + {"terms": ["exp(u)-1", "u**3"], "signs": {"exp(u)-1": 1, "u**3": -1}}, + {"terms": ["log(1+u**2)", "u**4"], "signs": {"log(1+u**2)": 1, "u**4": -1}}, + {"terms": ["tanh(u)", "u**5"], "signs": {"tanh(u)": 1, "u**5": -1}}, + {"terms": ["u/(1+u**2)", "u**5"], "signs": {"u/(1+u**2)": 1, "u**5": -1}}, + {"terms": ["u**2/(1+u**2)", "u**4"], "signs": {"u**2/(1+u**2)": 1, "u**4": -1}}, + {"family": "phase_field", "terms": ["u", "u**7"], "signs": {"u": 1, "u**7": -1}}, + {"family": "periodic_potential", "terms": ["sin(u)", "sin(2*u)"], "signs": {"sin(u)": 1, "sin(2*u)": -1}}, + {"family": "hill_kinetics", "terms": ["u**2/(1+u**2)", "u**4/(1+u**4)"], "signs": {"u**2/(1+u**2)": 1, "u**4/(1+u**4)": -1}}, + {"family": "exponential_absorption", "terms": ["exp(u)-1", "u**5"], "signs": {"exp(u)-1": 1, "u**5": -1}}, + {"family": "polynomial_logistic", "terms": ["u", "u**2", "u**3"], "signs": {"u": 1, "u**2": 1, "u**3": -1}}, + {"family": "bounded_polynomial", "terms": ["u", "tanh(u)", "u**3"], "signs": {"u": 1, "tanh(u)": 1, "u**3": -1}}, + {"family": "phase_field", "terms": ["u", "u**3", "u**5"], "signs": {"u": 1, "u**3": -1, "u**5": -1}}, + {"family": "periodic_potential_with_absorption", "terms": ["sin(u)", "sin(2*u)", "u**3"], "signs": {"sin(u)": 1, "sin(2*u)": -1, "u**3": -1}}, + {"family": "higher_order_phase_field", "terms": ["u", "u**3", "u**7"], "signs": {"u": 1, "u**3": 1, "u**7": -1}}, + {"family": "saturating_high_order_kinetics", "terms": ["u**2/(1+u**2)", "u**4/(1+u**4)", "u**5"], "signs": {"u**2/(1+u**2)": 1, "u**4/(1+u**4)": 1, "u**5": -1}} + ] +} diff --git a/trinity/common/workflows/connect_the_dots/pde_discovery/ground_truth_families/physical_full_train.json b/trinity/common/workflows/connect_the_dots/pde_discovery/ground_truth_families/physical_full_train.json new file mode 100644 index 00000000000..b85470d4e70 --- /dev/null +++ b/trinity/common/workflows/connect_the_dots/pde_discovery/ground_truth_families/physical_full_train.json @@ -0,0 +1,46 @@ +{ + "name": "physical_full_train", + "description": "Training split of the expanded physical PDE family. It covers every dictionary basis with one- and two-term reactions. The agent does not see this file.", + "templates": [ + {"terms": ["u"], "signs": {"u": -1}}, + {"terms": ["u**3"], "signs": {"u**3": -1}}, + {"terms": ["sin(u)"], "signs": {"sin(u)": -1}}, + {"terms": ["tanh(u)"], "signs": {"tanh(u)": -1}}, + {"terms": ["u**2/(1+u**2)"], "signs": {"u**2/(1+u**2)": -1}}, + {"terms": ["u", "u**2"], "signs": {"u": 1, "u**2": -1}}, + {"terms": ["u", "u**3"], "signs": {"u": 1, "u**3": -1}}, + {"terms": ["u", "u**5"], "signs": {"u": 1, "u**5": -1}}, + {"terms": ["u**2", "u**4"], "signs": {"u**2": 1, "u**4": -1}}, + {"terms": ["sin(u)", "u**3"], "signs": {"sin(u)": 1, "u**3": -1}}, + {"terms": ["exp(u)-1", "u**2"], "signs": {"exp(u)-1": 1, "u**2": -1}}, + {"terms": ["log(1+u**2)", "u**2"], "signs": {"log(1+u**2)": 1, "u**2": -1}}, + {"terms": ["tanh(u)", "u**2"], "signs": {"tanh(u)": 1, "u**2": -1}}, + {"terms": ["tanh(u)", "u**3"], "signs": {"tanh(u)": 1, "u**3": -1}}, + {"terms": ["u/(1+u**2)", "u**2"], "signs": {"u/(1+u**2)": 1, "u**2": -1}}, + {"terms": ["u/(1+u**2)", "u**3"], "signs": {"u/(1+u**2)": 1, "u**3": -1}}, + {"terms": ["u**2/(1+u**2)", "u"], "signs": {"u**2/(1+u**2)": 1, "u": -1}}, + {"terms": ["u**2/(1+u**2)", "u**3"], "signs": {"u**2/(1+u**2)": 1, "u**3": -1}}, + {"family": "phase_field", "terms": ["u**3", "u**5"], "signs": {"u**3": 1, "u**5": -1}}, + {"family": "mean_field_relaxation", "terms": ["u", "tanh(u)"], "signs": {"u": -1, "tanh(u)": 1}}, + {"family": "hill_kinetics", "terms": ["u", "u**4/(1+u**4)"], "signs": {"u": -1, "u**4/(1+u**4)": 1}}, + {"family": "haldane_kinetics", "terms": ["u", "u/(1+u+u**2)"], "signs": {"u": -1, "u/(1+u+u**2)": 1}}, + {"family": "periodic_potential_with_absorption", "terms": ["sin(2*u)", "u**2"], "signs": {"sin(2*u)": 1, "u**2": -1}}, + {"family": "periodic_potential_with_absorption", "terms": ["sin(2*u)", "u**3"], "signs": {"sin(2*u)": 1, "u**3": -1}}, + {"family": "periodic_potential_with_absorption", "terms": ["sin(2*u)", "u**5"], "signs": {"sin(2*u)": 1, "u**5": -1}}, + {"family": "periodic_potential_with_absorption", "terms": ["sin(u)", "u**4"], "signs": {"sin(u)": 1, "u**4": -1}}, + {"family": "bounded_activation", "terms": ["tanh(u)", "u**4"], "signs": {"tanh(u)": 1, "u**4": -1}}, + {"family": "substrate_inhibition", "terms": ["u/(1+u**2)", "u**4"], "signs": {"u/(1+u**2)": 1, "u**4": -1}}, + {"family": "haldane_kinetics", "terms": ["u/(1+u+u**2)", "u**2"], "signs": {"u/(1+u+u**2)": 1, "u**2": -1}}, + {"family": "haldane_kinetics", "terms": ["u/(1+u+u**2)", "u**3"], "signs": {"u/(1+u+u**2)": 1, "u**3": -1}}, + {"family": "periodic_potential_with_absorption", "terms": ["sin(2*u)", "u**2"], "signs": {"sin(2*u)": -1, "u**2": -1}}, + {"family": "periodic_potential_with_absorption", "terms": ["sin(2*u)", "u**3"], "signs": {"sin(2*u)": -1, "u**3": -1}}, + {"family": "periodic_potential_with_absorption", "terms": ["sin(2*u)", "u**5"], "signs": {"sin(2*u)": -1, "u**5": -1}}, + {"family": "periodic_potential_with_absorption", "terms": ["sin(u)", "u**4"], "signs": {"sin(u)": -1, "u**4": -1}}, + {"family": "periodic_potential_with_absorption", "terms": ["sin(u)", "u**3"], "signs": {"sin(u)": -1, "u**3": -1}}, + {"family": "substrate_inhibition", "terms": ["u", "u/(1+u**2)"], "signs": {"u": -1, "u/(1+u**2)": 1}}, + {"family": "hill_kinetics", "terms": ["u**2", "u**2/(1+u**2)"], "signs": {"u**2": -1, "u**2/(1+u**2)": 1}}, + {"family": "higher_order_phase_field", "terms": ["u**3", "u**7"], "signs": {"u**3": 1, "u**7": -1}}, + {"family": "higher_order_population", "terms": ["u**2", "u**5"], "signs": {"u**2": 1, "u**5": -1}}, + {"family": "hill_kinetics", "terms": ["u**2/(1+u**2)", "u**5"], "signs": {"u**2/(1+u**2)": 1, "u**5": -1}} + ] +} diff --git a/trinity/common/workflows/connect_the_dots/pde_discovery/ground_truth_families/physical_general.json b/trinity/common/workflows/connect_the_dots/pde_discovery/ground_truth_families/physical_general.json new file mode 100644 index 00000000000..65785c588f4 --- /dev/null +++ b/trinity/common/workflows/connect_the_dots/pde_discovery/ground_truth_families/physical_general.json @@ -0,0 +1,38 @@ +{ + "name": "physical_general", + "description": "A broader hidden PDE reaction template family over the full nonlinear dictionary. The agent does not see this file.", + "templates": [ + {"terms": ["u"], "signs": {"u": -1}}, + {"terms": ["u**2"], "signs": {"u**2": -1}}, + {"terms": ["u**3"], "signs": {"u**3": -1}}, + {"terms": ["u**4"], "signs": {"u**4": -1}}, + {"terms": ["u**5"], "signs": {"u**5": -1}}, + {"terms": ["sin(u)"], "signs": {"sin(u)": -1}}, + {"terms": ["exp(u)-1"], "signs": {"exp(u)-1": -1}}, + {"terms": ["log(1+u**2)"], "signs": {"log(1+u**2)": -1}}, + {"terms": ["tanh(u)"], "signs": {"tanh(u)": -1}}, + {"terms": ["u/(1+u**2)"], "signs": {"u/(1+u**2)": -1}}, + {"terms": ["u**2/(1+u**2)"], "signs": {"u**2/(1+u**2)": -1}}, + {"terms": ["u", "u**2"], "signs": {"u": 1, "u**2": -1}}, + {"terms": ["u", "u**3"], "signs": {"u": 1, "u**3": -1}}, + {"terms": ["u", "u**4"], "signs": {"u": 1, "u**4": -1}}, + {"terms": ["u", "u**5"], "signs": {"u": 1, "u**5": -1}}, + {"terms": ["u**2", "u**3"], "signs": {"u**2": 1, "u**3": -1}}, + {"terms": ["u**2", "u**4"], "signs": {"u**2": 1, "u**4": -1}}, + {"terms": ["sin(u)", "u**3"], "signs": {"sin(u)": 1, "u**3": -1}}, + {"terms": ["sin(u)", "u**5"], "signs": {"sin(u)": 1, "u**5": -1}}, + {"terms": ["exp(u)-1", "u**2"], "signs": {"exp(u)-1": 1, "u**2": -1}}, + {"terms": ["exp(u)-1", "u**3"], "signs": {"exp(u)-1": 1, "u**3": -1}}, + {"terms": ["log(1+u**2)", "u**2"], "signs": {"log(1+u**2)": 1, "u**2": -1}}, + {"terms": ["log(1+u**2)", "u**4"], "signs": {"log(1+u**2)": 1, "u**4": -1}}, + {"terms": ["tanh(u)", "u**2"], "signs": {"tanh(u)": 1, "u**2": -1}}, + {"terms": ["tanh(u)", "u**3"], "signs": {"tanh(u)": 1, "u**3": -1}}, + {"terms": ["tanh(u)", "u**5"], "signs": {"tanh(u)": 1, "u**5": -1}}, + {"terms": ["u/(1+u**2)", "u**2"], "signs": {"u/(1+u**2)": 1, "u**2": -1}}, + {"terms": ["u/(1+u**2)", "u**3"], "signs": {"u/(1+u**2)": 1, "u**3": -1}}, + {"terms": ["u/(1+u**2)", "u**5"], "signs": {"u/(1+u**2)": 1, "u**5": -1}}, + {"terms": ["u**2/(1+u**2)", "u"], "signs": {"u**2/(1+u**2)": 1, "u": -1}}, + {"terms": ["u**2/(1+u**2)", "u**3"], "signs": {"u**2/(1+u**2)": 1, "u**3": -1}}, + {"terms": ["u**2/(1+u**2)", "u**4"], "signs": {"u**2/(1+u**2)": 1, "u**4": -1}} + ] +} diff --git a/trinity/common/workflows/connect_the_dots/pde_discovery/ground_truth_families/physical_hard.json b/trinity/common/workflows/connect_the_dots/pde_discovery/ground_truth_families/physical_hard.json new file mode 100644 index 00000000000..e39d18e6a8b --- /dev/null +++ b/trinity/common/workflows/connect_the_dots/pde_discovery/ground_truth_families/physical_hard.json @@ -0,0 +1,13 @@ +{ + "name": "physical_hard", + "description": "Hard two-term hidden PDE reactions drawn from phase-field, population, periodic-potential, and saturating biochemical kinetics. The agent does not see this file.", + "templates": [ + {"family": "phase_field", "terms": ["u", "u**3"], "signs": {"u": 1, "u**3": -1}}, + {"family": "population", "terms": ["u", "u**2"], "signs": {"u": 1, "u**2": -1}}, + {"family": "periodic_potential", "terms": ["sin(u)", "sin(2*u)"], "signs": {"sin(u)": -1, "sin(2*u)": -1}}, + {"family": "periodic_potential", "terms": ["sin(u)", "sin(2*u)"], "signs": {"sin(u)": -1, "sin(2*u)": 1}}, + {"family": "hill_kinetics", "terms": ["u", "u**2/(1+u**2)"], "signs": {"u": -1, "u**2/(1+u**2)": 1}}, + {"family": "hill_kinetics", "terms": ["u", "u**4/(1+u**4)"], "signs": {"u": -1, "u**4/(1+u**4)": 1}}, + {"family": "haldane_kinetics", "terms": ["u", "u/(1+u+u**2)"], "signs": {"u": -1, "u/(1+u+u**2)": 1}} + ] +} diff --git a/trinity/common/workflows/connect_the_dots/pde_discovery/ground_truth_families/physical_simplified.json b/trinity/common/workflows/connect_the_dots/pde_discovery/ground_truth_families/physical_simplified.json new file mode 100644 index 00000000000..4c002045254 --- /dev/null +++ b/trinity/common/workflows/connect_the_dots/pde_discovery/ground_truth_families/physical_simplified.json @@ -0,0 +1,23 @@ +{ + "name": "physical_simplified", + "description": "Hidden PDE reaction templates for the simplified physical benchmark. The agent does not see this file.", + "templates": [ + {"terms": ["u"], "signs": {"u": -1}}, + {"terms": ["u**3"], "signs": {"u**3": -1}}, + {"terms": ["u**5"], "signs": {"u**5": -1}}, + {"terms": ["tanh(u)"], "signs": {"tanh(u)": -1}}, + {"terms": ["u/(1+u**2)"], "signs": {"u/(1+u**2)": -1}}, + {"terms": ["log(1+u**2)"], "signs": {"log(1+u**2)": -1}}, + {"terms": ["u**2/(1+u**2)"], "signs": {"u**2/(1+u**2)": -1}}, + {"terms": ["u", "u**2"], "signs": {"u": 1, "u**2": -1}}, + {"terms": ["u", "u**3"], "signs": {"u": 1, "u**3": -1}}, + {"terms": ["u", "u**5"], "signs": {"u": 1, "u**5": -1}}, + {"terms": ["u**2", "u**3"], "signs": {"u**2": 1, "u**3": -1}}, + {"terms": ["sin(u)", "u**3"], "signs": {"sin(u)": 1, "u**3": -1}}, + {"terms": ["tanh(u)", "u**3"], "signs": {"tanh(u)": 1, "u**3": -1}}, + {"terms": ["u/(1+u**2)", "u**3"], "signs": {"u/(1+u**2)": 1, "u**3": -1}}, + {"terms": ["tanh(u)", "u**2"], "signs": {"tanh(u)": 1, "u**2": -1}}, + {"terms": ["u/(1+u**2)", "u**2"], "signs": {"u/(1+u**2)": 1, "u**2": -1}}, + {"terms": ["log(1+u**2)", "u**2"], "signs": {"log(1+u**2)": 1, "u**2": -1}} + ] +} diff --git a/trinity/common/workflows/connect_the_dots/pde_discovery/pde_numeric.py b/trinity/common/workflows/connect_the_dots/pde_discovery/pde_numeric.py new file mode 100644 index 00000000000..367c5bec28e --- /dev/null +++ b/trinity/common/workflows/connect_the_dots/pde_discovery/pde_numeric.py @@ -0,0 +1,293 @@ +# -*- coding: utf-8 -*- +"""Numerical helpers for the PDE discovery workflow.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from typing import Callable + +import numpy as np + + +GROUND_TRUTH_STREAM = 0 +INITIAL_AMPLITUDE_STREAM = 1 +INITIAL_SHAPE_STREAM = 2 + + +@dataclass(frozen=True) +class InitialConditionShapeConfig: + """Configurable complexity prior for random smooth initial conditions.""" + + mode_count_range: tuple[int, int] + + def __post_init__(self) -> None: + value = self.mode_count_range + if len(value) != 2 or value[0] < 1 or value[1] < value[0]: + raise ValueError("mode_count_range must be a positive [min, max] pair") + + @classmethod + def from_mapping( + cls, + raw: Mapping[str, object] | None, + ) -> "InitialConditionShapeConfig": + if raw is None: + raise ValueError("initial_condition_shape configuration is required") + allowed = {"mode_count_range"} + unknown = set(raw) - allowed + if unknown: + raise ValueError( + "Unsupported initial_condition_shape keys: " + + ", ".join(sorted(unknown)) + ) + missing = allowed - set(raw) + if missing: + raise ValueError( + "Missing initial_condition_shape keys: " + + ", ".join(sorted(missing)) + ) + + def count_range(name: str) -> tuple[int, int]: + value = raw[name] + if ( + not isinstance(value, Sequence) + or isinstance(value, (str, bytes)) + or len(value) != 2 + ): + raise ValueError(f"{name} must be a two-element list") + return int(value[0]), int(value[1]) + + return cls(mode_count_range=count_range("mode_count_range")) + + +def initial_condition_amplitudes( + pack_seed: int, + trajectory_count: int, + amplitude_upper: float, +) -> list[float]: + """Sample one amplitude from each stratum below a calibrated upper bound.""" + count = max(1, int(trajectory_count)) + upper = float(amplitude_upper) + if not np.isfinite(upper) or upper <= 0.0: + raise ValueError("initial amplitude upper bound must be finite and positive") + band_edges = np.linspace(0.0, upper, count + 1) + rng = np.random.default_rng( + np.random.SeedSequence([int(pack_seed), INITIAL_AMPLITUDE_STREAM]) + ) + amplitudes = [ + float(rng.uniform(band_edges[index], band_edges[index + 1])) + for index in range(count) + ] + rng.shuffle(amplitudes) + return amplitudes + + +def _normalized_modes( + x_grid: np.ndarray, + rng: np.random.Generator, + mode_count: int, +) -> np.ndarray: + domain_length = float(x_grid[-1] - x_grid[0]) + if domain_length <= 0.0: + raise ValueError("initial-condition grid must span a positive domain") + coordinate = (x_grid - x_grid[0]) / domain_length + modes = np.zeros_like(x_grid, dtype=float) + coefficients = rng.normal(size=mode_count) + for mode, coefficient in enumerate(coefficients, start=1): + modes += coefficient * np.sin(mode * np.pi * coordinate) + scale = float(np.max(np.abs(modes))) + return modes / scale if scale > 1e-12 else modes + + +def _random_initial_condition( + x_grid: np.ndarray, + rng: np.random.Generator, + amplitude: float, + config: InitialConditionShapeConfig, +) -> np.ndarray: + domain_length = float(x_grid[-1] - x_grid[0]) + if domain_length <= 0.0: + raise ValueError("initial-condition grid must span a positive domain") + envelope = np.sin(np.pi * (x_grid - x_grid[0]) / domain_length) + mode_count = int( + rng.integers( + config.mode_count_range[0], + config.mode_count_range[1] + 1, + ) + ) + modes = _normalized_modes(x_grid, rng, mode_count) + shape = envelope * (modes - float(np.min(modes))) + shape_scale = max(float(np.max(shape)), 1e-12) + u0 = amplitude * shape / shape_scale + u0[0] = 0.0 + u0[-1] = 0.0 + return u0 + + +def initial_condition( + x_grid: np.ndarray, + pack_seed: int, + trajectory_index: int = 0, + trajectory_count: int = 1, + amplitude_upper: float | None = None, + shape_config: InitialConditionShapeConfig | None = None, +) -> np.ndarray: + seed_sequence = np.random.SeedSequence( + [int(pack_seed), INITIAL_SHAPE_STREAM, int(trajectory_index)] + ) + rng = np.random.default_rng(seed_sequence) + if amplitude_upper is None: + raise ValueError("initial conditions require a calibrated amplitude upper bound") + amplitudes = initial_condition_amplitudes( + pack_seed, + trajectory_count, + amplitude_upper=amplitude_upper, + ) + if not 0 <= trajectory_index < len(amplitudes): + raise ValueError( + f"trajectory_index={trajectory_index} is outside " + f"trajectory_count={trajectory_count}" + ) + if shape_config is None: + raise ValueError("initial_condition_shape configuration is required") + return _random_initial_condition( + x_grid, + rng, + amplitudes[trajectory_index], + shape_config, + ) + + +def solve_tridiagonal( + lower: np.ndarray, + diag: np.ndarray, + upper: np.ndarray, + rhs: np.ndarray, +) -> np.ndarray: + n = len(diag) + c_prime = np.zeros(max(0, n - 1), dtype=float) + d_prime = np.zeros(n, dtype=float) + denom = diag[0] + if n > 1: + c_prime[0] = upper[0] / denom + d_prime[0] = rhs[0] / denom + for i in range(1, n): + denom = diag[i] - lower[i - 1] * c_prime[i - 1] + if i < n - 1: + c_prime[i] = upper[i] / denom + d_prime[i] = (rhs[i] - lower[i - 1] * d_prime[i - 1]) / denom + solution = np.zeros(n, dtype=float) + solution[-1] = d_prime[-1] + for i in range(n - 2, -1, -1): + solution[i] = d_prime[i] - c_prime[i] * solution[i + 1] + return solution + + +def second_derivative_x(values: np.ndarray, dx: float) -> np.ndarray: + deriv2 = np.zeros_like(values) + deriv2[:, 1:-1] = ( + values[:, 2:] - 2.0 * values[:, 1:-1] + values[:, :-2] + ) / (dx * dx) + deriv2[:, 0] = deriv2[:, 1] + deriv2[:, -1] = deriv2[:, -2] + return deriv2 + + +def scheme_consistent_reaction_residual( + u_grid: np.ndarray, + dx: float, + dt: float, +) -> np.ndarray: + y_grid = np.zeros_like(u_grid, dtype=float) + if len(u_grid) < 2: + return y_grid + u_xx_next = second_derivative_x(u_grid[1:], dx) + y_grid[:-1, 1:-1] = ( + (u_grid[1:, 1:-1] - u_grid[:-1, 1:-1]) / dt + - u_xx_next[:, 1:-1] + ) + return y_grid + + +def simulate_dense_trajectory( + x_grid: np.ndarray, + t_grid: np.ndarray, + dx: float, + dt: float, + lower: np.ndarray, + diag: np.ndarray, + upper: np.ndarray, + trajectory_index: int, + reaction_fn: Callable[[np.ndarray], np.ndarray], + pack_seed: int, + trajectory_count: int = 1, + initial_amplitude_upper: float | None = None, + initial_condition_shape: InitialConditionShapeConfig | None = None, + state_abs_limit: float | None = None, +) -> dict: + if state_abs_limit is None or not np.isfinite(state_abs_limit) or state_abs_limit <= 0.0: + raise ValueError("state_abs_limit must be finite and positive") + nx = len(x_grid) + nt = len(t_grid) + u_grid = np.zeros((nt, nx), dtype=float) + u_grid[0] = initial_condition( + x_grid, + pack_seed, + trajectory_index, + trajectory_count=trajectory_count, + amplitude_upper=initial_amplitude_upper, + shape_config=initial_condition_shape, + ) + u_grid[:, 0] = 0.0 + u_grid[:, -1] = 0.0 + + for step in range(nt - 1): + rhs = u_grid[step, 1:-1] + dt * reaction_fn(u_grid[step, 1:-1]) + rhs = np.clip(rhs, -state_abs_limit, state_abs_limit) + u_next = solve_tridiagonal(lower, diag, upper, rhs) + u_grid[step + 1, 1:-1] = np.clip( + u_next, + -state_abs_limit, + state_abs_limit, + ) + + reaction_residual = scheme_consistent_reaction_residual(u_grid, dx, dt) + # The final state has no forward step, so evaluate f(u) directly. + reaction_residual[-1, 1:-1] = reaction_fn(u_grid[-1, 1:-1]) + return { + "x_grid": x_grid, + "t_grid": t_grid, + "u": u_grid, + "y": reaction_residual, + } + + +def interpolate_dense_field( + field: dict, + field_name: str, + x: np.ndarray, + t: np.ndarray, +) -> np.ndarray: + x_grid = field["x_grid"] + t_grid = field["t_grid"] + values = field[field_name] + + xi = np.searchsorted(x_grid, x, side="right") - 1 + ti = np.searchsorted(t_grid, t, side="right") - 1 + xi = np.clip(xi, 0, len(x_grid) - 2) + ti = np.clip(ti, 0, len(t_grid) - 2) + x0 = x_grid[xi] + x1 = x_grid[xi + 1] + t0 = t_grid[ti] + t1 = t_grid[ti + 1] + wx = (x - x0) / np.maximum(x1 - x0, 1e-12) + wt = (t - t0) / np.maximum(t1 - t0, 1e-12) + + v00 = values[ti, xi] + v01 = values[ti, xi + 1] + v10 = values[ti + 1, xi] + v11 = values[ti + 1, xi + 1] + return ( + (1.0 - wt) * ((1.0 - wx) * v00 + wx * v01) + + wt * ((1.0 - wx) * v10 + wx * v11) + ) diff --git a/trinity/common/workflows/connect_the_dots/pde_discovery/prompts/__init__.py b/trinity/common/workflows/connect_the_dots/pde_discovery/prompts/__init__.py new file mode 100644 index 00000000000..4c8bae2838d --- /dev/null +++ b/trinity/common/workflows/connect_the_dots/pde_discovery/prompts/__init__.py @@ -0,0 +1,28 @@ +# -*- coding: utf-8 -*- +"""Prompt management for CoD PDE discovery.""" + +from pathlib import Path + +from jinja2 import Environment, FileSystemLoader + +PROMPTS_DIR = Path(__file__).parent + + +def get_jinja_env() -> Environment: + return Environment( + loader=FileSystemLoader(PROMPTS_DIR), + trim_blocks=True, + lstrip_blocks=True, + ) + + +def load_system_prompt(**kwargs) -> str: + env = get_jinja_env() + template = env.get_template("system.jinja2") + return template.render(**kwargs) + + +def load_user_prompt(**kwargs) -> str: + env = get_jinja_env() + template = env.get_template("user.jinja2") + return template.render(**kwargs) diff --git a/trinity/common/workflows/connect_the_dots/pde_discovery/prompts/system.jinja2 b/trinity/common/workflows/connect_the_dots/pde_discovery/prompts/system.jinja2 new file mode 100644 index 00000000000..83dd5351f04 --- /dev/null +++ b/trinity/common/workflows/connect_the_dots/pde_discovery/prompts/system.jinja2 @@ -0,0 +1,127 @@ +You are a scientific discovery agent. Your goal is to identify the hidden nonlinear reaction term f(u) in + + partial_t u = partial_xx u + f(u) + +## Scientific setting + +Every task in one pack uses the same hidden reaction term f(u) but observes a different simulated trajectory generated from a different initial condition. Sampled datasets remain available until the pack ends. + +At a sampled point, u_sampled is the observed state and y_sampled is the noisy reaction target that f(u_sampled) should predict. + +## Candidate function library + +Model the unknown reaction as a sparse linear combination f(u) = sum_k c_k phi_k(u). The dictionary below lists the candidate basis functions phi_k(u): + +{{ dictionary_terms | join(", ") }} + +An equation's support is the subset of dictionary terms with nonzero coefficients. The hidden equation has {{ support_size_hint }} active terms; their identities and coefficients are not revealed. + +## Available actions + +The examples below illustrate XML syntax only. Their concrete values are not recommendations; choose arguments appropriate to the current task instead of copying the examples. + +### sample_pde_data + +Choose spatial-temporal coordinate pairs at which to sample the current trajectory. + + + + + + + + + + +point_grid is required. Each point is one pair (x, t) satisfying 0 < x < 1 and 0 <= t <= 1, and the number of points must fit the remaining budget. The action stores the measurements and returns their dataset_id plus coverage summaries; use that id in run_sparse_regression. + +### summarize_pack_evidence + +Inspect hypotheses recorded in previous tasks from the same pack. + + + + false + + + +reveal_equations is optional and defaults to false. Set it to true to include earlier agent-recorded equations, which are hypotheses rather than ground truth. + +### run_sparse_regression + +Fit and compare sparse equations on sampled data using one of two input modes: + +- dictionary: provide a non-empty subset of the candidate library. Regression enumerates and fits supports formed from these terms up to the stated support-size limit. +- candidate_equations: provide one to eight explicit equations with supports to compare. Scoring is restricted to those supports, and every support's coefficients are refitted on the selected data. + +Dictionary mode example: + + + + merged_all + 0.05 + 0.05 + + u + u**2 + u**3 + + + + +Candidate-equation mode example: + + + + merged_all + + 1.2*u - 0.8*u**3 + 0.9*sin(u) - 0.7*u**3 + + + + +- dataset_id: optional; defaults to merged_all, which combines all datasets stored so far in the pack. You may instead use a dataset_id returned by sample_pde_data. +- alpha: optional ridge strength for coefficient screening; larger values can improve stability under noise or correlated terms but shrink coefficients more. Default: 0.05. +- threshold: optional minimum absolute coefficient retained during support selection; larger values favor sparser equations, while smaller values retain weaker terms. Default: 0.05. + +The action returns the best fitted equation, the closest scored equation with different support, and sampled-data fit and conditioning diagnostics. It does not evaluate candidates on hidden ground truth. + +### update_scientific_context + +Finalize the current task and record a tested equation for later tasks in the pack. + + + + 1.2*u - 0.8*u**3 + The cubic support best matches the sampled evidence. + + sin(u) + + + + +preferred_equation is required and must match an equation scored by the latest run_sparse_regression call; copy the returned equation without changing or re-rounding its coefficients. note is an optional short reason for the choice, and uncertain_terms is an optional list of up to four plausible but unresolved dictionary terms. These annotations do not alter preferred_equation. The action is accepted only after sampling and regression in the current task, and an accepted update finalizes the task. + +## Equation syntax + +Use explicit numeric coefficients, not placeholders such as c1, a, or alpha. Treat each dictionary entry as one complete basis function and parenthesize a term containing an internal plus or minus sign: write 0.7*(exp(u)-1), not 0.7*exp(u)-1. + +## Sampled-data feedback + +Regression feedback may include fit summaries and these values for each point in the selected dataset: + + (x, t, u_sampled, y_sampled, f_pred(u_sampled), abs_error) + +Here f_pred is the candidate prediction, abs_error = abs(y_sampled - f_pred), and sampled-data max error is the largest abs_error. Measurements may be noisy. Selected-support kappa and full-dictionary kappa are condition numbers; larger values mean stronger collinearity and weaker numerical distinguishability. The condition penalty is the non-negative log-scale excess of selected-support kappa over the configured threshold. + +The closest different-support candidate is the alternative support with the lowest sampled-data max error. Its relative margin is (alternative_error - best_error) / best_error: a value near zero means the sampled observations poorly distinguish the supports, while a larger positive value means clearer separation. All diagnostics use sampled observations only. + +## Response protocol + +Regardless of how much reasoning you write, every response must end with exactly +one ... block wrapping exactly one action from the list above. +Keep any reasoning before the block brief. Do not stop after the reasoning: if +your analysis is unfinished or the budget is nearly exhausted, still emit your +best available action inside rather than ending without one. Never write +text after the closing , and never output more than one block. diff --git a/trinity/common/workflows/connect_the_dots/pde_discovery/prompts/user.jinja2 b/trinity/common/workflows/connect_the_dots/pde_discovery/prompts/user.jinja2 new file mode 100644 index 00000000000..6def6029cd7 --- /dev/null +++ b/trinity/common/workflows/connect_the_dots/pde_discovery/prompts/user.jinja2 @@ -0,0 +1,53 @@ +Task {{ task_position }}/{{ pack_size }} | Step {{ current_step }}/{{ max_steps }} +Point budget: {{ budget_remaining }}/{{ point_budget }} remaining + +{% if task_description %} +Task: {{ task_description }} +{% endif %} + +{% if available_datasets %} +Datasets: {{ available_datasets | join(", ") }} +{% else %} +Datasets: none +{% endif %} +{% if last_sampled_max_error is not none %} +Latest sampled-data max error: {{ "%.6g" | format(last_sampled_max_error) }} +{% else %} +Latest sampled-data max error: none +{% endif %} + +{% if evidence_memory %} +Prior task evidence: +- task_evidence_status: {{ evidence_memory.get("task_evidence_status", "unknown") }} +{% set observation_summary = evidence_memory.get("observation_summary", {}) %} +{% if observation_summary.get("n") is not none %} +- sampled_points: {{ observation_summary.get("n", 0) }} +{% endif %} +{% if observation_summary.get("u_span") is not none %} +- sampled_u_span: {{ "%.6g" | format(observation_summary.get("u_span", 0.0)) }} +{% endif %} +{% if observation_summary.get("preferred_sampled_max_error") is not none %} +- preferred_sampled_max_error: {{ "%.6g" | format(observation_summary.get("preferred_sampled_max_error", 0.0)) }} +{% endif %} +{% if observation_summary.get("relative_margin") is not none %} +- relative_margin: {{ "%.6g" | format(observation_summary.get("relative_margin", 0.0)) }} +{% endif %} +{% endif %} + +{% if action_feedback %} +Action feedback: +{{ action_feedback }} +{% else %} +No action has been taken in this task yet. +{% endif %} + +Choose the next action. + +{% if budget_remaining > 0 %} +Sampling is available. sample_pde_data must use point_grid coordinates and cannot request more than {{ budget_remaining }} points. +{% else %} +No sampling budget remains. If regression has not been run after the latest samples, run_sparse_regression next; otherwise use update_scientific_context or inspect prior evidence. +{% endif %} +run_sparse_regression may be called whenever you have enough sampled evidence. update_scientific_context requires current-task sampling and at least one current-task run_sparse_regression call. You may decide whether to sample more points, run another regression, or finalize based on the available evidence. + +Include exactly one XML action wrapped in tags. diff --git a/trinity/common/workflows/connect_the_dots/pde_discovery/regression.py b/trinity/common/workflows/connect_the_dots/pde_discovery/regression.py new file mode 100644 index 00000000000..69ef5651186 --- /dev/null +++ b/trinity/common/workflows/connect_the_dots/pde_discovery/regression.py @@ -0,0 +1,417 @@ +# -*- coding: utf-8 -*- +"""Sparse-regression helpers for PDE discovery. + +This module follows the SINDy-style sparse-identification pattern: build a +candidate-library matrix, select sparse supports, then refit coefficients on +the selected terms. This style of sparse regression is commonly used for +data-driven governing-equation and PDE discovery. + +References: Brunton, Steven L., Joshua L. Proctor, and J. Nathan Kutz. +"Discovering governing equations from data: Sparse identification of nonlinear +dynamical systems." arXiv preprint arXiv:1509.03580 (SINDy); Rudy, Samuel H., +et al. "Data-driven discovery of partial differential equations." arXiv +preprint arXiv:1609.06401 (PDE-FIND). +""" + +from __future__ import annotations + +import math +from itertools import combinations +from typing import List, Optional, Tuple + +import numpy as np + +from . import candidate as candidate_utils + + +KAPPA_CAP = 1e12 +DEBIASED_REFIT_ALPHA = 1e-8 + + +def max_abs_error(y_true: np.ndarray, y_pred: np.ndarray) -> float: + if len(y_true) == 0: + return 1.0 + return float(np.max(np.abs(y_true - y_pred))) + + +def relative_error_margin(second_best_error: float, best_error: float) -> float: + gap = max(0.0, float(second_best_error) - float(best_error)) + if best_error > 0.0: + return float(gap / float(best_error)) + return float("inf") if gap > 0.0 else 0.0 + + +def effective_kappa(u: np.ndarray, dictionary: List[str]) -> float: + theta, _ = candidate_utils.dictionary_matrix(u, dictionary) + if theta.shape[1] == 0: + return KAPPA_CAP + scales = np.linalg.norm(theta, axis=0) + valid = scales > 1e-12 + if np.count_nonzero(valid) == 0: + return KAPPA_CAP + theta_norm = theta[:, valid] / scales[valid] + try: + singular_values = np.linalg.svd(theta_norm, compute_uv=False) + except np.linalg.LinAlgError: + return KAPPA_CAP + if len(singular_values) == 0: + return KAPPA_CAP + sigma_min = float(singular_values[-1]) + sigma_max = float(singular_values[0]) + if sigma_min <= 1e-12: + return KAPPA_CAP + return min(float(sigma_max / sigma_min), KAPPA_CAP) + + +def sparse_regression_result( + u: np.ndarray, + y: np.ndarray, + u_objective: np.ndarray, + y_objective: np.ndarray, + dictionary: list, + alpha: float, + threshold: float, + max_reaction_terms: int, + default_blind_penalty: float, + kappa_threshold: float, + allowed_supports: Optional[List[List[str]]] = None, +) -> dict: + theta, terms = candidate_utils.dictionary_matrix(u, dictionary) + if theta.shape[1] == 0: + target_y = y_objective if len(y_objective) else y + max_error = max_abs_error(target_y, np.zeros_like(target_y, dtype=float)) + penalty = max_error + float(default_blind_penalty) + return { + "support": [], + "coefficients": {}, + "equation": "0", + "max_error": max_error, + "condition_ratio": 1.0, + "kappa": KAPPA_CAP, + "penalty": penalty, + "reward_penalty": penalty, + "relative_margin": 0.0, + "second_best_max_error": float("nan"), + "second_best_equation": "not evaluated: no valid dictionary terms", + "diagnostics": {}, + "candidate_diagnostics": [], + "feedback": "Sparse regression failed: no valid dictionary terms.", + } + + dictionary_kappa = effective_kappa(u, terms) + scales = np.linalg.norm(theta, axis=0) + valid = scales > 1e-12 + theta_valid = theta[:, valid] + terms_valid = [term for term, keep in zip(terms, valid) if keep] + scales_valid = scales[valid] + theta_scaled = theta_valid / scales_valid + theta_objective_valid, _ = candidate_utils.dictionary_matrix(u_objective, terms_valid) + active = np.ones(theta_scaled.shape[1], dtype=bool) + coef_scaled_full = np.zeros(theta_scaled.shape[1], dtype=float) + ridge = max(alpha, 0.0) + + for _ in range(8): + if not np.any(active): + break + design = theta_scaled[:, active] + gram = design.T @ design + ridge * np.eye(design.shape[1]) + rhs = design.T @ y + try: + coef_active = np.linalg.solve(gram, rhs) + except np.linalg.LinAlgError: + coef_active = np.linalg.lstsq(gram, rhs, rcond=None)[0] + coef_unscaled_active = coef_active / scales_valid[active] + keep_active = np.abs(coef_unscaled_active) >= threshold + coef_scaled_full[:] = 0.0 + coef_scaled_full[np.where(active)[0]] = coef_active + if np.all(keep_active): + break + active_indices = np.where(active)[0] + active[active_indices[~keep_active]] = False + + hard_cap_terms = max(1, min(max_reaction_terms, len(terms_valid))) + coefficient_threshold = max(float(threshold), 0.0) + + def objective( + max_error_value: float, + support_terms: List[str], + ) -> Tuple[float, float, float]: + support_kappa = effective_kappa(u, support_terms) if support_terms else dictionary_kappa + condition_ratio = max( + 0.0, math.log10(support_kappa) - math.log10(kappa_threshold) + ) + return ( + max_error_value, + support_kappa, + condition_ratio, + ) + + def fit_support(indices: List[int]) -> dict: + support_indices = sorted(set(indices)) + coef_unscaled = np.zeros(theta_scaled.shape[1], dtype=float) + + if support_indices: + design = theta_valid[:, support_indices] + rhs = design.T @ y + try: + if DEBIASED_REFIT_ALPHA > 0.0: + gram = design.T @ design + DEBIASED_REFIT_ALPHA * np.eye( + design.shape[1] + ) + coef_selected = np.linalg.solve(gram, rhs) + else: + coef_selected = np.linalg.lstsq(design, y, rcond=None)[0] + except np.linalg.LinAlgError: + coef_selected = np.linalg.lstsq(design, y, rcond=None)[0] + + kept_pairs = [ + (idx, float(coef)) + for idx, coef in zip(support_indices, coef_selected) + if abs(float(coef)) >= coefficient_threshold + ] + kept_indices = [idx for idx, _ in kept_pairs] + if kept_indices != support_indices: + return fit_support(kept_indices) + for idx, coef in kept_pairs: + coef_unscaled[idx] = coef + + coefficients = { + terms_valid[idx]: round(float(coef_unscaled[idx]), 3) + for idx in support_indices + } + # Score the same rounded equation exposed to the agent. + equation = candidate_utils.format_equation(coefficients) + coefficients = candidate_utils.candidate_coefficients(equation, terms_valid) + coef_unscaled = np.array( + [coefficients.get(term, 0.0) for term in terms_valid], dtype=float + ) + pred_objective = theta_objective_valid @ coef_unscaled + max_error_value = max_abs_error(y_objective, pred_objective) + score, support_kappa, condition_ratio = objective( + max_error_value, + list(coefficients), + ) + return { + "indices": support_indices, + "support": list(coefficients), + "coefficients": coefficients, + "coef_unscaled": coef_unscaled, + "max_error": max_error_value, + "condition_ratio": condition_ratio, + "kappa": support_kappa, + "penalty": score, + } + + initial_indices = [ + int(idx) + for idx, keep in enumerate(active) + if keep and abs(float(coef_scaled_full[idx] / scales_valid[idx])) >= coefficient_threshold + ] + if len(initial_indices) > hard_cap_terms: + seed_coefficients = coef_scaled_full / scales_valid + initial_indices = sorted( + initial_indices, + key=lambda idx: abs(float(seed_coefficients[idx])), + reverse=True, + )[:hard_cap_terms] + + fit_cache = {} + + def cached_fit(indices: List[int]) -> dict: + key = tuple(sorted(set(indices))) + if key not in fit_cache: + fit_cache[key] = fit_support(list(key)) + return fit_cache[key] + + candidate_supports = set() + if allowed_supports is not None: + term_to_index = {term: idx for idx, term in enumerate(terms_valid)} + for support_terms in allowed_supports: + support_indices = [ + term_to_index[term] for term in support_terms if term in term_to_index + ] + candidate_supports.add(tuple(sorted(set(support_indices)))) + else: + if initial_indices: + candidate_supports.add(tuple(sorted(initial_indices))) + for target_size in range(1, hard_cap_terms + 1): + for support_indices in combinations(range(len(terms_valid)), target_size): + candidate_supports.add(tuple(support_indices)) + if not candidate_supports: + candidate_supports.add(tuple()) + + support_candidates = [ + cached_fit(list(candidate_support)) + for candidate_support in sorted(candidate_supports) + ] + + def add_relative_margin(candidate: dict) -> dict: + candidate = dict(candidate) + support_key = tuple(candidate["support"]) + alternatives = [ + other + for other in support_candidates + if tuple(other["support"]) != support_key + ] + if alternatives: + second_best = min(alternatives, key=lambda item: item["max_error"]) + second_best_max_error = float(second_best["max_error"]) + second_best_equation = candidate_utils.format_equation( + second_best["coefficients"] + ) + relative_margin = relative_error_margin( + second_best_max_error, + float(candidate["max_error"]), + ) + else: + second_best_max_error = float("nan") + second_best_equation = "not evaluated: no alternative support supplied" + relative_margin = 0.0 + candidate.update( + { + "relative_margin": float(relative_margin), + "second_best_max_error": second_best_max_error, + "second_best_equation": second_best_equation, + "reward_penalty": float(candidate["max_error"]), + "diagnostics": { + "condition_ratio": float(candidate["condition_ratio"]), + "relative_margin": float(relative_margin), + "second_best_max_error": second_best_max_error, + "second_best_equation": second_best_equation, + }, + "penalty": float(candidate["max_error"]), + } + ) + return candidate + + scored_candidates = [add_relative_margin(candidate) for candidate in support_candidates] + scored_candidates.sort(key=lambda item: item["penalty"]) + best = scored_candidates[0] + equation = candidate_utils.format_equation(best["coefficients"]) + selection_scope = ( + "agent-specified refit set" + if allowed_supports is not None + else "agent-specified dictionary" + ) + selection_note = ( + f" Scored {len(support_candidates)} support(s) from the {selection_scope}." + ) + if len({tuple(candidate["support"]) for candidate in support_candidates}) < 2: + selection_note += ( + " No different-support alternative was supplied, so relative_margin is " + "reported as 0." + ) + + candidate_diagnostics = [ + { + "equation": candidate_utils.format_equation(candidate["coefficients"]), + "support": list(candidate["support"]), + "training_max_error": float(candidate["max_error"]), + "condition_penalty": float(candidate["condition_ratio"]), + "relative_margin": float(candidate["relative_margin"]), + "second_best_equation": candidate["second_best_equation"], + "diagnostics": dict(candidate["diagnostics"]), + "step_penalty": float(candidate["reward_penalty"]), + } + for candidate in sorted(scored_candidates, key=lambda item: item["penalty"]) + ] + candidate_summary = ( + f"best equation={equation}; " + f"sampled-data max error={best['max_error']:.3g}; " + f"condition penalty={best['condition_ratio']:.3g}; " + f"closest different-support candidate={best['second_best_equation']}; " + f"relative margin={best['relative_margin']:.3g}" + ) + feedback = ( + "Sparse regression with thresholded ridge screening, hard-capped support " + "enumeration, and debiased coefficient refit. " + f"Selected-support kappa = {best['kappa']:.3g}; " + f"full-dictionary kappa = {dictionary_kappa:.3g}. " + f"Best fit: f(u) = {equation}. " + f"sampled-data max error = {best['max_error']:.3g}. " + f"Coefficient threshold = {coefficient_threshold:.3g}; " + f"Non-zero terms: {len(best['support'])}. " + f"condition penalty term = {best['condition_ratio']:.3g}. " + f"Closest scored different-support equation {best['second_best_equation']} " + f"has max error = {best['second_best_max_error']:.3g}; " + f"relative margin = {best['relative_margin']:.3g}. " + f"{selection_note} " + f"Candidate sampled-data diagnostics: {candidate_summary}. " + "This tool updates only the latest regression candidate; the scientific " + "conclusion inherited by future tasks is set only by update_scientific_context." + ) + return { + "support": best["support"], + "coefficients": best["coefficients"], + "equation": equation, + "max_error": best["max_error"], + "condition_ratio": best["condition_ratio"], + "relative_margin": best["relative_margin"], + "second_best_max_error": best["second_best_max_error"], + "second_best_equation": best["second_best_equation"], + "kappa": best["kappa"], + "penalty": best["reward_penalty"], + "reward_penalty": best["reward_penalty"], + "diagnostics": dict(best["diagnostics"]), + "candidate_diagnostics": candidate_diagnostics, + "feedback": feedback, + } + + +def candidate_objective_diagnostics( + candidate: str, + data: dict, + objective_data: dict, + dictionary: List[str], + dictionary_terms: List[str], + last_candidate_diagnostics: List[dict], + kappa_threshold: float, +) -> dict: + coefficients = candidate_utils.candidate_coefficients(candidate, dictionary_terms) + support = [term for term in dictionary if abs(coefficients.get(term, 0.0)) > 0.0] + pred_objective = candidate_utils.reaction_value(objective_data["u"], coefficients) + max_error = max_abs_error(objective_data["y"], pred_objective) + support_kappa = effective_kappa(data["u"], support) if support else KAPPA_CAP + condition_penalty = max( + 0.0, + math.log10(support_kappa) - math.log10(kappa_threshold), + ) + support_key = tuple(support) + alternatives = [ + item + for item in last_candidate_diagnostics + if tuple(item.get("support", [])) != support_key + ] + if alternatives: + second_best = min( + alternatives, + key=lambda item: item["training_max_error"], + ) + second_best_max_error = float(second_best["training_max_error"]) + second_best_equation = str(second_best["equation"]) + relative_margin = relative_error_margin(second_best_max_error, max_error) + else: + second_best_max_error = float("nan") + second_best_equation = "not evaluated: no alternative support supplied" + relative_margin = 0.0 + + step_penalty = max_error + diagnostics = { + "condition_ratio": condition_penalty, + "relative_margin": relative_margin, + "second_best_max_error": second_best_max_error, + "second_best_equation": second_best_equation, + } + return { + "equation": candidate, + "support": support, + "coefficients": coefficients, + "training_max_error": max_error, + "condition_penalty": condition_penalty, + "relative_margin": relative_margin, + "second_best_max_error": second_best_max_error, + "second_best_equation": second_best_equation, + "step_penalty": step_penalty, + "reward_penalty": step_penalty, + "diagnostics": diagnostics, + } diff --git a/trinity/common/workflows/connect_the_dots/pde_discovery/update_context_workflow.py b/trinity/common/workflows/connect_the_dots/pde_discovery/update_context_workflow.py new file mode 100644 index 00000000000..d2604d6df91 --- /dev/null +++ b/trinity/common/workflows/connect_the_dots/pde_discovery/update_context_workflow.py @@ -0,0 +1,46 @@ +# -*- coding: utf-8 -*- +"""PDE-specific context update without exposing the training reward.""" + +from typing import List + +from trinity.common.workflows.connect_the_dots.update_context_workflow import ( + AsyncCoDUpdateContextWorkflow, +) + + +class PDEUpdateContextWorkflow(AsyncCoDUpdateContextWorkflow): + """Build the standard context-update prompt without its reward block.""" + + def build_messages(self) -> List[dict]: + messages = super().build_messages() + user_prompt = messages[-1]["content"] + feedback_block = ( + f"\n\nEnvironment feedback: {self.feedback}\n\n--- Your job ---" + ) + before_feedback, feedback_marker, after_feedback = user_prompt.rpartition( + feedback_block + ) + before_reward, reward_marker, reward_text = before_feedback.rpartition( + "\nReward: " + ) + if ( + not feedback_marker + or not reward_marker + or not reward_text.strip() + or "\n" in reward_text + ): + raise ValueError( + "Expected reward block was not found in context-update prompt" + ) + messages[-1]["content"] = ( + before_reward + feedback_marker + after_feedback + ) + messages[-1]["content"] += """ + +Protocol reminder: after your brief reasoning, end with exactly this block: +--- Start of updated hints --- +- Your concise, transferable hints go here. +--- End of updated hints --- +Do not shorten the opening delimiter to "--- Updated hints ---", omit either delimiter, or write anything after the end delimiter. +""" + return messages diff --git a/trinity/common/workflows/connect_the_dots/pde_discovery/workflow.py b/trinity/common/workflows/connect_the_dots/pde_discovery/workflow.py new file mode 100644 index 00000000000..a9f7f3d29fe --- /dev/null +++ b/trinity/common/workflows/connect_the_dots/pde_discovery/workflow.py @@ -0,0 +1,1360 @@ +# -*- coding: utf-8 -*- +"""CoD PDE discovery workflow. + +This is a lightweight implementation of the PDE discovery demo in +``pde_discovery_in_CoD``. The environment accepts XML actions in plain-text +responses, keeps raw sampled points in a stateful pack buffer, and +lets the outer CoD workflow autonomously update the context after each task. +""" + +from __future__ import annotations + +import json +import logging +import math +import os +import re +import weakref +from dataclasses import dataclass, field +from typing import Dict, List, Optional, Tuple + +import numpy as np + +from trinity.common.experience import Experience +from trinity.common.models.model import ModelWrapper +from trinity.common.workflows.connect_the_dots.base_workflow import ( + AsyncCoDMultiStepWorkflow, +) +from trinity.common.workflows.connect_the_dots.pde_discovery import candidate as candidate_utils +from trinity.common.workflows.connect_the_dots.pde_discovery import ground_truth +from trinity.common.workflows.connect_the_dots.pde_discovery import pde_numeric +from trinity.common.workflows.connect_the_dots.pde_discovery import regression +from trinity.common.workflows.connect_the_dots.pde_discovery.candidate import ( + DEFAULT_DICTIONARY, +) +from trinity.common.workflows.connect_the_dots.pde_discovery.prompts import ( + load_system_prompt, + load_user_prompt, +) +from trinity.common.workflows.connect_the_dots.utils import parse_xml_answer +from trinity.common.workflows.workflow import Task + + +KAPPA_CAP = regression.KAPPA_CAP +XML_LIST_TAGS = {"dictionary", "candidate_equations", "uncertain_terms"} + +logger = logging.getLogger(__name__) + + +@dataclass +class PDEPackState: + """Persistent state shared by tasks in one CoD pack.""" + + datasets: Dict[str, dict] = field(default_factory=dict) + sample_counter: int = 0 + reaction_terms: List[str] = field(default_factory=list) + reaction_coefficients: Dict[str, float] = field(default_factory=dict) + ground_truth_template_index: Optional[int] = None + environment_seed: Optional[int] = None + latest_regression_equation: str = "" + preferred_equation: str = "" + preferred_support: List[str] = field(default_factory=list) + preferred_coefficients: Dict[str, float] = field(default_factory=dict) + dense_field: Optional[dict] = None + last_context_update: Dict[str, object] = field(default_factory=dict) + evidence_history: List[dict] = field(default_factory=list) + next_task_position: int = 1 + + +def _normalize_point_grid( + raw_grid: object, +) -> Optional[List[Tuple[float, float]]]: + if not isinstance(raw_grid, dict): + return None + + raw_points = raw_grid.get("point") + if isinstance(raw_points, dict): + raw_points = [raw_points] + if not isinstance(raw_points, list): + return None + + try: + x_values = [float(point["x"]) for point in raw_points] + t_values = [float(point["t"]) for point in raw_points] + except (KeyError, TypeError, ValueError): + return None + if not x_values or not t_values: + return None + if any(not 0.0 < value < 1.0 for value in x_values): + return None + if any(not 0.0 <= value <= 1.0 for value in t_values): + return None + if len(x_values) != len(t_values): + return None + return list(zip(x_values, t_values)) + + +class CoDPDEDiscoveryWorkflow(AsyncCoDMultiStepWorkflow): + """CoD workflow for discovering nonlinear PDE reaction terms.""" + + is_async: bool = True + # use weakref to avoid memory leakage + _PACK_STATES: weakref.WeakValueDictionary[str, PDEPackState] = ( + weakref.WeakValueDictionary() + ) + + def __init__( + self, + model: ModelWrapper, + task: Task, + auxiliary_models: Optional[List] = None, + use_openai_client: bool = False, + ): + super().__init__( + task=task, + model=model, + auxiliary_models=auxiliary_models, + use_openai_client=use_openai_client, + ) + self.reset(task) + + def reset(self, task: Task): + if isinstance(getattr(task, "workflow_args", None), dict): + task.workflow_args.setdefault("context_compression_mode", "keep_all") + super().reset(task) + args = task.workflow_args if hasattr(task, "workflow_args") else {} + self.dictionary_terms = args.get("dictionary_terms", DEFAULT_DICTIONARY) + self.min_reaction_terms = max(1, int(args.get("min_reaction_terms", 1))) + self.max_reaction_terms = max( + self.min_reaction_terms, + int(args.get("max_reaction_terms", 2)), + ) + self.dictionary_terms = [ + str(term) for term in self.dictionary_terms if str(term) in DEFAULT_DICTIONARY + ] or list(DEFAULT_DICTIONARY) + self.ground_truth_family = str( + args.get("ground_truth_family", ground_truth.DEFAULT_HIDDEN_GT_FAMILY_FILE) + ) + self.initial_condition_shape = ( + pde_numeric.InitialConditionShapeConfig.from_mapping( + args.get("initial_condition_shape") + ) + ) + if "pde_state_abs_limit" not in args: + raise ValueError("pde_state_abs_limit workflow argument is required") + self.pde_state_abs_limit = float(args["pde_state_abs_limit"]) + if ( + not math.isfinite(self.pde_state_abs_limit) + or self.pde_state_abs_limit <= 0.0 + ): + raise ValueError("pde_state_abs_limit must be finite and positive") + self.hidden_gt_templates = self._load_hidden_gt_templates() + self.initial_amplitude_upper = ( + ground_truth.calibrate_family_initial_amplitude_upper( + self.hidden_gt_templates, + self.min_reaction_terms, + self.max_reaction_terms, + state_abs_limit=self.pde_state_abs_limit, + ) + ) + self.max_steps = args.get("max_steps", 10) + self.point_budget = int(args.get("point_budget", 15)) + self.default_blind_penalty = args.get("blind_submission_penalty", 100.0) + self.kappa_threshold = float(args.get("kappa_threshold", 100.0)) + if not math.isfinite(self.kappa_threshold) or self.kappa_threshold <= 0.0: + raise ValueError("kappa_threshold must be a finite positive number") + self.noise_level = float(args.get("noise_level", 0.002)) + if not math.isfinite(self.noise_level) or self.noise_level < 0.0: + raise ValueError("noise_level must be a finite non-negative number") + self.pde_grid_size = int(args.get("pde_grid_size", 257)) + self.pde_time_steps = int(args.get("pde_time_steps", 2001)) + self.seed = int(self.raw_task.get("seed", 42)) + self.task_position = ( + int(self.raw_task.get("task_idx", task.index.get("index", 0))) + 1 + ) + self.pack_size = int(self.raw_task.get("pack_size", args.get("pack_size", 1))) + if self.pack_size <= 0: + raise ValueError("PDE task pack_size must be positive") + self.trajectory_count = self.pack_size + trajectory_index = (max(1, self.task_position) - 1) % self.trajectory_count + self.current_trajectory_id = f"traj_{trajectory_index}" + self.pack_key = self._pack_key() + + if self.task_position == 1: + self.pack_state = PDEPackState() + self._PACK_STATES[self.pack_key] = self.pack_state + elif self.pack_key in self._PACK_STATES: + self.pack_state = self._PACK_STATES[self.pack_key] + else: + raise ValueError( + "PDE pack tasks must be processed sequentially from task_idx=0" + ) + if self.task_position != self.pack_state.next_task_position: + raise ValueError( + "PDE pack task_idx values must be sequential without gaps or repeats" + ) + self._ensure_hidden_reaction() + self.pack_state.next_task_position += 1 + + self.budget_remaining = self.point_budget + self.done = False + self.current_step = 0 + self.action_feedback: Optional[str] = None + self.final_reward = 0.0 + self.last_sampled_max_error: Optional[float] = None + self.hidden_formula_l1_error: Optional[float] = None + self.hidden_formula_extra_support_penalty: Optional[float] = None + self.hidden_formula_reward_loss: Optional[float] = None + self.last_regression_dataset_id = "merged_all" + self.last_regression_dictionary = [] + self.last_candidate_diagnostics: List[dict] = [] + self.has_task_regression = False + self.has_task_explicit_sampling = False + self.has_task_context_update = False + self.early_termination_by_format_issue = False + self.dump_trajectories = bool(args.get("dump_trajectories", True)) + self.trajectory_dump_path = self._trajectory_dump_path(args) + + def _pack_key(self) -> str: + if "pack_seed" in self.raw_task: + # Runtime-injected pack_seed is the complete pack identity. Ignore + # legacy dataset pack_id so pre-generated rows remain safe when + # shuffled, mixed with another domain, or repacked at a new size. + return f"pde-seed-{self.raw_task['pack_seed']}" + return f"pde-pack-{self.task.batch_id}-{self.task.task_id}" + + def _environment_seed(self) -> int: + """Return an optional benchmark-controlled seed hidden from the model.""" + return int( + self.raw_task.get( + "pde_environment_seed", + self.raw_task.get("pack_seed", self.seed), + ) + ) + + def _support_size_hint(self) -> str: + lower = max(1, min(self.min_reaction_terms, len(self.dictionary_terms))) + upper = max(lower, min(self.max_reaction_terms, len(self.dictionary_terms))) + if lower == upper: + return f"exactly {lower}" + return f"{lower} to {upper}" + + def _ensure_hidden_reaction(self) -> None: + """Create the hidden sparse PDE reaction for this CoD pack. + + The target is deliberately kept out of task rows and prompts. It is + regenerated from the pack seed inside the environment so the agent can + only infer it through sampled-data feedback. + """ + raw_template_index = self.raw_task.get("ground_truth_template_index") + template_index = ( + int(raw_template_index) if raw_template_index is not None else None + ) + environment_seed = self._environment_seed() + if self.pack_state.reaction_coefficients: + if template_index != self.pack_state.ground_truth_template_index: + raise ValueError( + "Every task in a PDE eval pack must use the same " + "ground_truth_template_index" + ) + if environment_seed != self.pack_state.environment_seed: + raise ValueError( + "Every task in a PDE eval pack must use the same " + "pde_environment_seed" + ) + return + pack_seed = environment_seed + seed_sequence = np.random.SeedSequence( + [pack_seed, pde_numeric.GROUND_TRUTH_STREAM] + ) + rng = np.random.default_rng(seed_sequence) + min_terms = max(1, min(self.min_reaction_terms, len(self.dictionary_terms))) + max_terms = max(min_terms, min(self.max_reaction_terms, len(self.dictionary_terms))) + terms: List[str] = [] + coefficients: Dict[str, float] = {} + max_attempts = ground_truth.GT_STABILITY_MAX_ATTEMPTS + for attempt in range(max_attempts): + terms, coefficients = self._sample_hidden_reaction( + rng, + min_terms, + max_terms, + template_index=template_index, + ) + if self._is_stable_reaction_candidate(coefficients): + break + else: + logger.warning( + "No stable hidden PDE reaction found after %d attempts; " + "using the last sampled template candidate.", + max_attempts, + ) + + self.pack_state.reaction_terms = terms + self.pack_state.reaction_coefficients = coefficients + self.pack_state.ground_truth_template_index = template_index + self.pack_state.environment_seed = environment_seed + + def _load_hidden_gt_templates(self) -> List[dict]: + return ground_truth.load_hidden_gt_templates( + self.dictionary_terms, + family=self.ground_truth_family, + ) + + def _sample_hidden_reaction( + self, + rng: np.random.Generator, + min_terms: int, + max_terms: int, + template_index: Optional[int] = None, + ) -> Tuple[List[str], Dict[str, float]]: + return ground_truth.sample_hidden_reaction( + rng, + self.hidden_gt_templates, + min_terms, + max_terms, + template_index=template_index, + ) + + def _is_stable_reaction_candidate( + self, + coefficients: Dict[str, float], + ) -> bool: + return ground_truth.is_stable_reaction_candidate( + coefficients=coefficients, + pde_grid_size=self.pde_grid_size, + pde_time_steps=self.pde_time_steps, + trajectory_count=self.trajectory_count, + pack_seed=self._environment_seed(), + initial_amplitude_upper=self.initial_amplitude_upper, + initial_condition_shape=self.initial_condition_shape, + state_abs_limit=self.pde_state_abs_limit, + ) + + def _trajectory_dump_path(self, args: dict) -> str: + root = args.get("checkpoint_job_dir") or os.path.join( + "logs", "research_cod", "pde_discovery" + ) + dump_dir = os.path.join(str(root), "trajectory_dumps") + raw_uid = str(self.raw_task.get("uid", f"task_{self.task.task_id}")) + safe_uid = re.sub(r"[^A-Za-z0-9_.-]+", "_", raw_uid) + safe_pack = re.sub(r"[^A-Za-z0-9_.-]+", "_", self.pack_key) + file_name = ( + f"{safe_pack}_{safe_uid}_pos{self.task_position}_" + f"pid{os.getpid()}_obj{id(self)}.jsonl" + ) + return os.path.join(dump_dir, file_name) + + def _dump_trajectory_record(self, record: dict) -> None: + if not self.dump_trajectories: + return + try: + os.makedirs(os.path.dirname(self.trajectory_dump_path), exist_ok=True) + payload = { + "uid": self.raw_task.get("uid"), + "pack_key": self.pack_key, + "task_position": self.task_position, + "event": record.get("event"), + "step": record.get("step"), + "budget_remaining": self.budget_remaining, + "latest_regression_equation": self.pack_state.latest_regression_equation, + "current_equation": self.pack_state.preferred_equation, + "reward_snapshot": self.final_reward, + "format_error": self.early_termination_by_format_issue, + **record, + } + with open(self.trajectory_dump_path, "a", encoding="utf-8") as f: + f.write(json.dumps(payload, ensure_ascii=False) + "\n") + except Exception as exc: + logger.warning("Failed to write PDE trajectory dump: %s", exc) + + async def run_async(self) -> List[Experience]: + self.memory.clear() + sys_prompt = load_system_prompt( + dictionary_terms=self.dictionary_terms, + support_size_hint=self._support_size_hint(), + ) + sys_prompt = self._augment_system_prompt(sys_prompt) + self.memory.append({"role": "system", "content": sys_prompt}) + try: + return await super().run_async() + except BaseException: + self._PACK_STATES.pop(self.pack_key, None) + raise + + async def step_async(self, step_num: int) -> Tuple[bool, List[Experience]]: + if self.done: + return False, [] + + user_content = load_user_prompt( + task_position=self.task_position, + pack_size=self.pack_size, + current_step=step_num + 1, + max_steps=self.max_steps, + point_budget=self.point_budget, + budget_remaining=self.budget_remaining, + task_description=self.task_desc or "", + available_datasets=sorted(self.pack_state.datasets.keys()), + last_sampled_max_error=self.last_sampled_max_error, + evidence_memory=self.pack_state.last_context_update.get( + "evidence_memory", {} + ), + action_feedback=self.action_feedback, + ) + if self.icl_examples and step_num == 0: + user_content = ( + f"{user_content}\n\nHere are some reference examples:\n\n" + f"{self.icl_examples}" + ) + + self.memory.append({"role": "user", "content": user_content}) + if self.reply_prefix: + self.memory.append({"role": "assistant", "content": self.reply_prefix}) + + experiences = await self.model.chat_async(self.memory, **self.rollout_args) + response_text = experiences[0].response_text or "" + self.memory.append({"role": "assistant", "content": response_text}) + + sys_prompt = self.memory[0]["content"] if self.memory else "" + for exp in experiences: + exp.info["sys_prompt"] = sys_prompt + exp.info["user_prompt"] = user_content + + action, parse_error = parse_xml_answer(response_text, XML_LIST_TAGS) + if action is None: + response_excerpt = response_text.replace("\n", "\\n")[:500] + logger.warning( + "PDE discovery format error (%s). " + "Response excerpt: %s", + parse_error, + response_excerpt, + ) + for exp in experiences: + exp.info["pde_parse_error"] = parse_error + exp.info["pde_parse_error_response_excerpt"] = response_excerpt + self.action_feedback = ( + "Invalid format: expected exactly one XML action wrapped in " + ".... Game over." + ) + self.final_reward = 0.0 + self.done = True + self.current_step = step_num + 1 + self.early_termination_by_format_issue = True + terminated = True + format_error = True + feedback = self.action_feedback + else: + feedback, terminated, format_error = self._execute_action(action) + self.action_feedback = feedback + self.done = terminated + self.current_step = step_num + 1 + self.early_termination_by_format_issue = format_error + if format_error: + self.final_reward = 0.0 + + self._dump_trajectory_record( + { + "event": "step", + "step": step_num + 1, + "response": response_text, + "action": action, + "action_feedback": feedback, + "terminated": terminated, + "format_error": format_error, + "parse_error": parse_error, + "reward_snapshot": self.final_reward, + } + ) + return not self.done and self.current_step < self.max_steps, experiences + + def _execute_action(self, payload: dict) -> Tuple[str, bool, bool]: + action = payload.get("action") + args = payload.get("args", {}) + if not isinstance(args, dict): + return "Invalid action arguments. Game over.", True, True + + handlers = { + "sample_pde_data": self._action_sample_pde_data, + "summarize_pack_evidence": self._action_summarize_pack_evidence, + "run_sparse_regression": self._action_run_sparse_regression, + "update_scientific_context": self._action_update_scientific_context, + } + handler = handlers.get(action) + if handler is None: + return f"Invalid action: unknown action '{action}'. Game over.", True, True + return handler(args) + + def _action_summarize_pack_evidence(self, args: dict) -> Tuple[str, bool, bool]: + reveal_value = args.get("reveal_equations", "false") + if str(reveal_value).lower() not in {"true", "false"}: + return "Error: reveal_equations must be true or false.", False, False + reveal_equations = str(reveal_value).lower() == "true" + history = list(self.pack_state.evidence_history) + if not history: + return ( + "Pack evidence summary: no prior task evidence has been recorded " + "in this pack. Start with current-task measurements.", + False, + False, + ) + + slots: Dict[str, dict] = {} + for entry in history: + equation = str(entry.get("preferred_equation") or "").strip() + if not equation: + continue + if equation not in slots: + slots[equation] = { + "slot_id": f"H{len(slots) + 1}", + "equation": equation, + "tasks": set(), + } + slot = slots[equation] + if entry.get("task_position") is not None: + slot["tasks"].add(str(entry["task_position"])) + + if not slots: + return ( + "Pack evidence summary: prior notes exist, but no tested " + "candidate slots are available. Use current-task actions to build evidence.", + False, + False, + ) + + lines = [ + "Pack evidence summary; historical hypotheses, not ground truth.", + ] + for equation, slot in slots.items(): + pieces = [ + f"{slot['slot_id']}: tasks={','.join(sorted(slot['tasks'])) or 'unknown'}", + ] + if reveal_equations: + pieces.append(f"agent_recorded_equation={slot['equation']}") + lines.append("; ".join(pieces) + ".") + + if not reveal_equations: + lines.append( + "Call summarize_pack_evidence with reveal_equations=true only if " + "you need to actively reuse prior agent-recorded candidates." + ) + return "\n".join(lines), False, False + + def _action_sample_pde_data(self, args: dict) -> Tuple[str, bool, bool]: + has_points = args.get("points") is not None + has_point_grid = args.get("point_grid") is not None + if has_points: + return ( + "Error: explicit points are not supported. Use point_grid with " + "one point element per numeric x and t coordinate pair.", + False, + False, + ) + grid_points = _normalize_point_grid(args.get("point_grid")) + if not has_point_grid: + return ( + "Error: sample_pde_data supports only point_grid. Provide " + "one point element per numeric x and t coordinate pair.", + False, + False, + ) + if has_point_grid and grid_points is None: + return ( + "Error: each point in point_grid must contain numeric x and t " + "values with 0 < x < 1 and 0 <= t <= 1.", + False, + False, + ) + + grid_sample_points = grid_points or [] + num_points = len(grid_sample_points) + if num_points <= 0: + return "Error: point_grid must produce at least one point.", False, False + if num_points > self.budget_remaining: + return ( + f"Error: requested {num_points} points but only " + f"{self.budget_remaining} budget remains.", + False, + False, + ) + + trajectory_id = self.current_trajectory_id + self.pack_state.sample_counter += 1 + dataset_id = f"ds_t{self.task_position}_{self.pack_state.sample_counter}" + data = self._sample_grid_points(grid_sample_points, trajectory_id) + x_values = data["x"] + t_values = data["t"] + region = { + "x": [float(np.min(x_values)), float(np.max(x_values))], + "t": [float(np.min(t_values)), float(np.max(t_values))], + } + sampling_mode = "point_grid" + self.has_task_explicit_sampling = True + data["region"] = region + data["sampling_strategy"] = sampling_mode + data["task_position"] = self.task_position + data["trajectory_id"] = np.array([trajectory_id] * num_points, dtype=object) + self.pack_state.datasets[dataset_id] = data + self.budget_remaining -= num_points + + t = data["t"] + summary = ( + f"Sampled {num_points} points. Dataset: {dataset_id}. " + f"Point budget remaining: {self.budget_remaining}/{self.point_budget}. " + f"Sampling mode: {sampling_mode}. " + f"Region: x in [{region['x'][0]:.3g}, {region['x'][1]:.3g}], " + f"t in [{region['t'][0]:.3g}, {region['t'][1]:.3g}]. " + f"mean t = {np.mean(t):.3g}. " + f"Dataset coverage: {self._coverage_report(data)}." + ) + return summary, False, False + + def _action_run_sparse_regression(self, args: dict) -> Tuple[str, bool, bool]: + dataset_id = str(args.get("dataset_id", "merged_all")) + data = self._resolve_dataset(dataset_id) + if data is None: + return f"Error: dataset_id '{dataset_id}' was not found.", False, False + + raw_dictionary = args.get("dictionary") + raw_candidates = args.get("candidate_equations") + has_candidates = isinstance(raw_candidates, list) and bool(raw_candidates) + if has_candidates and len(raw_candidates) > 8: + return "Error: candidate_equations accepts at most 8 equations.", False, False + if raw_dictionary is None and not has_candidates: + return ( + "Error: run_sparse_regression requires an agent-specified " + "dictionary or candidate_equations.", + False, + False, + ) + if raw_dictionary is not None and ( + not isinstance(raw_dictionary, list) or not raw_dictionary + ): + return ( + "Error: dictionary must be a non-empty list of supported terms.", + False, + False, + ) + dictionary = self._sanitize_dictionary(raw_dictionary) + if raw_dictionary is not None and not any( + str(term) in self.dictionary_terms for term in raw_dictionary + ): + return ( + "Error: dictionary did not contain any supported basis terms.", + False, + False, + ) + allowed_supports = None + if has_candidates: + candidate_strings = [ + str(candidate).strip() + for candidate in raw_candidates[:8] + if str(candidate).strip() + ] + allowed_supports = [] + for candidate in candidate_strings: + support = self._candidate_support(candidate, dictionary) + if support: + allowed_supports.append(support) + if not allowed_supports: + return ( + "Error: candidate_equations did not contain any supported " + "dictionary terms.", + False, + False, + ) + if raw_dictionary is None: + support_terms = { + term for support in allowed_supports for term in support + } + dictionary = [ + term + for term in self.dictionary_terms + if term in support_terms + ] + if not dictionary: + return ( + "Error: run_sparse_regression has no selected supported terms " + "to score.", + False, + False, + ) + alpha_arg = args.get("alpha") + threshold_arg = args.get("threshold") + try: + alpha = float(0.05 if alpha_arg is None else alpha_arg) + threshold = float(0.05 if threshold_arg is None else threshold_arg) + except (TypeError, ValueError): + return ( + "Error: alpha and threshold must be numeric when provided.", + False, + False, + ) + if not all( + math.isfinite(value) and value >= 0.0 for value in (alpha, threshold) + ): + return ( + "Error: alpha and threshold must be finite non-negative numbers.", + False, + False, + ) + result = self._regression_result( + data["u"], + data["y"], + data["u"], + data["y"], + dictionary, + alpha, + threshold, + allowed_supports=allowed_supports, + ) + result["feedback"] += ( + f" Fit dataset coverage: {self._coverage_report(data)}. " + "Objective scope: current sampled dataset only; use it as a local " + "fit and conditioning check before committing an equation. " + f"{self._point_error_report(data, result['coefficients'])}" + ) + self.last_sampled_max_error = result["penalty"] + self.last_regression_dataset_id = dataset_id + self.last_regression_dictionary = list(dictionary) + self.last_candidate_diagnostics = result.get("candidate_diagnostics", []) + self.pack_state.latest_regression_equation = result["equation"] + self.has_task_regression = True + return result["feedback"], False, False + + def _task_local_observation_stats( + self, + data: Optional[dict], + ) -> dict: + if data is None: + return {"n": 0, "u_span": 0.0} + u_values = np.asarray(data.get("u", []), dtype=float) + u_values = u_values[np.isfinite(u_values)] + if len(u_values) == 0: + return {"n": 0, "u_span": 0.0} + return { + "n": int(len(u_values)), + "u_span": float(np.max(u_values) - np.min(u_values)), + } + + def _structured_evidence_memory( + self, + *, + preferred_diagnostics: dict, + observation_stats: dict, + ) -> dict: + preferred_error = float( + preferred_diagnostics.get("training_max_error", KAPPA_CAP) + ) + + return { + "task_evidence_status": "recorded", + "observation_summary": { + "n": observation_stats.get("n", 0), + "u_span": observation_stats.get("u_span", 0.0), + "relative_margin": preferred_diagnostics.get("relative_margin", 0.0), + "preferred_sampled_max_error": preferred_error, + }, + } + + def _action_update_scientific_context(self, args: dict) -> Tuple[str, bool, bool]: + short_note = re.sub(r"\s+", " ", str(args.get("note", "") or "")).strip() + preferred_equation = str(args.get("preferred_equation", "")).strip() + uncertain_terms = args.get("uncertain_terms", []) + if not isinstance(uncertain_terms, list): + uncertain_terms = [str(uncertain_terms)] + if len(uncertain_terms) > 4: + return "Error: uncertain_terms accepts at most 4 terms.", False, False + uncertain_terms = [ + re.sub(r"\s+", " ", str(term or "")).strip() + for term in uncertain_terms + if str(term or "").strip() + ] + if not preferred_equation: + return ( + "Error: update_scientific_context requires preferred_equation. " + "Choose the equation you want to finalize from the sampled-data " + "diagnostics. This may be the " + "latest regression candidate or another candidate scored by " + "run_sparse_regression.", + False, + False, + ) + for field_name, equation in ( + ("preferred_equation", preferred_equation), + ): + symbolic_tokens = candidate_utils.symbolic_coefficient_tokens(equation) + if symbolic_tokens: + tokens = ", ".join(symbolic_tokens) + return ( + f"Error: {field_name} contains symbolic coefficient(s): " + f"{tokens}. update_scientific_context requires explicit " + "numeric coefficients from a tested candidate, not " + "placeholders such as c1/c2/a/b.", + False, + False, + ) + if not self.has_task_explicit_sampling: + return ( + "Error: update_scientific_context requires current-task " + "sample_pde_data evidence in the current task. Choose " + "point_grid coordinates, sample them, then run sparse " + "regression before updating scientific context.", + False, + False, + ) + if not self.has_task_regression: + return ( + "Error: update_scientific_context requires at least one " + "run_sparse_regression call in the current task. Run regression " + "on an available dataset first so the task has a current " + "equation and sampled-data evidence.", + False, + False, + ) + preferred_coefficients = self._candidate_coefficients(preferred_equation) + if not preferred_coefficients: + return ( + "Error: preferred_equation did not parse into dictionary terms. " + "Use an explicit equation built from supported dictionary terms " + "with numeric coefficients.", + False, + False, + ) + evaluated_candidates = [ + self._candidate_coefficients(str(item.get("equation", ""))) + for item in self.last_candidate_diagnostics + ] + if not any( + all( + math.isclose( + preferred_coefficients.get(term, 0.0), + candidate.get(term, 0.0), + rel_tol=1e-9, + abs_tol=1e-12, + ) + for term in self.dictionary_terms + ) + for candidate in evaluated_candidates + ): + return ( + "Error: preferred_equation must match a candidate scored by the " + "latest run_sparse_regression call.", + False, + False, + ) + data = self._resolve_dataset(self.last_regression_dataset_id) + if data is None: + data = self._resolve_dataset("merged_all") + preferred_diagnostics = self._candidate_objective_diagnostics( + preferred_equation, + data, + self.last_regression_dictionary, + ) + preferred_support = list(preferred_diagnostics["support"]) + observation_stats = self._task_local_observation_stats(data) + reward_diagnostics = preferred_diagnostics + self.last_sampled_max_error = float(reward_diagnostics["step_penalty"]) + hidden_reward_diagnostics = self._hidden_formula_reward_diagnostics( + preferred_equation + ) + self.hidden_formula_l1_error = float( + hidden_reward_diagnostics["normalized_l1_error"] + ) + self.hidden_formula_extra_support_penalty = float( + hidden_reward_diagnostics["extra_support_penalty"] + ) + self.hidden_formula_reward_loss = float(hidden_reward_diagnostics["loss"]) + self.pack_state.preferred_equation = preferred_equation + self.pack_state.preferred_support = preferred_support + self.pack_state.preferred_coefficients = dict( + preferred_diagnostics["coefficients"] + ) + self.final_reward = float(hidden_reward_diagnostics["reward"]) + evidence_memory = self._structured_evidence_memory( + preferred_diagnostics=preferred_diagnostics, + observation_stats=observation_stats, + ) + preferred_diagnostics_summary = { + "training_max_error": reward_diagnostics["training_max_error"], + "condition_penalty": reward_diagnostics["condition_penalty"], + "relative_margin": reward_diagnostics["relative_margin"], + "second_best_equation": reward_diagnostics["second_best_equation"], + "step_penalty": reward_diagnostics["step_penalty"], + } + self.pack_state.last_context_update = { + "evidence_memory": evidence_memory, + "preferred_equation": preferred_equation, + "preferred_diagnostics": preferred_diagnostics_summary, + "latest_regression_equation": self.pack_state.latest_regression_equation, + "note": short_note, + "uncertain_terms": [str(term) for term in uncertain_terms], + } + self.pack_state.evidence_history.append( + { + "task_position": self.task_position, + "evidence_memory": dict(evidence_memory), + "preferred_equation": preferred_equation, + } + ) + self.pack_state.evidence_history = self.pack_state.evidence_history[-self.pack_size :] + self.has_task_context_update = True + return ( + "Task terminated. " + f"reward_equation={preferred_equation}; " + f"sampled-data max error={reward_diagnostics['training_max_error']:.3g}; " + f"condition penalty={reward_diagnostics['condition_penalty']:.3g}; " + f"closest different-support candidate={reward_diagnostics['second_best_equation']}; " + f"relative margin={reward_diagnostics['relative_margin']:.3g}.", + True, + False, + ) + + def _sample_grid_points( + self, + points: List[Tuple[float, float]], + trajectory_id: str, + ) -> dict: + seed_sequence = np.random.SeedSequence( + [ + self.seed, + self.task_position, + self.pack_state.sample_counter, + ] + ) + rng = np.random.default_rng(seed_sequence) + self._ensure_dense_field() + x = np.array([point[0] for point in points], dtype=float) + t = np.array([point[1] for point in points], dtype=float) + return self._sample_values_at_points(x, t, trajectory_id, rng) + + def _sample_values_at_points( + self, + x: np.ndarray, + t: np.ndarray, + trajectory_id: str, + rng: np.random.Generator, + ) -> dict: + # The regression target is the scheme-consistent reaction residual, + # sampled from the dense hidden PDE field rather than direct oracle f(u). + u = self._interpolate_dense_field("u", x, t, trajectory_id) + y = self._interpolate_dense_field("y", x, t, trajectory_id) + point_count = len(x) + u += rng.normal(0.0, self.noise_level, point_count) + y += rng.normal(0.0, self.noise_level * 0.25, point_count) + return {"x": x, "t": t, "u": u, "y": y} + + def _ensure_dense_field(self) -> None: + if self.pack_state.dense_field is not None: + return + + nx = max(17, self.pde_grid_size) + nt = max(51, self.pde_time_steps) + if nx % 2 == 0: + nx += 1 + x_grid = np.linspace(0.0, 1.0, nx) + t_grid = np.linspace(0.0, 1.0, nt) + dx = float(x_grid[1] - x_grid[0]) + dt = float(t_grid[1] - t_grid[0]) + r = dt / (dx * dx) + + interior = nx - 2 + lower = -r * np.ones(interior - 1, dtype=float) + diag = (1.0 + 2.0 * r) * np.ones(interior, dtype=float) + upper = -r * np.ones(interior - 1, dtype=float) + + trajectories = {} + for traj_idx in range(self.trajectory_count): + trajectory_id = f"traj_{traj_idx}" + trajectories[trajectory_id] = self._simulate_dense_trajectory( + x_grid, + t_grid, + dx, + dt, + lower, + diag, + upper, + traj_idx, + ) + + default_field = trajectories["traj_0"] + self.pack_state.dense_field = { + "trajectories": trajectories, + "x_grid": x_grid, + "t_grid": t_grid, + "u": default_field["u"], + "y": default_field["y"], + } + + def _simulate_dense_trajectory( + self, + x_grid: np.ndarray, + t_grid: np.ndarray, + dx: float, + dt: float, + lower: np.ndarray, + diag: np.ndarray, + upper: np.ndarray, + trajectory_index: int, + ) -> dict: + return pde_numeric.simulate_dense_trajectory( + x_grid=x_grid, + t_grid=t_grid, + dx=dx, + dt=dt, + lower=lower, + diag=diag, + upper=upper, + trajectory_index=trajectory_index, + reaction_fn=self._reaction_value, + pack_seed=self._environment_seed(), + trajectory_count=self.trajectory_count, + initial_amplitude_upper=self.initial_amplitude_upper, + initial_condition_shape=self.initial_condition_shape, + state_abs_limit=self.pde_state_abs_limit, + ) + + def _interpolate_dense_field( + self, + field_name: str, + x: np.ndarray, + t: np.ndarray, + trajectory_id: str = "traj_0", + ) -> np.ndarray: + self._ensure_dense_field() + return pde_numeric.interpolate_dense_field( + self._trajectory_field(trajectory_id), + field_name, + x, + t, + ) + + def _trajectory_field(self, trajectory_id: str) -> dict: + self._ensure_dense_field() + dense_field = self.pack_state.dense_field or {} + trajectories = dense_field.get("trajectories") + if isinstance(trajectories, dict) and trajectory_id in trajectories: + return trajectories[trajectory_id] + return dense_field + + def _sanitize_dictionary(self, dictionary: object) -> List[str]: + return candidate_utils.sanitize_dictionary(dictionary, self.dictionary_terms) + + def _reaction_value( + self, u: np.ndarray, coefficients: Optional[Dict[str, float]] = None + ) -> np.ndarray: + return candidate_utils.reaction_value( + u, + coefficients or self.pack_state.reaction_coefficients, + ) + + def _resolve_dataset(self, dataset_id: str) -> Optional[dict]: + if dataset_id in self.pack_state.datasets: + return self.pack_state.datasets[dataset_id] + if dataset_id == "merged_all" or dataset_id == "merged": + ids = sorted(self.pack_state.datasets) + else: + match = re.fullmatch(r"merged_t(\d+)_t(\d+)", dataset_id) + if not match: + return None + lo, hi = int(match.group(1)), int(match.group(2)) + ids = [ + dsid + for dsid, data in self.pack_state.datasets.items() + if lo <= int(data.get("task_position", 0)) <= hi + ] + if not ids: + return None + merged = { + key: np.concatenate([self.pack_state.datasets[dsid][key] for dsid in ids]) + for key in ["x", "t", "u", "y"] + } + if all("trajectory_id" in self.pack_state.datasets[dsid] for dsid in ids): + merged["trajectory_id"] = np.concatenate( + [self.pack_state.datasets[dsid]["trajectory_id"] for dsid in ids] + ) + return merged + + def _coverage_report(self, data: dict) -> str: + u = np.asarray(data.get("u", []), dtype=float) + if len(u) == 0: + return "empty dataset" + x = np.asarray(data.get("x", []), dtype=float) + t = np.asarray(data.get("t", []), dtype=float) + pieces = [ + f"n={len(u)}", + f"u=[{np.min(u):.3g},{np.max(u):.3g}]", + ] + if len(x): + pieces.append(f"x_span={np.max(x) - np.min(x):.3g}") + if len(t): + pieces.append(f"t_span={np.max(t) - np.min(t):.3g}") + return "; ".join(pieces) + + def _point_error_report(self, data: dict, coefficients: Dict[str, float]) -> str: + x = np.asarray(data.get("x", []), dtype=float) + t = np.asarray(data.get("t", []), dtype=float) + u = np.asarray(data.get("u", []), dtype=float) + y = np.asarray(data.get("y", []), dtype=float) + if not (len(x) == len(t) == len(u) == len(y)) or len(u) == 0: + return "Sampled-point errors: unavailable." + pred = self._reaction_value(u, coefficients=coefficients) + errors = np.abs(y - pred) + rows = [ + ( + f"(x={float(x_i):.4g}, t={float(t_i):.4g}, " + f"u={float(u_i):.4g}, y={float(y_i):.4g}, " + f"pred={float(pred_i):.4g}, abs_error={float(err_i):.4g})" + ) + for x_i, t_i, u_i, y_i, pred_i, err_i in zip(x, t, u, y, pred, errors) + ] + return "Sampled-point errors: " + "; ".join(rows) + "." + + def _hidden_formula_reward_diagnostics(self, candidate: str) -> dict: + """Compute the training-only reward from hidden coefficient vectors.""" + pred_coefficients = self._candidate_coefficients(candidate) + target_coefficients = self.pack_state.reaction_coefficients + target_l1 = sum(abs(float(value)) for value in target_coefficients.values()) + diff_l1 = 0.0 + for term in self.dictionary_terms: + diff_l1 += abs( + float(pred_coefficients.get(term, 0.0)) + - float(target_coefficients.get(term, 0.0)) + ) + normalized_l1 = diff_l1 / max(1e-8, target_l1) + pred_support_size = sum( + 1 for value in pred_coefficients.values() if abs(float(value)) > 1e-12 + ) + target_support_size = len(self.pack_state.reaction_terms) + extra_support_penalty = float( + max(0, pred_support_size - target_support_size) + ) + loss = float(normalized_l1 + extra_support_penalty) + return { + "normalized_l1_error": float(normalized_l1), + "extra_support_penalty": extra_support_penalty, + "loss": loss, + "reward": float(1.0 / (1.0 + loss)), + } + + def _regression_result( + self, + u: np.ndarray, + y: np.ndarray, + u_objective: np.ndarray, + y_objective: np.ndarray, + dictionary: list, + alpha: float, + threshold: float, + allowed_supports: Optional[List[List[str]]] = None, + ) -> dict: + return regression.sparse_regression_result( + u=u, + y=y, + u_objective=u_objective, + y_objective=y_objective, + dictionary=dictionary, + alpha=alpha, + threshold=threshold, + max_reaction_terms=self.max_reaction_terms, + default_blind_penalty=self.default_blind_penalty, + kappa_threshold=self.kappa_threshold, + allowed_supports=allowed_supports, + ) + + def _candidate_coefficients(self, candidate: str) -> Dict[str, float]: + return candidate_utils.candidate_coefficients( + candidate, + self.dictionary_terms, + ) + + def _candidate_support(self, candidate: str, dictionary: List[str]) -> List[str]: + return candidate_utils.candidate_support( + candidate, + dictionary, + self.dictionary_terms, + ) + + def _candidate_objective_diagnostics( + self, + candidate: str, + sampled_data: dict, + dictionary: List[str], + ) -> dict: + return regression.candidate_objective_diagnostics( + candidate=candidate, + data=sampled_data, + objective_data=sampled_data, + dictionary=dictionary, + dictionary_terms=self.dictionary_terms, + last_candidate_diagnostics=self.last_candidate_diagnostics, + kappa_threshold=self.kappa_threshold, + ) + + def _get_feedback(self) -> str: + if self.early_termination_by_format_issue: + return self.action_feedback or "Invalid action." + return "\n".join( + [ + "Latest sampled-data diagnostic: " + f"{self.last_sampled_max_error if self.last_sampled_max_error is not None else 'none'}", + ( + "Latest finalized equation: " + f"{self.pack_state.preferred_equation or 'none'}" + ), + ( + "Latest regression candidate: " + f"{self.pack_state.latest_regression_equation or 'none'}" + ), + ] + ) + + def _hidden_accuracy_metrics(self) -> dict: + if not self.has_task_context_update: + return { + "structure_recovery": 0.0, + "coefficient_relative_error": 1.0, + } + coeffs = self.pack_state.preferred_coefficients + target = self.pack_state.reaction_coefficients + target_support = set(self.pack_state.reaction_terms) + support = self.pack_state.preferred_support + structure_recovery = float(set(support) == target_support) + coefficient_error = 1.0 + if structure_recovery: + coefficient_error = sum( + abs(coeffs.get(term, 0.0) - target[term]) + / max(1e-6, abs(target[term])) + for term in target_support + ) / max(1, len(target_support)) + return { + "structure_recovery": structure_recovery, + "coefficient_relative_error": coefficient_error, + } + + async def reward_async(self, exps: List[Experience]) -> float: + if not self.early_termination_by_format_issue and not self.has_task_context_update: + self.final_reward = 0.0 + if self.last_sampled_max_error is None: + self.last_sampled_max_error = self.default_blind_penalty + self.action_feedback = ( + self.action_feedback or "" + ) + "\nTask ended without update_scientific_context; no scientific context was committed." + hidden_metrics = self._hidden_accuracy_metrics() + reward = await super().reward_async(exps) + if not exps: + if self.task_position >= self.pack_size: + self._PACK_STATES.pop(self.pack_key, None) + return reward + + metrics = exps[-1].metrics or {} + metrics.update( + { + "pde_structure_recovery": hidden_metrics["structure_recovery"], + "pde_coefficient_relative_error": hidden_metrics[ + "coefficient_relative_error" + ], + "pde_sample_efficiency_points": float( + self.point_budget - self.budget_remaining + ), + "pde_last_sampled_max_error": float( + self.last_sampled_max_error + if self.last_sampled_max_error is not None + else self.default_blind_penalty + ), + "pde_hidden_formula_normalized_l1_error": float( + self.hidden_formula_l1_error + if self.hidden_formula_l1_error is not None + else self.default_blind_penalty + ), + "pde_hidden_formula_extra_support_penalty": float( + self.hidden_formula_extra_support_penalty + if self.hidden_formula_extra_support_penalty is not None + else 0.0 + ), + "pde_hidden_formula_reward_loss": float( + self.hidden_formula_reward_loss + if self.hidden_formula_reward_loss is not None + else self.default_blind_penalty + ), + "pde_ground_truth_support_size": float( + len(self.pack_state.reaction_terms) + ), + "pde_ground_truth_template_index": float( + self.pack_state.ground_truth_template_index + if self.pack_state.ground_truth_template_index is not None + else -1 + ), + "pde_ground_truth_template_instance": float( + self.raw_task.get("ground_truth_template_instance", -1) + ), + "pde_eval_pack_index": float( + self.raw_task.get("eval_pack_index", -1) + ), + } + ) + support_size = len(self.pack_state.reaction_terms) + metrics.update( + { + f"pde_structure_recovery_support_size_{support_size}": hidden_metrics[ + "structure_recovery" + ], + f"pde_hidden_formula_l1_support_size_{support_size}": float( + self.hidden_formula_l1_error + if self.hidden_formula_l1_error is not None + else self.default_blind_penalty + ), + f"pde_hidden_formula_reward_support_size_{support_size}": float( + self.final_reward + ), + } + ) + template_index = self.pack_state.ground_truth_template_index + if template_index is not None: + metrics.update( + { + f"pde_structure_recovery_template_{template_index}": hidden_metrics[ + "structure_recovery" + ], + f"pde_hidden_formula_l1_template_{template_index}": float( + self.hidden_formula_l1_error + if self.hidden_formula_l1_error is not None + else self.default_blind_penalty + ), + f"pde_hidden_formula_reward_template_{template_index}": float( + self.final_reward + ), + } + ) + exps[-1].metrics = metrics + self._dump_trajectory_record( + { + "event": "final", + "step": self.current_step, + "reward": reward, + "final_reward": self.final_reward, + "feedback": self._get_feedback(), + "metrics": dict(metrics), + } + ) + if self.task_position >= self.pack_size: + self._PACK_STATES.pop(self.pack_key, None) + return reward + + def _compress_memory(self) -> None: + super()._compress_memory() + old_feedback_cutoff = max(1, self.current_step - 1) + seen_user_messages = 0 + for msg in self.memory: + if msg.get("role") != "user": + continue + seen_user_messages += 1 + content = msg.get("content", "") + if seen_user_messages >= old_feedback_cutoff or "Action feedback:" not in content: + continue + msg["content"] = ( + content.split("Action feedback:")[0] + + "Action feedback:\n[History compressed.]" + ) + + @property + def max_step_num(self) -> int: + return self.max_steps diff --git a/trinity/common/workflows/connect_the_dots/terminal/__init__.py b/trinity/common/workflows/connect_the_dots/terminal/__init__.py new file mode 100644 index 00000000000..0e5518b1bc8 --- /dev/null +++ b/trinity/common/workflows/connect_the_dots/terminal/__init__.py @@ -0,0 +1,2 @@ +# -*- coding: utf-8 -*- +"""Simulated terminal environment for RL training of file operation skills.""" diff --git a/trinity/common/workflows/connect_the_dots/terminal/commands.py b/trinity/common/workflows/connect_the_dots/terminal/commands.py new file mode 100644 index 00000000000..ec7a2df5ab1 --- /dev/null +++ b/trinity/common/workflows/connect_the_dots/terminal/commands.py @@ -0,0 +1,1101 @@ +# -*- coding: utf-8 -*- +""" +Command handlers for the simulated terminal environment. + +Each handler is a class with an ``execute(args, env)`` method that returns +a string (the terminal output). Handlers read/write the virtual filesystem +via ``env.current_machine.fs`` and ``env.current_machine`` state. +""" + +from __future__ import annotations + +import json +from abc import ABC, abstractmethod +from typing import TYPE_CHECKING, Dict, List, Set + +from trinity.common.workflows.connect_the_dots.terminal.env import OSType + +if TYPE_CHECKING: + from trinity.common.workflows.connect_the_dots.terminal.env import TerminalEnv + + +# --------------------------------------------------------------------------- +# Base +# --------------------------------------------------------------------------- + +class CommandHandler(ABC): + """Base class for command handlers.""" + + name: str = "" + # OS where this command is available. Empty set = all OS. + available_os: Set[OSType] = set() + # If True, command can only run when NOT ssh-connected (on local machine). + local_only: bool = False + + def is_available(self, os_type: OSType, ssh_connected: bool) -> bool: + if self.local_only and ssh_connected: + return False + if ssh_connected: + # Remote is always Linux + return True + if self.available_os and os_type not in self.available_os: + return False + return True + + @abstractmethod + def execute(self, args: List[str], env: TerminalEnv) -> str: + ... + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +def _machine(env: TerminalEnv): + return env.current_machine + + +def _resolve(env: TerminalEnv, path: str) -> str: + return _machine(env).resolve_path(path) + + +def _display(env: TerminalEnv, memfs_path: str) -> str: + return _machine(env).to_display_path(memfs_path) + + +def _format_size(size: int) -> str: + if size < 1024: + return f"{size}" + elif size < 1024 * 1024: + return f"{size / 1024:.1f}K" + else: + return f"{size / (1024 * 1024):.1f}M" + + +def _parse_scp_target(arg: str): + """Parse ``user@host:path``, ``user@host:``, or ``user@host``. Returns (user, host, path) or (None, None, path).""" + if ":" in arg and "@" in arg.split(":")[0]: + user_host, path = arg.split(":", 1) + user, host = user_host.split("@", 1) + return user, host, path + # user@host without colon — treat as remote home directory + if "@" in arg and "/" not in arg and "\\" not in arg: + user, host = arg.split("@", 1) + return user, host, "" + return None, None, arg + + +# --------------------------------------------------------------------------- +# ls +# --------------------------------------------------------------------------- + +class LsHandler(CommandHandler): + name = "ls" + + def execute(self, args: List[str], env: TerminalEnv) -> str: + m = _machine(env) + fs = m.fs + + show_long = False + show_all = False + paths = [] + for a in args: + if a.startswith("--"): + if a == "--all": + show_all = True + elif a == "--long": + show_long = True + elif a.startswith("-"): + for ch in a[1:]: + if ch == "l": + show_long = True + elif ch == "a": + show_all = True + else: + paths.append(a) + + if not paths: + paths = [m.to_display_path(m.cwd)] + + all_output = [] + for p in paths: + mp = _resolve(env, p) + if not fs.exists(mp): + all_output.append(f"ls: cannot access '{p}': No such file or directory") + continue + if fs.isfile(mp): + if show_long: + all_output.append(self._long_entry(fs, mp, mp.split("/")[-1])) + else: + all_output.append(mp.split("/")[-1]) + continue + try: + entries = fs.listdir(mp) + except Exception: + all_output.append(f"ls: cannot access '{p}': No such file or directory") + continue + + if not show_all: + entries = [e for e in entries if not e.startswith(".")] + + # Truncate long listings + truncated = False + if len(entries) > 50: + total = len(entries) + entries = entries[:50] + truncated = True + + if show_long: + lines = [] + for name in entries: + child = f"{mp}/{name}" if mp else name + lines.append(self._long_entry(fs, child, name)) + if truncated: + lines.append(f"... ({total - 50} more entries)") + all_output.append("\n".join(lines)) + else: + if m.os_type == OSType.WINDOWS and not env.ssh_connected: + all_output.append(self._dir_format(fs, mp, entries, truncated, + total if truncated else len(entries))) + else: + line = " ".join( + f"{name}/" if fs.isdir(f"{mp}/{name}" if mp else name) else name + for name in entries + ) + if truncated: + line += f"\n... ({total - 50} more entries)" + all_output.append(line) + + return "\n".join(all_output) + + def _long_entry(self, fs, memfs_path, name): + meta = fs.get_meta(memfs_path) + is_dir = fs.isdir(memfs_path) + perm_str = "d" if is_dir else "-" + mode = int(meta.permissions, 8) if meta.permissions.isdigit() else 0o644 + for shift in (6, 3, 0): + bits = (mode >> shift) & 7 + perm_str += "r" if bits & 4 else "-" + perm_str += "w" if bits & 2 else "-" + perm_str += "x" if bits & 1 else "-" + size = meta.size if not is_dir else 4096 + return f"{perm_str} 1 {meta.owner} {meta.group} {size:>8} Jan 15 10:30 {name}{'/' if is_dir else ''}" + + def _dir_format(self, fs, mp, entries, truncated, total_count): + """Windows dir format.""" + dir_display = _machine_stub_display(mp) + lines = [ + " Volume in drive C has no label.", + f" Directory of {dir_display}", + "", + ] + file_count = 0 + dir_count = 0 + total_size = 0 + for name in entries: + child = f"{mp}/{name}" if mp else name + is_dir = fs.isdir(child) + meta = fs.get_meta(child) + if is_dir: + dir_count += 1 + lines.append(f"01/15/2024 10:30 AM {name}") + else: + file_count += 1 + total_size += meta.size + lines.append(f"01/15/2024 10:30 AM {meta.size:>8} {name}") + if truncated: + lines.append(f"... ({total_count - 50} more entries)") + lines.append(f" {file_count} File(s) {total_size:>10} bytes") + lines.append(f" {dir_count} Dir(s) 500,000,000 bytes free") + return "\n".join(lines) + + +def _machine_stub_display(mp): + """Quick display for dir header.""" + parts = mp.split("/") + if parts and len(parts[0]) == 1 and parts[0].isalpha(): + drive = parts[0].upper() + return f"{drive}:\\" + "\\".join(parts[1:]) + return "/" + mp + + +# --------------------------------------------------------------------------- +# cd +# --------------------------------------------------------------------------- + +class CdHandler(CommandHandler): + name = "cd" + + def execute(self, args: List[str], env: TerminalEnv) -> str: + m = _machine(env) + if not args: + if m.os_type == OSType.WINDOWS and not env.ssh_connected: + return m.to_display_path(m.cwd) + m.cwd = m.home_dir + self._update_pwd(m) + return "" + target = args[0] + mp = _resolve(env, target) + if not m.fs.exists(mp): + if m.os_type == OSType.WINDOWS and not env.ssh_connected: + return "The system cannot find the path specified." + return f"bash: cd: {target}: No such file or directory" + if not m.fs.isdir(mp): + if m.os_type == OSType.WINDOWS and not env.ssh_connected: + return "The directory name is invalid." + return f"bash: cd: {target}: Not a directory" + m.cwd = mp + self._update_pwd(m) + if m.os_type == OSType.WINDOWS and not env.ssh_connected: + return m.to_display_path(m.cwd) + return "" + + @staticmethod + def _update_pwd(m): + display = m.to_display_path(m.cwd) + if m.os_type == OSType.WINDOWS: + m.env_vars["CD"] = display + else: + m.env_vars["PWD"] = display + + +# --------------------------------------------------------------------------- +# pwd +# --------------------------------------------------------------------------- + +class PwdHandler(CommandHandler): + name = "pwd" + + def execute(self, args: List[str], env: TerminalEnv) -> str: + return _machine(env).to_display_path(_machine(env).cwd) + + +# --------------------------------------------------------------------------- +# cat +# --------------------------------------------------------------------------- + +class CatHandler(CommandHandler): + name = "cat" + + def execute(self, args: List[str], env: TerminalEnv) -> str: + if not args: + return "cat: missing operand" + m = _machine(env) + outputs = [] + for f in args: + mp = _resolve(env, f) + if not m.fs.exists(mp): + outputs.append(f"cat: {f}: No such file or directory") + elif m.fs.isdir(mp): + outputs.append(f"cat: {f}: Is a directory") + else: + meta = m.fs.get_meta(mp) + if meta.archive_type: + outputs.append(f"[Binary file - {meta.archive_type} archive, " + f"{len(meta.archive_entries or {})} entries]") + else: + outputs.append(m.fs.readtext(mp)) + return "\n".join(outputs) + + +# --------------------------------------------------------------------------- +# mkdir +# --------------------------------------------------------------------------- + +class MkdirHandler(CommandHandler): + name = "mkdir" + + def execute(self, args: List[str], env: TerminalEnv) -> str: + m = _machine(env) + create_parents = False + paths = [] + for a in args: + if a in ("-p", "--parents"): + create_parents = True + else: + paths.append(a) + if not paths: + return "mkdir: missing operand" + outputs = [] + for p in paths: + mp = _resolve(env, p) + try: + if create_parents: + m.fs.makedirs(mp) + else: + m.fs.mkdir(mp) + except FileExistsError: + outputs.append(f"mkdir: cannot create directory '{p}': File exists") + except FileNotFoundError: + outputs.append(f"mkdir: cannot create directory '{p}': No such file or directory") + except NotADirectoryError: + outputs.append(f"mkdir: cannot create directory '{p}': Not a directory") + return "\n".join(outputs) + + +# --------------------------------------------------------------------------- +# chmod +# --------------------------------------------------------------------------- + +class ChmodHandler(CommandHandler): + name = "chmod" + available_os = {OSType.LINUX, OSType.MAC} + + def execute(self, args: List[str], env: TerminalEnv) -> str: + if len(args) < 2: + return "chmod: missing operand" + recursive = False + mode_str = None + paths = [] + for a in args: + if a in ("-R", "--recursive"): + recursive = True + elif mode_str is None and len(a) == 3 and all(c in "01234567" for c in a): + mode_str = a + else: + paths.append(a) + if mode_str is None: + return f"chmod: invalid mode: '{args[0]}'" + if not paths: + return "chmod: missing operand" + + m = _machine(env) + outputs = [] + for p in paths: + mp = _resolve(env, p) + if not m.fs.exists(mp): + outputs.append(f"chmod: cannot access '{p}': No such file or directory") + continue + if recursive and m.fs.isdir(mp): + for dirpath, dirs, files in m.fs.walk(mp): + for name in dirs + files: + child = f"{dirpath}/{name}" if dirpath else name + meta = m.fs.get_meta(child) + meta.permissions = mode_str + else: + meta = m.fs.get_meta(mp) + meta.permissions = mode_str + return "\n".join(outputs) + + +# --------------------------------------------------------------------------- +# cp +# --------------------------------------------------------------------------- + +class CpHandler(CommandHandler): + name = "cp" + + def execute(self, args: List[str], env: TerminalEnv) -> str: + recursive = False + paths = [] + for a in args: + if a in ("-r", "-R", "--recursive"): + recursive = True + else: + paths.append(a) + if len(paths) < 2: + return "cp: missing destination operand" + + m = _machine(env) + src = _resolve(env, paths[0]) + dst = _resolve(env, paths[1]) + + if not m.fs.exists(src): + return f"cp: cannot stat '{paths[0]}': No such file or directory" + if m.fs.isdir(src) and not recursive: + return f"cp: -r not specified; omitting directory '{paths[0]}'" + + # If dst is an existing dir, copy into it + if m.fs.isdir(dst): + name = src.split("/")[-1] + dst = f"{dst}/{name}" + + try: + if m.fs.isdir(src): + m.fs.copy_tree(src, dst) + else: + # Ensure parent exists + parent = "/".join(dst.split("/")[:-1]) + if parent and not m.fs.exists(parent): + return f"cp: cannot create regular file '{paths[1]}': No such file or directory" + m.fs.copy_file(src, dst) + except Exception as e: + return f"cp: {e}" + return "" + + +# --------------------------------------------------------------------------- +# mv +# --------------------------------------------------------------------------- + +class MvHandler(CommandHandler): + name = "mv" + + def execute(self, args: List[str], env: TerminalEnv) -> str: + paths = [a for a in args if not a.startswith("-")] + if len(paths) < 2: + return "mv: missing destination operand" + + m = _machine(env) + src = _resolve(env, paths[0]) + dst = _resolve(env, paths[1]) + + if not m.fs.exists(src): + return f"mv: cannot stat '{paths[0]}': No such file or directory" + + # If dst is existing dir, move into it + if m.fs.isdir(dst): + name = src.split("/")[-1] + dst = f"{dst}/{name}" + + # Check parent of dst + dst_parent = "/".join(dst.split("/")[:-1]) + if dst_parent and not m.fs.isdir(dst_parent): + return f"mv: cannot move '{paths[0]}' to '{paths[1]}': No such file or directory" + + try: + if m.fs.isdir(src): + m.fs.copy_tree(src, dst) + else: + m.fs.copy_file(src, dst) + m.fs.remove_recursive(src) + except Exception as e: + return f"mv: {e}" + return "" + + +# --------------------------------------------------------------------------- +# rm +# --------------------------------------------------------------------------- + +class RmHandler(CommandHandler): + name = "rm" + + def execute(self, args: List[str], env: TerminalEnv) -> str: + recursive = False + force = False + paths = [] + for a in args: + if a.startswith("-"): + flags = a.lstrip("-") + if "r" in flags or "R" in flags: + recursive = True + if "f" in flags: + force = True + else: + paths.append(a) + if not paths: + return "rm: missing operand" + + m = _machine(env) + outputs = [] + for p in paths: + mp = _resolve(env, p) + if not m.fs.exists(mp): + if not force: + outputs.append(f"rm: cannot remove '{p}': No such file or directory") + continue + if m.fs.isdir(mp) and not recursive: + outputs.append(f"rm: cannot remove '{p}': Is a directory") + continue + try: + if recursive: + m.fs.remove_recursive(mp) + else: + m.fs.remove(mp) + except Exception as e: + outputs.append(f"rm: {e}") + return "\n".join(outputs) + + +# --------------------------------------------------------------------------- +# touch +# --------------------------------------------------------------------------- + +class TouchHandler(CommandHandler): + name = "touch" + available_os = {OSType.LINUX, OSType.MAC} + + def execute(self, args: List[str], env: TerminalEnv) -> str: + if not args: + return "touch: missing file operand" + m = _machine(env) + for f in args: + if f.startswith("-"): + continue + mp = _resolve(env, f) + if not m.fs.exists(mp): + parent = "/".join(mp.split("/")[:-1]) + if parent and not m.fs.isdir(parent): + return f"touch: cannot touch '{f}': No such file or directory" + m.fs.writetext(mp, "") + return "" + + +# --------------------------------------------------------------------------- +# echo +# --------------------------------------------------------------------------- + +class EchoHandler(CommandHandler): + name = "echo" + + def execute(self, args: List[str], env: TerminalEnv) -> str: + return " ".join(args) + + +# --------------------------------------------------------------------------- +# whoami +# --------------------------------------------------------------------------- + +class WhoamiHandler(CommandHandler): + name = "whoami" + + def execute(self, args: List[str], env: TerminalEnv) -> str: + return _machine(env).username + + +# --------------------------------------------------------------------------- +# ssh +# --------------------------------------------------------------------------- + +class SshHandler(CommandHandler): + name = "ssh" + + def execute(self, args: List[str], env: TerminalEnv) -> str: + if env.ssh_connected: + return "bash: already connected to remote host. Use 'exit' first." + + # Parse: ssh [-p port] [-i key] user@host + target = None + i = 0 + while i < len(args): + if args[i] in ("-p", "-i") and i + 1 < len(args): + i += 2 # skip flag and value + continue + if not args[i].startswith("-"): + target = args[i] + break + i += 1 + + if target is None: + return "usage: ssh [user@]hostname" + + if "@" not in target: + return f"ssh: Could not resolve hostname {target}: Name or service not known" + + user, host = target.split("@", 1) + # Check against known remote + if host != env.remote.hostname and host != env.remote.env_vars.get("IP", ""): + return f"ssh: connect to host {host} port 22: Connection refused" + if user != env.remote.username: + return f"Permission denied (publickey,password)." + + env.ssh_connected = True + env.current_machine = env.remote + return f"Welcome to Ubuntu 22.04 LTS ({env.remote.hostname})\nLast login: Mon Jan 15 10:30:00 2024" + + +# --------------------------------------------------------------------------- +# exit +# --------------------------------------------------------------------------- + +class ExitHandler(CommandHandler): + name = "exit" + + def execute(self, args: List[str], env: TerminalEnv) -> str: + if env.ssh_connected: + host = env.remote.hostname + env.ssh_connected = False + env.current_machine = env.local + return f"Connection to {host} closed." + return "" + + +# --------------------------------------------------------------------------- +# scp +# --------------------------------------------------------------------------- + +class ScpHandler(CommandHandler): + name = "scp" + + def execute(self, args: List[str], env: TerminalEnv) -> str: + if env.ssh_connected: + return "bash: scp must be run from local machine. Use 'exit' to disconnect first." + + recursive = False + positional = [] + i = 0 + while i < len(args): + a = args[i] + if a in ("-r", "-R"): + recursive = True + elif a in ("-P", "-i") and i + 1 < len(args): + i += 1 # skip value + elif not a.startswith("-"): + positional.append(a) + i += 1 + + if len(positional) < 2: + return "usage: scp [-r] source ... target" + + src_arg = positional[0] + dst_arg = positional[1] + + src_user, src_host, src_path = _parse_scp_target(src_arg) + dst_user, dst_host, dst_path = _parse_scp_target(dst_arg) + + # Determine source/dest machines + if src_host and dst_host: + return "scp: copying between two remote hosts is not supported" + elif src_host: + # Download: remote -> local + src_machine = env.remote + dst_machine = env.local + if src_host != env.remote.hostname and src_host != env.remote.env_vars.get("IP", ""): + return f"ssh: connect to host {src_host} port 22: Connection refused" + elif dst_host: + # Upload: local -> remote + src_machine = env.local + dst_machine = env.remote + if dst_host != env.remote.hostname and dst_host != env.remote.env_vars.get("IP", ""): + return f"ssh: connect to host {dst_host} port 22: Connection refused" + else: + return "scp: use cp for local-to-local copy" + + src_mp = src_machine.resolve_path(src_path) + dst_mp = dst_machine.resolve_path(dst_path) + + if not src_machine.fs.exists(src_mp): + return f"scp: {src_path}: No such file or directory" + + if src_machine.fs.isdir(src_mp) and not recursive: + return f"scp: {src_path}: not a regular file" + + # If dst is existing dir, copy into it + if dst_machine.fs.isdir(dst_mp): + name = src_mp.split("/")[-1] + dst_mp = f"{dst_mp}/{name}" + + # Check dst parent exists + dst_parent = "/".join(dst_mp.split("/")[:-1]) + if dst_parent and not dst_machine.fs.isdir(dst_parent): + return f"scp: {dst_path}: No such file or directory" + + try: + self._cross_copy(src_machine, src_mp, dst_machine, dst_mp, recursive) + except Exception as e: + return f"scp: {e}" + + filename = src_mp.split("/")[-1] + size = src_machine.fs.get_meta(src_mp).size + return f"{filename} 100% {_format_size(size)} transferred" + + def _cross_copy(self, src_m, src_mp, dst_m, dst_mp, recursive): + """Copy across two different MachineStates.""" + src_fs = src_m.fs + dst_fs = dst_m.fs + if src_fs.isfile(src_mp): + content = src_fs.readtext(src_mp) + dst_fs.writetext(dst_mp, content, create_parents=False) + src_meta = src_fs.get_meta(src_mp) + dst_meta = dst_fs.get_meta(dst_mp) + dst_meta.permissions = src_meta.permissions + dst_meta.owner = dst_m.username + dst_meta.group = dst_m.username + dst_meta.archive_type = src_meta.archive_type + dst_meta.archive_entries = ( + dict(src_meta.archive_entries) if src_meta.archive_entries else None + ) + elif recursive: + dst_fs.makedirs(dst_mp) + for name in src_fs.listdir(src_mp): + s = f"{src_mp}/{name}" + d = f"{dst_mp}/{name}" + self._cross_copy(src_m, s, dst_m, d, True) + + +# --------------------------------------------------------------------------- +# rsync +# --------------------------------------------------------------------------- + +class RsyncHandler(CommandHandler): + name = "rsync" + available_os = {OSType.LINUX, OSType.MAC} + + def execute(self, args: List[str], env: TerminalEnv) -> str: + if env.ssh_connected: + return "bash: rsync must be run from local machine. Use 'exit' to disconnect first." + + # Strip flags, find positional args + positional = [] + i = 0 + while i < len(args): + a = args[i] + if a in ("-e",) and i + 1 < len(args): + i += 1 # skip value + elif not a.startswith("-"): + positional.append(a) + i += 1 + + if len(positional) < 2: + return "usage: rsync [options] source destination" + + src_arg = positional[0] + dst_arg = positional[1] + + # Reuse SCP logic for cross-machine copy + scp = ScpHandler() + # Build equivalent scp args + scp_args = ["-r", src_arg, dst_arg] + result = scp.execute(scp_args, env) + + if "transferred" in result: + filename = src_arg.rstrip("/").split("/")[-1] + return (f"sending incremental file list\n" + f"{filename}\n" + f"\nsent 1024 bytes received 42 bytes 2132.00 bytes/sec\n" + f"total size is 1024 speedup is 0.96") + return result + + +# --------------------------------------------------------------------------- +# tar +# --------------------------------------------------------------------------- + +class TarHandler(CommandHandler): + name = "tar" + available_os = {OSType.LINUX, OSType.MAC} + + def execute(self, args: List[str], env: TerminalEnv) -> str: + if not args: + return "tar: You must specify one of the '-Acdtrux' options" + + m = _machine(env) + fs = m.fs + + # Collect all flags and positional args. + # tar accepts flags in many forms: -czf, czf, -c -z -f, --create, etc. + mode = None + gzip = False + archive_path = None + file_args = [] + expect_archive = False # next positional arg is the archive path + + for a in args: + if expect_archive: + archive_path = a + expect_archive = False + continue + if a.startswith("--"): + if a == "--create": + mode = "create" + elif a in ("--extract", "--get"): + mode = "extract" + elif a == "--list": + mode = "list" + elif a in ("--gzip", "--gunzip", "--ungzip"): + gzip = True + elif a.startswith("--file="): + archive_path = a.split("=", 1)[1] + elif a == "--file": + expect_archive = True + # --verbose and other long flags are silently ignored + continue + if a.startswith("-") or (a == args[0] and len(a) <= 6 and a[0] in "cxtzvf"): + # Short flags: -czf, czf, -c, -z, -f + chars = a.lstrip("-") + for ch in chars: + if ch == "c": + mode = "create" + elif ch == "x": + mode = "extract" + elif ch == "t": + mode = "list" + elif ch == "z": + gzip = True + elif ch == "f": + expect_archive = True + # v and other single-char flags silently ignored + continue + # Positional arg + file_args.append(a) + + if archive_path is None: + return "tar: Refusing to read archive contents from terminal" + + archive_mp = _resolve(env, archive_path) + archive_type = "tar.gz" if gzip else "tar" + + if mode == "create": + if not file_args: + return "tar: Cowardly refusing to create an empty archive" + entries = {} + for f in file_args: + fmp = _resolve(env, f) + if not fs.exists(fmp): + return f"tar: {f}: Cannot open: No such file or directory" + if fs.isfile(fmp): + basename = fmp.split("/")[-1] + entries[basename] = fs.readtext(fmp) + elif fs.isdir(fmp): + base = fmp.rstrip("/").split("/")[-1] + for dp, dirs, files in fs.walk(fmp): + for fname in files: + child = f"{dp}/{fname}" if dp else fname + rel = child[len(fmp):].lstrip("/") if child.startswith(fmp) else child + full_rel = f"{base}/{rel}" if rel else base + entries[full_rel] = fs.readtext(child) + + content = json.dumps({"type": archive_type, "entries": entries}) + parent = "/".join(archive_mp.split("/")[:-1]) + if parent and not fs.isdir(parent): + return f"tar: {archive_path}: Cannot open: No such file or directory" + fs.writetext(archive_mp, content) + meta = fs.get_meta(archive_mp) + meta.archive_type = archive_type + meta.archive_entries = entries + meta.size = sum(len(v) for v in entries.values()) + return "" + + elif mode == "extract": + if not fs.exists(archive_mp): + return f"tar: {archive_path}: Cannot open: No such file or directory" + meta = fs.get_meta(archive_mp) + if not meta.archive_type or "tar" not in meta.archive_type: + return f"tar: {archive_path}: This does not look like a tar archive" + entries = meta.archive_entries or {} + for rel_path, content in entries.items(): + out_path = _resolve(env, rel_path) + parent = "/".join(out_path.split("/")[:-1]) + if parent: + fs.makedirs(parent) + fs.writetext(out_path, content) + out_meta = fs.get_meta(out_path) + out_meta.owner = m.username + out_meta.group = m.username + return "" + + elif mode == "list": + if not fs.exists(archive_mp): + return f"tar: {archive_path}: Cannot open: No such file or directory" + meta = fs.get_meta(archive_mp) + if not meta.archive_type or "tar" not in meta.archive_type: + return f"tar: {archive_path}: This does not look like a tar archive" + entries = meta.archive_entries or {} + return "\n".join(sorted(entries.keys())) + + return "tar: You must specify one of the '-Acdtrux' options" + + +# --------------------------------------------------------------------------- +# zip +# --------------------------------------------------------------------------- + +class ZipHandler(CommandHandler): + name = "zip" + + def execute(self, args: List[str], env: TerminalEnv) -> str: + if len(args) < 2: + return "zip: missing archive name or files" + + m = _machine(env) + fs = m.fs + recursive = False + positional = [] + for a in args: + if a in ("-r", "-R"): + recursive = True + elif not a.startswith("-"): + positional.append(a) + + if len(positional) < 2: + return "zip: missing archive name or files" + + archive_name = positional[0] + file_args = positional[1:] + archive_mp = _resolve(env, archive_name) + + entries = {} + output_lines = [] + for f in file_args: + fmp = _resolve(env, f) + if not fs.exists(fmp): + return f"zip error: Nothing to do! ({f} not found)" + if fs.isfile(fmp): + basename = fmp.split("/")[-1] + entries[basename] = fs.readtext(fmp) + output_lines.append(f" adding: {basename} (stored 0%)") + elif fs.isdir(fmp): + if not recursive: + output_lines.append(f" adding: {f}/ (stored 0%)") + continue + base = fmp.rstrip("/").split("/")[-1] + for dp, dirs, files in fs.walk(fmp): + for fname in files: + child = f"{dp}/{fname}" if dp else fname + rel = child[len(fmp):].lstrip("/") if child.startswith(fmp) else child + full_rel = f"{base}/{rel}" if rel else base + entries[full_rel] = fs.readtext(child) + output_lines.append(f" adding: {full_rel} (stored 0%)") + + content = json.dumps({"type": "zip", "entries": entries}) + parent = "/".join(archive_mp.split("/")[:-1]) + if parent and not fs.isdir(parent): + fs.makedirs(parent) + fs.writetext(archive_mp, content) + meta = fs.get_meta(archive_mp) + meta.archive_type = "zip" + meta.archive_entries = entries + meta.size = sum(len(v) for v in entries.values()) + + return "\n".join(output_lines) + + +# --------------------------------------------------------------------------- +# unzip +# --------------------------------------------------------------------------- + +class UnzipHandler(CommandHandler): + name = "unzip" + + def execute(self, args: List[str], env: TerminalEnv) -> str: + if not args: + return "unzip: missing archive name" + + m = _machine(env) + fs = m.fs + archive_name = None + dest_dir = None + i = 0 + while i < len(args): + if args[i] == "-d" and i + 1 < len(args): + dest_dir = args[i + 1] + i += 2 + continue + if not args[i].startswith("-"): + archive_name = args[i] + i += 1 + + if archive_name is None: + return "unzip: missing archive name" + + archive_mp = _resolve(env, archive_name) + if not fs.exists(archive_mp): + return f"unzip: cannot find or open {archive_name}" + + meta = fs.get_meta(archive_mp) + if meta.archive_type != "zip": + return f"unzip: {archive_name} is not a zip archive" + + entries = meta.archive_entries or {} + output_lines = [f"Archive: {archive_name}"] + + for rel_path, content in entries.items(): + if dest_dir: + out_path = _resolve(env, f"{dest_dir}/{rel_path}") + else: + out_path = _resolve(env, rel_path) + parent = "/".join(out_path.split("/")[:-1]) + if parent: + fs.makedirs(parent) + fs.writetext(out_path, content) + out_meta = fs.get_meta(out_path) + out_meta.owner = m.username + out_meta.group = m.username + output_lines.append(f" extracting: {rel_path}") + + return "\n".join(output_lines) + + +# --------------------------------------------------------------------------- +# gzip / gunzip +# --------------------------------------------------------------------------- + +class GzipHandler(CommandHandler): + name = "gzip" + available_os = {OSType.LINUX, OSType.MAC} + + def execute(self, args: List[str], env: TerminalEnv) -> str: + paths = [a for a in args if not a.startswith("-")] + if not paths: + return "gzip: missing file operand" + + m = _machine(env) + fs = m.fs + for f in paths: + mp = _resolve(env, f) + if not fs.exists(mp): + return f"gzip: {f}: No such file or directory" + if fs.isdir(mp): + return f"gzip: {f}: Is a directory" + content = fs.readtext(mp) + gz_path = mp + ".gz" + gz_content = json.dumps({"type": "gzip", "entries": {f: content}}) + fs.writetext(gz_path, gz_content) + gz_meta = fs.get_meta(gz_path) + gz_meta.archive_type = "gzip" + gz_meta.archive_entries = {f: content} + gz_meta.size = len(content) + fs.remove(mp) + return "" + + +class GunzipHandler(CommandHandler): + name = "gunzip" + available_os = {OSType.LINUX, OSType.MAC} + + def execute(self, args: List[str], env: TerminalEnv) -> str: + paths = [a for a in args if not a.startswith("-")] + if not paths: + return "gunzip: missing file operand" + + m = _machine(env) + fs = m.fs + for f in paths: + mp = _resolve(env, f) + if not fs.exists(mp): + return f"gunzip: {f}: No such file or directory" + meta = fs.get_meta(mp) + if meta.archive_type != "gzip": + return f"gunzip: {f}: not in gzip format" + entries = meta.archive_entries or {} + for orig_name, content in entries.items(): + out_path = _resolve(env, orig_name) + fs.writetext(out_path, content) + fs.remove(mp) + return "" + + +# --------------------------------------------------------------------------- +# Registry +# --------------------------------------------------------------------------- + +def build_command_registry() -> Dict[str, CommandHandler]: + """Build the command name -> handler mapping.""" + handlers = [ + LsHandler(), + CdHandler(), + PwdHandler(), + CatHandler(), + MkdirHandler(), + ChmodHandler(), + CpHandler(), + MvHandler(), + RmHandler(), + TouchHandler(), + EchoHandler(), + WhoamiHandler(), + SshHandler(), + ExitHandler(), + ScpHandler(), + RsyncHandler(), + TarHandler(), + ZipHandler(), + UnzipHandler(), + GzipHandler(), + GunzipHandler(), + ] + registry = {} + for h in handlers: + registry[h.name] = h + return registry diff --git a/trinity/common/workflows/connect_the_dots/terminal/env.py b/trinity/common/workflows/connect_the_dots/terminal/env.py new file mode 100644 index 00000000000..6ffbfc2ca85 --- /dev/null +++ b/trinity/common/workflows/connect_the_dots/terminal/env.py @@ -0,0 +1,636 @@ +# -*- coding: utf-8 -*- +""" +Virtual terminal environment with in-memory filesystem. + +No real commands are executed. All state is held in Python dicts. +""" + +import re +import shlex +from dataclasses import dataclass, field +from enum import Enum +from typing import Any, Dict, List, Optional, Tuple + + +# --------------------------------------------------------------------------- +# Enums +# --------------------------------------------------------------------------- + +class OSType(Enum): + WINDOWS = "windows" + MAC = "mac" + LINUX = "linux" + + +# --------------------------------------------------------------------------- +# Virtual Filesystem +# --------------------------------------------------------------------------- + + + +@dataclass +class FileMeta: + """Metadata attached to every file/directory.""" + permissions: str = "644" + owner: str = "user" + group: str = "user" + size: int = 0 + # For simulated archives + archive_type: Optional[str] = None # "tar", "zip", "tar.gz", "gzip" + archive_entries: Optional[Dict[str, str]] = None + + +class VirtualFS: + """Minimal in-memory filesystem backed by nested dicts. + + Internal layout:: + + _tree = { + "home": { + "user": { + "file.txt": "content string", + "subdir": { ... }, + } + } + } + + - A *str* value => regular file (the string is its content). + - A *dict* value => directory. + + All paths inside VirtualFS use ``/`` separators with no leading slash. + The caller (MachineState) is responsible for converting OS-specific paths + to this canonical form *before* calling VirtualFS methods. + """ + + def __init__(self): + self._tree: Dict[str, Any] = {} + self._meta: Dict[str, FileMeta] = {} # canonical_path -> metadata + + # -- helpers -------------------------------------------------------- + + @staticmethod + def _split(path: str) -> List[str]: + """Split a canonical path into parts, filtering blanks.""" + return [p for p in path.replace("\\", "/").split("/") if p] + + def _navigate(self, parts: List[str], create_parents: bool = False): + """Walk *parts* from root. Returns (parent_dict, last_key). + + If *create_parents* is True, intermediate dirs are created. + Raises FileNotFoundError if a segment doesn't exist (and create is off). + Raises NotADirectoryError if a segment is a file, not a dir. + """ + node = self._tree + for i, part in enumerate(parts[:-1]): + child = node.get(part) + if child is None: + if create_parents: + node[part] = {} + dir_path = "/".join(parts[: i + 1]) + self._meta[dir_path] = FileMeta(permissions="755", size=4096) + node = node[part] + else: + raise FileNotFoundError("/".join(parts[: i + 1])) + elif isinstance(child, dict): + node = child + else: + raise NotADirectoryError("/".join(parts[: i + 1])) + return (node, parts[-1]) if parts else (self._tree, "") + + # -- public API ----------------------------------------------------- + + def exists(self, path: str) -> bool: + parts = self._split(path) + if not parts: + return True # root always exists + try: + parent, key = self._navigate(parts) + return key in parent + except (FileNotFoundError, NotADirectoryError): + return False + + def isdir(self, path: str) -> bool: + parts = self._split(path) + if not parts: + return True + try: + parent, key = self._navigate(parts) + return isinstance(parent.get(key), dict) + except (FileNotFoundError, NotADirectoryError): + return False + + def isfile(self, path: str) -> bool: + parts = self._split(path) + if not parts: + return False + try: + parent, key = self._navigate(parts) + val = parent.get(key) + return val is not None and not isinstance(val, dict) + except (FileNotFoundError, NotADirectoryError): + return False + + def listdir(self, path: str) -> List[str]: + parts = self._split(path) + node = self._tree + for part in parts: + child = node.get(part) + if child is None: + raise FileNotFoundError(path) + if not isinstance(child, dict): + raise NotADirectoryError(path) + node = child + return sorted(node.keys()) + + def readtext(self, path: str) -> str: + parts = self._split(path) + if not parts: + raise IsADirectoryError(path) + parent, key = self._navigate(parts) + val = parent.get(key) + if val is None: + raise FileNotFoundError(path) + if isinstance(val, dict): + raise IsADirectoryError(path) + return val + + def writetext(self, path: str, content: str, create_parents: bool = False) -> None: + parts = self._split(path) + if not parts: + raise IsADirectoryError(path) + parent, key = self._navigate(parts, create_parents=create_parents) + parent[key] = content + meta = self._meta.get(path) + if meta is None: + meta = FileMeta() + self._meta[path] = meta + meta.size = len(content) + + def makedirs(self, path: str) -> None: + parts = self._split(path) + if not parts: + return + node = self._tree + for i, part in enumerate(parts): + child = node.get(part) + if child is None: + node[part] = {} + dir_path = "/".join(parts[: i + 1]) + self._meta[dir_path] = FileMeta(permissions="755", size=4096) + node = node[part] + elif isinstance(child, dict): + node = child + else: + raise NotADirectoryError("/".join(parts[: i + 1])) + + def mkdir(self, path: str) -> None: + """Create a single directory (parent must exist).""" + parts = self._split(path) + if not parts: + return + parent, key = self._navigate(parts) + if key in parent: + raise FileExistsError(path) + parent[key] = {} + self._meta[path] = FileMeta(permissions="755", size=4096) + + def remove(self, path: str) -> None: + parts = self._split(path) + if not parts: + raise PermissionError("cannot remove root") + parent, key = self._navigate(parts) + if key not in parent: + raise FileNotFoundError(path) + val = parent[key] + if isinstance(val, dict) and val: + raise OSError(f"directory not empty: {path}") + del parent[key] + self._meta.pop(path, None) + + def remove_recursive(self, path: str) -> None: + parts = self._split(path) + if not parts: + self._tree.clear() + self._meta.clear() + return + parent, key = self._navigate(parts) + if key not in parent: + raise FileNotFoundError(path) + # Remove all metadata under this path + prefix = path.rstrip("/") + "/" + to_del = [k for k in self._meta if k == path or k.startswith(prefix)] + for k in to_del: + del self._meta[k] + del parent[key] + + def get_meta(self, path: str) -> FileMeta: + parts = self._split(path) + canon = "/".join(parts) if parts else "" + meta = self._meta.get(canon) + if meta is None: + # Auto-create metadata + meta = FileMeta() + if self.isdir(path): + meta.permissions = "755" + meta.size = 4096 + else: + try: + content = self.readtext(path) + meta.size = len(content) + except Exception: + pass + self._meta[canon] = meta + return meta + + + def walk(self, path: str = "") -> List[Tuple[str, List[str], List[str]]]: + """Walk filesystem tree. Yields (dir_path, [subdirs], [files]).""" + parts = self._split(path) + node = self._tree + for part in parts: + child = node.get(part) + if child is None or not isinstance(child, dict): + return [] + node = child + + result = [] + self._walk_recursive(node, path, result) + return result + + def _walk_recursive(self, node: dict, prefix: str, result: list): + dirs = [] + files = [] + for name, val in sorted(node.items()): + if isinstance(val, dict): + dirs.append(name) + else: + files.append(name) + result.append((prefix, dirs, files)) + for d in dirs: + child_prefix = f"{prefix}/{d}" if prefix else d + self._walk_recursive(node[d], child_prefix, result) + + def copy_file(self, src: str, dst: str) -> None: + """Copy a single file.""" + content = self.readtext(src) + src_meta = self.get_meta(src) + self.writetext(dst, content, create_parents=False) + dst_meta = self.get_meta(dst) + dst_meta.permissions = src_meta.permissions + dst_meta.owner = src_meta.owner + dst_meta.group = src_meta.group + dst_meta.archive_type = src_meta.archive_type + dst_meta.archive_entries = ( + dict(src_meta.archive_entries) if src_meta.archive_entries else None + ) + + def copy_tree(self, src: str, dst: str) -> None: + """Recursively copy directory.""" + if self.isfile(src): + self.copy_file(src, dst) + return + self.makedirs(dst) + for name in self.listdir(src): + s = f"{src}/{name}" if src else name + d = f"{dst}/{name}" if dst else name + if self.isdir(s): + self.copy_tree(s, d) + else: + self.copy_file(s, d) + + +# --------------------------------------------------------------------------- +# Machine State +# --------------------------------------------------------------------------- + +@dataclass +class MachineState: + """State of a single virtual machine.""" + os_type: OSType + hostname: str + username: str + home_dir: str # canonical memfs path, e.g. "home/user" or "C/Users/user" + cwd: str # canonical memfs path + fs: VirtualFS = field(default_factory=VirtualFS) + env_vars: Dict[str, str] = field(default_factory=dict) + + # -- path conversion ------------------------------------------------ + + def to_memfs_path(self, user_path: str) -> str: + """Convert a user-visible path to canonical MemFS path (no leading /).""" + if self.os_type == OSType.WINDOWS: + return self._win_to_memfs(user_path) + return self._unix_to_memfs(user_path) + + def to_display_path(self, memfs_path: str) -> str: + """Convert canonical MemFS path to user-visible path.""" + if self.os_type == OSType.WINDOWS: + return self._memfs_to_win(memfs_path) + return "/" + memfs_path if memfs_path else "/" + + def _win_to_memfs(self, p: str) -> str: + p = p.replace("/", "\\") + # Absolute: C:\... -> C/... + if len(p) >= 2 and p[1] == ":": + drive = p[0].upper() + rest = p[2:].lstrip("\\") + parts = [drive] + [x for x in rest.split("\\") if x] + return "/".join(parts) + # Relative + parts = [x for x in p.split("\\") if x] + if not parts: + return self.cwd + return self.cwd + "/" + "/".join(parts) + + def _memfs_to_win(self, p: str) -> str: + parts = p.split("/") + if parts and len(parts[0]) == 1 and parts[0].isalpha(): + drive = parts[0].upper() + rest = "\\".join(parts[1:]) + return f"{drive}:\\{rest}" if rest else f"{drive}:\\" + return "\\".join(parts) + + def _unix_to_memfs(self, p: str) -> str: + if p.startswith("/"): + return p.lstrip("/") + # Relative + if not p or p == ".": + return self.cwd + if p == "..": + parts = self.cwd.split("/") + return "/".join(parts[:-1]) if len(parts) > 1 else "" + return f"{self.cwd}/{p}" if self.cwd else p + + def resolve_path(self, user_path: str) -> str: + """Resolve a user-visible path to canonical MemFS path, handling . and ..""" + raw = self.to_memfs_path(user_path) + parts = raw.split("/") + resolved = [] + for p in parts: + if p == "" or p == ".": + continue + elif p == "..": + if resolved: + resolved.pop() + else: + resolved.append(p) + return "/".join(resolved) + + def get_display_cwd(self) -> str: + """CWD for display in terminal prompt, with ~ substitution (Unix only).""" + display = self.to_display_path(self.cwd) + if self.os_type == OSType.WINDOWS: + return display + home_display = self.to_display_path(self.home_dir) + if display == home_display: + return "~" + if display.startswith(home_display + "/"): + return "~" + display[len(home_display):] + return display + + def get_prompt_string(self) -> str: + """Render the terminal prompt.""" + if self.os_type == OSType.WINDOWS: + return f"{self.to_display_path(self.cwd)}>" + elif self.os_type == OSType.MAC: + return f"{self.username}@MacBook {self.get_display_cwd()} % " + else: + return f"{self.username}@{self.hostname}:{self.get_display_cwd()}$ " + + +# --------------------------------------------------------------------------- +# Terminal Environment +# --------------------------------------------------------------------------- + +class TerminalEnv: + """Complete terminal environment with local and remote machines. + + The *command_handlers* dict is injected by the caller (see commands.py). + """ + + def __init__( + self, + local: MachineState, + remote: MachineState, + command_handlers: Dict[str, Any], + max_steps: int = 15, + ): + self.local = local + self.remote = remote + self.command_handlers = command_handlers + self.max_steps = max_steps + self.current_machine: MachineState = self.local + self.ssh_connected: bool = False + self.step_count: int = 0 + self.command_history: List[str] = [] + + # -- prompt ---------------------------------------------------------- + + def get_prompt_string(self) -> str: + return self.current_machine.get_prompt_string() + + # -- execution ------------------------------------------------------- + + def step(self, raw_command: str, goal_check_fn=None): + """Gym-style step. Returns (observation, reward, done, info). + + Args: + raw_command: The command string to execute. + goal_check_fn: Optional callable ``fn(env) -> float`` that returns + 1.0 if the task is complete, 0.0 otherwise. When it returns + 1.0, the episode terminates early. + """ + observation = self.execute(raw_command) + reward = 0.0 + done = False + info = {} + + if goal_check_fn is not None: + reward = goal_check_fn(self) + if reward > 0.5: + done = True + info["early_completion"] = True + + if self.step_count >= self.max_steps: + done = True + + return observation, reward, done, info + + def execute(self, raw_command: str) -> str: + """Execute a raw command string. Returns terminal output.""" + raw_command = raw_command.strip() + if not raw_command: + return "" + + self.command_history.append(raw_command) + self.step_count += 1 + + # Expand environment variables + expanded = self._expand_env_vars(raw_command) + + # Check for unsupported operators + if _contains_unquoted(expanded, "|"): + return "Error: pipes (|) are not supported in this environment." + if _contains_unquoted(expanded, ">") or _contains_unquoted(expanded, "<"): + return "Error: redirects (>, >>, <) are not supported in this environment." + + # Split on && and ; + segments = _split_command_chain(expanded) + + output_parts = [] + last_was_error = False + prev_operator = None + for cmd_str, operator in segments: + cmd_str = cmd_str.strip() + if not cmd_str: + prev_operator = operator + continue + # && : skip if previous command failed + if prev_operator == "&&" and last_was_error: + break + result, last_was_error = self._execute_single(cmd_str) + if result: + output_parts.append(result) + prev_operator = operator + + return "\n".join(output_parts) + + def _execute_single(self, cmd_str: str) -> Tuple[str, bool]: + """Execute one command. Returns (output, is_error).""" + try: + posix = self.current_machine.os_type != OSType.WINDOWS or self.ssh_connected + tokens = shlex.split(cmd_str, posix=posix) + except ValueError: + return "bash: syntax error: unexpected end of file", True + + if not tokens: + return "", False + + cmd_name = tokens[0] + args = tokens[1:] + + # Resolve aliases + cmd_key = self._resolve_alias(cmd_name.lower()) + + handler = self.command_handlers.get(cmd_key) + if handler is None: + return self._unknown_cmd_error(cmd_name), True + + # Check OS availability + if not handler.is_available(self.current_machine.os_type, self.ssh_connected): + return self._unknown_cmd_error(cmd_name), True + + try: + result = handler.execute(args, self) + # Detect error by common patterns: "cmd: error...", "Error:", "bash:", + # Windows errors, or empty success + is_err = bool(result) and ( + result.startswith(("Error:", "bash:", "'", "The system", "Could Not")) + or (": " in result.split("\n")[0] + and result.split(":")[0].split()[-1].lower() in ( + "ls", "cd", "cat", "mkdir", "chmod", "cp", "mv", "rm", + "scp", "rsync", "tar", "zip", "unzip", "gzip", "gunzip", + "ssh", "touch", "del", + )) + ) + return result, is_err + except Exception as e: + return f"Error: {e}", True + + def _expand_env_vars(self, cmd: str) -> str: + machine = self.current_machine + if machine.os_type == OSType.WINDOWS and not self.ssh_connected: + for var, val in machine.env_vars.items(): + cmd = cmd.replace(f"%{var}%", val) + else: + # ~ expansion: only standalone ~ or ~/... at token boundaries + # Do NOT expand ~ after : (e.g., scp file user@host:~/dir) + home = machine.env_vars.get("HOME", machine.to_display_path(machine.home_dir)) + cmd = re.sub(r"(? str: + aliases = { + "dir": "ls", + "type": "cat", + "copy": "cp", + "move": "mv", + "del": "rm", + "ren": "mv", + "rename": "mv", + } + return aliases.get(name, name) + + def _unknown_cmd_error(self, cmd_name: str) -> str: + if self.current_machine.os_type == OSType.WINDOWS and not self.ssh_connected: + return ( + f"'{cmd_name}' is not recognized as an internal or external command,\n" + f"operable program or batch file." + ) + return f"bash: {cmd_name}: command not found" + + +# --------------------------------------------------------------------------- +# Helpers for command chain parsing +# --------------------------------------------------------------------------- + +def _contains_unquoted(text: str, char: str) -> bool: + """Check if *char* appears outside of quotes.""" + in_single = False + in_double = False + i = 0 + while i < len(text): + c = text[i] + if c == "'" and not in_double: + in_single = not in_single + elif c == '"' and not in_single: + in_double = not in_double + elif c == char and not in_single and not in_double: + # For &&, check double char + if char == "&" and i + 1 < len(text) and text[i + 1] == "&": + return False # && is handled separately + if char in (">", "<", "|"): + return True + i += 1 + return False + + +def _split_command_chain(text: str) -> List[Tuple[str, Optional[str]]]: + """Split on ``&&`` and ``;``. Returns [(cmd, operator), ...].""" + result = [] + current = [] + in_single = False + in_double = False + i = 0 + while i < len(text): + c = text[i] + if c == "'" and not in_double: + in_single = not in_single + current.append(c) + elif c == '"' and not in_single: + in_double = not in_double + current.append(c) + elif not in_single and not in_double: + if c == ";" : + result.append(("".join(current), ";")) + current = [] + elif c == "&" and i + 1 < len(text) and text[i + 1] == "&": + result.append(("".join(current), "&&")) + current = [] + i += 1 # skip second & + else: + current.append(c) + else: + current.append(c) + i += 1 + + remaining = "".join(current).strip() + if remaining: + result.append((remaining, None)) + return result diff --git a/trinity/common/workflows/connect_the_dots/terminal/goal_check.py b/trinity/common/workflows/connect_the_dots/terminal/goal_check.py new file mode 100644 index 00000000000..b6e781b6e17 --- /dev/null +++ b/trinity/common/workflows/connect_the_dots/terminal/goal_check.py @@ -0,0 +1,77 @@ +# -*- coding: utf-8 -*- +""" +Goal checking for terminal tasks. + +Uses a unified check-list approach: every task (single or composite) defines +a list of ``goal_checks``, each a dict like:: + + {"type": "file_exists", "machine": "remote", "path": "...", "content": "..."} + +All checks must pass for reward = 1.0; any failure gives 0.0. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, Dict, List + +if TYPE_CHECKING: + from trinity.common.workflows.connect_the_dots.terminal.env import TerminalEnv + from trinity.common.workflows.connect_the_dots.terminal.task_gen import TerminalTask + + +def check_goal(task: TerminalTask, env: TerminalEnv) -> float: + """Return 1.0 if all goal checks pass, 0.0 otherwise.""" + for check in task.goal_checks: + if not _evaluate_check(check, env): + return 0.0 + return 1.0 + + +def _evaluate_check(check: Dict[str, Any], env: TerminalEnv) -> bool: + """Evaluate a single goal check.""" + machine = env.remote if check["machine"] == "remote" else env.local + fs = machine.fs + path = check["path"] + check_type = check["type"] + + if check_type == "file_exists": + if not fs.isfile(path): + return False + if "content" in check: + try: + actual = fs.readtext(path) + if actual != check["content"]: + return False + except Exception: + return False + return True + + elif check_type == "file_not_exists": + return not fs.exists(path) + + elif check_type == "dir_exists": + return fs.isdir(path) + + elif check_type == "permission_equals": + if not fs.exists(path): + return False + meta = fs.get_meta(path) + return meta.permissions == check["value"] + + elif check_type == "archive_contains": + if not fs.isfile(path): + return False + meta = fs.get_meta(path) + if meta.archive_type is None: + return False + expected_entries = check.get("entries", {}) + actual_entries = meta.archive_entries or {} + for rel_path, expected_content in expected_entries.items(): + if rel_path not in actual_entries: + return False + if actual_entries[rel_path] != expected_content: + return False + return True + + else: + raise ValueError(f"Unknown check type: {check_type}") diff --git a/trinity/common/workflows/connect_the_dots/terminal/prompts/__init__.py b/trinity/common/workflows/connect_the_dots/terminal/prompts/__init__.py new file mode 100644 index 00000000000..a2f5f124acc --- /dev/null +++ b/trinity/common/workflows/connect_the_dots/terminal/prompts/__init__.py @@ -0,0 +1,41 @@ +# -*- coding: utf-8 -*- +"""Prompt management for Terminal workflow using Jinja2 templates.""" + +from pathlib import Path +from typing import Optional + +from jinja2 import Environment, FileSystemLoader + +PROMPTS_DIR = Path(__file__).parent + + +def get_jinja_env() -> Environment: + return Environment( + loader=FileSystemLoader(PROMPTS_DIR), + trim_blocks=True, + lstrip_blocks=True, + ) + + +def load_system_prompt(**kwargs) -> str: + env = get_jinja_env() + template = env.get_template("system.jinja2") + return template.render(**kwargs) + + +def load_user_prompt( + current_step: int, + max_steps: int, + command_output: Optional[str] = None, + terminal_prompt: Optional[str] = None, + **kwargs, +) -> str: + env = get_jinja_env() + template = env.get_template("user.jinja2") + return template.render( + current_step=current_step, + max_steps=max_steps, + command_output=command_output, + terminal_prompt=terminal_prompt, + **kwargs, + ) diff --git a/trinity/common/workflows/connect_the_dots/terminal/prompts/system.jinja2 b/trinity/common/workflows/connect_the_dots/terminal/prompts/system.jinja2 new file mode 100644 index 00000000000..1e4ec8e29e4 --- /dev/null +++ b/trinity/common/workflows/connect_the_dots/terminal/prompts/system.jinja2 @@ -0,0 +1,29 @@ +You are interacting with a simulated terminal environment. Your goal is to complete file operations by typing commands. + +## Environment +- Local machine: {{ os_name }} (type: {{ os_type }}, hostname: localhost) +- Remote server: {{ remote_host }} (Linux, SSH pre-configured) +- Remote username: {{ remote_user }} + +## Available Commands +ls, dir, cd, pwd, cat, type, mkdir, chmod, cp, copy, mv, move, rm, del, ren, touch, scp, rsync, ssh, exit, whoami, echo, tar, zip, unzip, gzip, gunzip + +Note: Not all commands are available on all systems. + +## Task +{{ task_description }} + +## Response Format +First think step by step about which command to run, then provide your command using ... tags at the end. Format your complete response as follows: +``` +[THINKING] +COMMAND +``` +where [THINKING] should be replaced with your thinking process, while COMMAND is a single shell command line to execute this turn. + +Other requirements: +- Be concise and avoid overthinking in your thinking process. +- Issue exactly ONE command per turn. You may use `&&` or `;` to chain commands within the single `` block. +- If your response cannot be parsed (missing `` tags), the episode ends immediately. +- You have {{ max_steps }} steps to complete the task. +- The `` tag must appear at the END of your response. diff --git a/trinity/common/workflows/connect_the_dots/terminal/prompts/user.jinja2 b/trinity/common/workflows/connect_the_dots/terminal/prompts/user.jinja2 new file mode 100644 index 00000000000..f2404607d6f --- /dev/null +++ b/trinity/common/workflows/connect_the_dots/terminal/prompts/user.jinja2 @@ -0,0 +1,7 @@ +Step {{ current_step }}/{{ max_steps }} +{% if command_output %} + +{{ command_output }} +{% endif %} + +{{ terminal_prompt }} \ No newline at end of file diff --git a/trinity/common/workflows/connect_the_dots/terminal/task_gen.py b/trinity/common/workflows/connect_the_dots/terminal/task_gen.py new file mode 100644 index 00000000000..b9221d11d32 --- /dev/null +++ b/trinity/common/workflows/connect_the_dots/terminal/task_gen.py @@ -0,0 +1,996 @@ +# -*- coding: utf-8 -*- +""" +Task generation for the simulated terminal environment. + +Every task is fully deterministic given a seed: the seed controls the OS type, +task type, file names, paths, remote server details, and filesystem layout. +""" + +from __future__ import annotations + +import json +import random +from dataclasses import dataclass, field +from enum import Enum +from typing import Any, Dict, List, Optional, Tuple + +from trinity.common.workflows.connect_the_dots.terminal.env import ( + FileMeta, + MachineState, + OSType, + TerminalEnv, + VirtualFS, +) + +# --------------------------------------------------------------------------- +# Task types +# --------------------------------------------------------------------------- + + +class TaskType(Enum): + UPLOAD = "upload" + DOWNLOAD = "download" + RENAME = "rename" + MOVE = "move" + CHMOD = "chmod" + DELETE = "delete" + COPY = "copy" + PACK = "pack" + MKDIR = "mkdir" + + +class CompositeTemplate(Enum): + PACK_UPLOAD = "pack_upload" + DOWNLOAD_EXTRACT = "download_extract" + MKDIR_UPLOAD = "mkdir_upload" + UPLOAD_CHMOD = "upload_chmod" + UPLOAD_DELETE_SOURCE = "upload_delete_source" + PACK_UPLOAD_EXTRACT = "pack_upload_extract" + DOWNLOAD_RENAME = "download_rename" + BACKUP_REPLACE = "backup_replace" + + +# --------------------------------------------------------------------------- +# Pools for randomisation +# --------------------------------------------------------------------------- + +FILENAMES = [ + "report.txt", "data.csv", "image.png", "config.json", "notes.md", + "script.sh", "log.txt", "readme.txt", "database.db", "output.log", + "results.csv", "document.pdf", "settings.ini", "main.py", "index.html", + "style.css", "app.js", "server.py", "deploy.sh", "backup.sql", + "photo.jpg", "diagram.svg", "metrics.json", "requirements.txt", "Makefile", +] + +DIRNAMES = [ + "projects", "documents", "downloads", "backups", "config", + "data", "logs", "workspace", "reports", "media", + "scripts", "output", "staging", "archive", "temp", + "src", "build", "dist", "assets", "uploads", +] + +REMOTE_USERS = ["admin", "deploy", "user", "webmaster", "devops", "ops", "ubuntu"] + +PERMISSION_MODES = ["644", "755", "600", "700", "664", "640", "444", "750"] + +ARCHIVE_FORMATS = ["tar", "tar.gz", "zip"] + + +# --------------------------------------------------------------------------- +# Dataclass +# --------------------------------------------------------------------------- + + +@dataclass +class TerminalTask: + """Complete specification for one terminal task.""" + seed: int + local_os: OSType + task_category: str # "single" or "composite" + task_type: str # TaskType value or CompositeTemplate value + remote_host: str + remote_user: str + description: str # human-readable task description + # goal state as a list of checks + goal_checks: List[Dict[str, Any]] = field(default_factory=list) + # initial filesystem snapshots (for building MachineState at runtime) + local_fs_spec: List[Dict[str, Any]] = field(default_factory=list) + remote_fs_spec: List[Dict[str, Any]] = field(default_factory=list) + # extra metadata + params: Dict[str, Any] = field(default_factory=dict) + + +# --------------------------------------------------------------------------- +# Filesystem helpers +# --------------------------------------------------------------------------- + + +def _home_dir(os_type: OSType) -> str: + """Canonical memfs home directory.""" + if os_type == OSType.WINDOWS: + return "C/Users/user" + elif os_type == OSType.MAC: + return "Users/user" + return "home/user" + + +def _display_path(memfs_path: str, os_type: OSType) -> str: + """Convert memfs path to OS-appropriate display path.""" + if os_type == OSType.WINDOWS: + parts = memfs_path.split("/") + if parts and len(parts[0]) == 1 and parts[0].isalpha(): + drive = parts[0].upper() + rest = "\\".join(parts[1:]) + return f"{drive}:\\{rest}" if rest else f"{drive}:\\" + return "\\".join(parts) + return "/" + memfs_path if memfs_path else "/" + + +def _make_file_content(filename: str, seed: int) -> str: + """Deterministic file content for a given filename and seed.""" + return f"[Content of {filename} | seed={seed} | size=1024]" + + +def _pick_subdir(rng: random.Random, os_type: OSType) -> str: + """Pick a random subdirectory under home.""" + home = _home_dir(os_type) + subdir = rng.choice(DIRNAMES) + return f"{home}/{subdir}" + + +def _pick_remote_dir(rng: random.Random, remote_user: str) -> str: + """Pick a random directory on the remote machine.""" + options = [ + f"home/{remote_user}/{rng.choice(DIRNAMES)}", + f"var/www/html", + f"opt/{rng.choice(DIRNAMES)}", + f"home/{remote_user}", + ] + return rng.choice(options) + + +def _add_distractor_files( + fs_spec: list, base_dir: str, rng: random.Random, seed: int, count: int = 3 +): + """Add some irrelevant files to make the task more realistic.""" + used = {f["path"] for f in fs_spec} + for _ in range(count): + name = rng.choice(FILENAMES) + path = f"{base_dir}/{name}" + if path not in used: + content = _make_file_content(name, seed + hash(name) % 10000) + fs_spec.append({"path": path, "content": content}) + used.add(path) + + +# --------------------------------------------------------------------------- +# Single task generators +# --------------------------------------------------------------------------- + + +def _gen_upload(rng: random.Random, seed: int, os_type: OSType, + remote_host: str, remote_user: str) -> TerminalTask: + filename = rng.choice(FILENAMES) + local_dir = _pick_subdir(rng, os_type) + remote_dir = _pick_remote_dir(rng, remote_user) + content = _make_file_content(filename, seed) + + local_fs = [{"path": f"{local_dir}/{filename}", "content": content}] + _add_distractor_files(local_fs, local_dir, rng, seed) + remote_fs = [{"path": remote_dir, "is_dir": True}] + + local_display = _display_path(f"{local_dir}/{filename}", os_type) + remote_display = f"{remote_user}@{remote_host}:/{remote_dir}/{filename}" + + return TerminalTask( + seed=seed, local_os=os_type, task_category="single", + task_type=TaskType.UPLOAD.value, + remote_host=remote_host, remote_user=remote_user, + description=f"Upload '{local_display}' to {remote_display}", + goal_checks=[ + {"type": "file_exists", "machine": "remote", + "path": f"{remote_dir}/{filename}", "content": content}, + ], + local_fs_spec=local_fs, remote_fs_spec=remote_fs, + params={"filename": filename, "local_dir": local_dir, "remote_dir": remote_dir}, + ) + + +def _gen_download(rng: random.Random, seed: int, os_type: OSType, + remote_host: str, remote_user: str) -> TerminalTask: + filename = rng.choice(FILENAMES) + local_dir = _pick_subdir(rng, os_type) + remote_dir = _pick_remote_dir(rng, remote_user) + content = _make_file_content(filename, seed) + + local_fs = [{"path": local_dir, "is_dir": True}] + _add_distractor_files(local_fs, local_dir, rng, seed) + remote_fs = [{"path": f"{remote_dir}/{filename}", "content": content}] + _add_distractor_files(remote_fs, remote_dir, rng, seed) + + local_display = _display_path(f"{local_dir}/{filename}", os_type) + remote_display = f"{remote_user}@{remote_host}:/{remote_dir}/{filename}" + + return TerminalTask( + seed=seed, local_os=os_type, task_category="single", + task_type=TaskType.DOWNLOAD.value, + remote_host=remote_host, remote_user=remote_user, + description=f"Download '{remote_display}' to '{local_display}'", + goal_checks=[ + {"type": "file_exists", "machine": "local", + "path": f"{local_dir}/{filename}", "content": content}, + ], + local_fs_spec=local_fs, remote_fs_spec=remote_fs, + params={"filename": filename, "local_dir": local_dir, "remote_dir": remote_dir}, + ) + + +def _gen_rename(rng: random.Random, seed: int, os_type: OSType, + remote_host: str, remote_user: str) -> TerminalTask: + is_remote = rng.choice([True, False]) + old_name = rng.choice(FILENAMES) + # Generate a different new name + new_name = rng.choice(FILENAMES) + while new_name == old_name: + new_name = rng.choice(FILENAMES) + content = _make_file_content(old_name, seed) + + if is_remote: + target_dir = _pick_remote_dir(rng, remote_user) + fs_spec = [{"path": f"{target_dir}/{old_name}", "content": content}] + _add_distractor_files(fs_spec, target_dir, rng, seed) + machine_label = "remote" + desc_path = f"{remote_user}@{remote_host}:/{target_dir}" + else: + target_dir = _pick_subdir(rng, os_type) + fs_spec = [{"path": f"{target_dir}/{old_name}", "content": content}] + _add_distractor_files(fs_spec, target_dir, rng, seed) + machine_label = "local" + desc_path = _display_path(target_dir, os_type) + + local_fs = fs_spec if not is_remote else [] + remote_fs = fs_spec if is_remote else [] + + return TerminalTask( + seed=seed, local_os=os_type, task_category="single", + task_type=TaskType.RENAME.value, + remote_host=remote_host, remote_user=remote_user, + description=(f"Rename '{old_name}' to '{new_name}' in " + f"{'remote ' if is_remote else ''}{desc_path}"), + goal_checks=[ + {"type": "file_not_exists", "machine": machine_label, + "path": f"{target_dir}/{old_name}"}, + {"type": "file_exists", "machine": machine_label, + "path": f"{target_dir}/{new_name}", "content": content}, + ], + local_fs_spec=local_fs, remote_fs_spec=remote_fs, + params={"is_remote": is_remote, "old_name": old_name, "new_name": new_name, + "target_dir": target_dir}, + ) + + +def _gen_move(rng: random.Random, seed: int, os_type: OSType, + remote_host: str, remote_user: str) -> TerminalTask: + is_remote = rng.choice([True, False]) + filename = rng.choice(FILENAMES) + content = _make_file_content(filename, seed) + + if is_remote: + src_dir = _pick_remote_dir(rng, remote_user) + dst_dir = _pick_remote_dir(rng, remote_user) + while dst_dir == src_dir: + dst_dir = _pick_remote_dir(rng, remote_user) + fs_spec = [ + {"path": f"{src_dir}/{filename}", "content": content}, + {"path": dst_dir, "is_dir": True}, + ] + _add_distractor_files(fs_spec, src_dir, rng, seed) + machine_label = "remote" + src_display = f"/{src_dir}/{filename}" + dst_display = f"/{dst_dir}/" + else: + src_dir = _pick_subdir(rng, os_type) + dst_dir = _pick_subdir(rng, os_type) + while dst_dir == src_dir: + dst_dir = _pick_subdir(rng, os_type) + fs_spec = [ + {"path": f"{src_dir}/{filename}", "content": content}, + {"path": dst_dir, "is_dir": True}, + ] + _add_distractor_files(fs_spec, src_dir, rng, seed) + machine_label = "local" + src_display = _display_path(f"{src_dir}/{filename}", os_type) + dst_display = _display_path(dst_dir, os_type) + + local_fs = fs_spec if not is_remote else [] + remote_fs = fs_spec if is_remote else [] + + return TerminalTask( + seed=seed, local_os=os_type, task_category="single", + task_type=TaskType.MOVE.value, + remote_host=remote_host, remote_user=remote_user, + description=(f"Move '{src_display}' to '{dst_display}'" + f"{' on remote server' if is_remote else ''}"), + goal_checks=[ + {"type": "file_not_exists", "machine": machine_label, + "path": f"{src_dir}/{filename}"}, + {"type": "file_exists", "machine": machine_label, + "path": f"{dst_dir}/{filename}", "content": content}, + ], + local_fs_spec=local_fs, remote_fs_spec=remote_fs, + params={"is_remote": is_remote, "filename": filename, + "src_dir": src_dir, "dst_dir": dst_dir}, + ) + + +def _gen_chmod(rng: random.Random, seed: int, os_type: OSType, + remote_host: str, remote_user: str) -> TerminalTask: + # chmod only works on Linux/Mac or remote + is_remote = rng.choice([True, False]) + if not is_remote and os_type == OSType.WINDOWS: + is_remote = True # force remote for Windows local + + filename = rng.choice(FILENAMES) + content = _make_file_content(filename, seed) + old_perm = rng.choice(PERMISSION_MODES) + new_perm = rng.choice(PERMISSION_MODES) + while new_perm == old_perm: + new_perm = rng.choice(PERMISSION_MODES) + + if is_remote: + target_dir = _pick_remote_dir(rng, remote_user) + remote_fs = [{"path": f"{target_dir}/{filename}", "content": content, + "permissions": old_perm}] + _add_distractor_files(remote_fs, target_dir, rng, seed) + local_fs = [] + machine_label = "remote" + desc_path = f"{remote_user}@{remote_host}:/{target_dir}/{filename}" + else: + target_dir = _pick_subdir(rng, os_type) + local_fs = [{"path": f"{target_dir}/{filename}", "content": content, + "permissions": old_perm}] + _add_distractor_files(local_fs, target_dir, rng, seed) + remote_fs = [] + machine_label = "local" + desc_path = _display_path(f"{target_dir}/{filename}", os_type) + + return TerminalTask( + seed=seed, local_os=os_type, task_category="single", + task_type=TaskType.CHMOD.value, + remote_host=remote_host, remote_user=remote_user, + description=f"Change permissions of '{desc_path}' to {new_perm}", + goal_checks=[ + {"type": "permission_equals", "machine": machine_label, + "path": f"{target_dir}/{filename}", "value": new_perm}, + ], + local_fs_spec=local_fs, remote_fs_spec=remote_fs, + params={"is_remote": is_remote, "filename": filename, + "target_dir": target_dir, "new_perm": new_perm}, + ) + + +def _gen_delete(rng: random.Random, seed: int, os_type: OSType, + remote_host: str, remote_user: str) -> TerminalTask: + is_remote = rng.choice([True, False]) + filename = rng.choice(FILENAMES) + content = _make_file_content(filename, seed) + + if is_remote: + target_dir = _pick_remote_dir(rng, remote_user) + remote_fs = [{"path": f"{target_dir}/{filename}", "content": content}] + _add_distractor_files(remote_fs, target_dir, rng, seed) + local_fs = [] + machine_label = "remote" + desc_path = f"{remote_user}@{remote_host}:/{target_dir}/{filename}" + else: + target_dir = _pick_subdir(rng, os_type) + local_fs = [{"path": f"{target_dir}/{filename}", "content": content}] + _add_distractor_files(local_fs, target_dir, rng, seed) + remote_fs = [] + machine_label = "local" + desc_path = _display_path(f"{target_dir}/{filename}", os_type) + + return TerminalTask( + seed=seed, local_os=os_type, task_category="single", + task_type=TaskType.DELETE.value, + remote_host=remote_host, remote_user=remote_user, + description=f"Delete the file '{desc_path}'", + goal_checks=[ + {"type": "file_not_exists", "machine": machine_label, + "path": f"{target_dir}/{filename}"}, + ], + local_fs_spec=local_fs, remote_fs_spec=remote_fs, + params={"is_remote": is_remote, "filename": filename, "target_dir": target_dir}, + ) + + +def _gen_copy(rng: random.Random, seed: int, os_type: OSType, + remote_host: str, remote_user: str) -> TerminalTask: + is_remote = rng.choice([True, False]) + filename = rng.choice(FILENAMES) + content = _make_file_content(filename, seed) + + if is_remote: + src_dir = _pick_remote_dir(rng, remote_user) + dst_dir = _pick_remote_dir(rng, remote_user) + while dst_dir == src_dir: + dst_dir = _pick_remote_dir(rng, remote_user) + fs_spec = [ + {"path": f"{src_dir}/{filename}", "content": content}, + {"path": dst_dir, "is_dir": True}, + ] + machine_label = "remote" + src_display = f"/{src_dir}/{filename}" + dst_display = f"/{dst_dir}/{filename}" + else: + src_dir = _pick_subdir(rng, os_type) + dst_dir = _pick_subdir(rng, os_type) + while dst_dir == src_dir: + dst_dir = _pick_subdir(rng, os_type) + fs_spec = [ + {"path": f"{src_dir}/{filename}", "content": content}, + {"path": dst_dir, "is_dir": True}, + ] + machine_label = "local" + src_display = _display_path(f"{src_dir}/{filename}", os_type) + dst_display = _display_path(f"{dst_dir}/{filename}", os_type) + + local_fs = fs_spec if not is_remote else [] + remote_fs = fs_spec if is_remote else [] + + return TerminalTask( + seed=seed, local_os=os_type, task_category="single", + task_type=TaskType.COPY.value, + remote_host=remote_host, remote_user=remote_user, + description=(f"Copy '{src_display}' to '{dst_display}'" + f"{' on remote server' if is_remote else ''}"), + goal_checks=[ + {"type": "file_exists", "machine": machine_label, + "path": f"{src_dir}/{filename}", "content": content}, + {"type": "file_exists", "machine": machine_label, + "path": f"{dst_dir}/{filename}", "content": content}, + ], + local_fs_spec=local_fs, remote_fs_spec=remote_fs, + params={"is_remote": is_remote, "filename": filename, + "src_dir": src_dir, "dst_dir": dst_dir}, + ) + + +def _gen_pack(rng: random.Random, seed: int, os_type: OSType, + remote_host: str, remote_user: str) -> TerminalTask: + is_remote = rng.choice([True, False]) + fmt = rng.choice(ARCHIVE_FORMATS) + # On Windows local, only zip is available (no tar) + if not is_remote and os_type == OSType.WINDOWS and fmt.startswith("tar"): + fmt = "zip" + + dir_name = rng.choice(DIRNAMES) + n_files = rng.randint(2, 4) + files = rng.sample(FILENAMES, min(n_files, len(FILENAMES))) + + if fmt == "zip": + archive_name = f"{dir_name}.zip" + elif fmt == "tar.gz": + archive_name = f"{dir_name}.tar.gz" + else: + archive_name = f"{dir_name}.tar" + + if is_remote: + base_dir = _pick_remote_dir(rng, remote_user) + machine_label = "remote" + else: + base_dir = _pick_subdir(rng, os_type) + machine_label = "local" + + src_dir = f"{base_dir}/{dir_name}" + archive_entries = {} + fs_spec = [] + for f in files: + content = _make_file_content(f, seed) + fs_spec.append({"path": f"{src_dir}/{f}", "content": content}) + archive_entries[f"{dir_name}/{f}"] = content + + local_fs = fs_spec if not is_remote else [] + remote_fs = fs_spec if is_remote else [] + + if is_remote: + desc = f"Pack the directory '/{src_dir}' into '{archive_name}' on the remote server" + else: + desc = f"Pack the directory '{_display_path(src_dir, os_type)}' into '{archive_name}'" + + return TerminalTask( + seed=seed, local_os=os_type, task_category="single", + task_type=TaskType.PACK.value, + remote_host=remote_host, remote_user=remote_user, + description=desc, + goal_checks=[ + {"type": "archive_contains", "machine": machine_label, + "path": f"{base_dir}/{archive_name}", + "entries": archive_entries, "archive_type": fmt}, + ], + local_fs_spec=local_fs, remote_fs_spec=remote_fs, + params={"is_remote": is_remote, "dir_name": dir_name, "archive_name": archive_name, + "fmt": fmt, "base_dir": base_dir, "files": files}, + ) + + +def _gen_mkdir(rng: random.Random, seed: int, os_type: OSType, + remote_host: str, remote_user: str) -> TerminalTask: + is_remote = rng.choice([True, False]) + depth = rng.randint(2, 4) + parts = rng.sample(DIRNAMES, min(depth, len(DIRNAMES))) + + if is_remote: + base = f"home/{remote_user}" + machine_label = "remote" + else: + base = _home_dir(os_type) + machine_label = "local" + + target_path = base + "/" + "/".join(parts) + + if is_remote: + desc = f"Create the directory structure '/{target_path}' on the remote server" + else: + desc = f"Create the directory structure '{_display_path(target_path, os_type)}'" + + return TerminalTask( + seed=seed, local_os=os_type, task_category="single", + task_type=TaskType.MKDIR.value, + remote_host=remote_host, remote_user=remote_user, + description=desc, + goal_checks=[ + {"type": "dir_exists", "machine": machine_label, "path": target_path}, + ], + local_fs_spec=[], remote_fs_spec=[], + params={"is_remote": is_remote, "target_path": target_path}, + ) + + +SINGLE_GENERATORS = { + TaskType.UPLOAD: _gen_upload, + TaskType.DOWNLOAD: _gen_download, + TaskType.RENAME: _gen_rename, + TaskType.MOVE: _gen_move, + TaskType.CHMOD: _gen_chmod, + TaskType.DELETE: _gen_delete, + TaskType.COPY: _gen_copy, + TaskType.PACK: _gen_pack, + TaskType.MKDIR: _gen_mkdir, +} + + +# --------------------------------------------------------------------------- +# Composite task generators +# --------------------------------------------------------------------------- + + +def _gen_pack_upload(rng, seed, os_type, remote_host, remote_user): + """Pack local files into archive, then upload archive to remote.""" + fmt = rng.choice(ARCHIVE_FORMATS) + if os_type == OSType.WINDOWS and fmt.startswith("tar"): + fmt = "zip" + + dir_name = rng.choice(DIRNAMES) + files = rng.sample(FILENAMES, rng.randint(2, 4)) + local_dir = _pick_subdir(rng, os_type) + remote_dir = _pick_remote_dir(rng, remote_user) + + ext = {"tar": ".tar", "tar.gz": ".tar.gz", "zip": ".zip"}[fmt] + archive_name = f"{dir_name}{ext}" + src_dir = f"{local_dir}/{dir_name}" + + local_fs = [] + archive_entries = {} + for f in files: + content = _make_file_content(f, seed) + local_fs.append({"path": f"{src_dir}/{f}", "content": content}) + archive_entries[f"{dir_name}/{f}"] = content + + remote_fs = [{"path": remote_dir, "is_dir": True}] + + local_display = _display_path(src_dir, os_type) + return TerminalTask( + seed=seed, local_os=os_type, task_category="composite", + task_type=CompositeTemplate.PACK_UPLOAD.value, + remote_host=remote_host, remote_user=remote_user, + description=(f"Pack '{local_display}' into a {fmt} archive and upload it to " + f"{remote_user}@{remote_host}:/{remote_dir}/"), + goal_checks=[ + {"type": "archive_contains", "machine": "remote", + "path": f"{remote_dir}/{archive_name}", + "entries": archive_entries, "archive_type": fmt}, + ], + local_fs_spec=local_fs, remote_fs_spec=remote_fs, + params={"fmt": fmt, "dir_name": dir_name, "archive_name": archive_name, + "local_dir": local_dir, "remote_dir": remote_dir, "files": files}, + ) + + +def _gen_download_extract(rng, seed, os_type, remote_host, remote_user): + """Download archive from remote, extract locally.""" + fmt = rng.choice(ARCHIVE_FORMATS) + if os_type == OSType.WINDOWS and fmt.startswith("tar"): + fmt = "zip" + + dir_name = rng.choice(DIRNAMES) + files = rng.sample(FILENAMES, rng.randint(2, 4)) + local_dir = _pick_subdir(rng, os_type) + remote_dir = _pick_remote_dir(rng, remote_user) + + ext = {"tar": ".tar", "tar.gz": ".tar.gz", "zip": ".zip"}[fmt] + archive_name = f"{dir_name}{ext}" + + archive_entries = {} + for f in files: + content = _make_file_content(f, seed) + archive_entries[f"{dir_name}/{f}"] = content + + archive_content = json.dumps({"type": fmt, "entries": archive_entries}) + remote_fs = [ + {"path": f"{remote_dir}/{archive_name}", "content": archive_content, + "archive_type": fmt, "archive_entries": archive_entries}, + ] + local_fs = [{"path": local_dir, "is_dir": True}] + + # Goal: all files extracted locally + checks = [] + for rel, content in archive_entries.items(): + checks.append({ + "type": "file_exists", "machine": "local", + "path": f"{local_dir}/{rel}", "content": content, + }) + + return TerminalTask( + seed=seed, local_os=os_type, task_category="composite", + task_type=CompositeTemplate.DOWNLOAD_EXTRACT.value, + remote_host=remote_host, remote_user=remote_user, + description=(f"Download '{archive_name}' from " + f"{remote_user}@{remote_host}:/{remote_dir}/ " + f"and extract its contents to " + f"'{_display_path(local_dir, os_type)}'"), + goal_checks=checks, + local_fs_spec=local_fs, remote_fs_spec=remote_fs, + params={"fmt": fmt, "dir_name": dir_name, "archive_name": archive_name, + "local_dir": local_dir, "remote_dir": remote_dir}, + ) + + +def _gen_mkdir_upload(rng, seed, os_type, remote_host, remote_user): + """Create directory on remote, then upload file there.""" + filename = rng.choice(FILENAMES) + content = _make_file_content(filename, seed) + local_dir = _pick_subdir(rng, os_type) + # Remote dir that does NOT exist yet + depth = rng.randint(2, 3) + parts = rng.sample(DIRNAMES, depth) + remote_dir = f"home/{remote_user}/" + "/".join(parts) + + local_fs = [{"path": f"{local_dir}/{filename}", "content": content}] + _add_distractor_files(local_fs, local_dir, rng, seed) + remote_fs = [] # remote dir intentionally does not exist + + return TerminalTask( + seed=seed, local_os=os_type, task_category="composite", + task_type=CompositeTemplate.MKDIR_UPLOAD.value, + remote_host=remote_host, remote_user=remote_user, + description=(f"Create the directory '/{remote_dir}' on the remote server, " + f"then upload '{_display_path(local_dir + '/' + filename, os_type)}' there"), + goal_checks=[ + {"type": "dir_exists", "machine": "remote", "path": remote_dir}, + {"type": "file_exists", "machine": "remote", + "path": f"{remote_dir}/{filename}", "content": content}, + ], + local_fs_spec=local_fs, remote_fs_spec=remote_fs, + params={"filename": filename, "local_dir": local_dir, "remote_dir": remote_dir}, + ) + + +def _gen_upload_chmod(rng, seed, os_type, remote_host, remote_user): + """Upload file to remote, then change its permissions.""" + filename = rng.choice(FILENAMES) + content = _make_file_content(filename, seed) + local_dir = _pick_subdir(rng, os_type) + remote_dir = _pick_remote_dir(rng, remote_user) + target_perm = rng.choice(["755", "700", "600", "444"]) + + local_fs = [{"path": f"{local_dir}/{filename}", "content": content}] + remote_fs = [{"path": remote_dir, "is_dir": True}] + + return TerminalTask( + seed=seed, local_os=os_type, task_category="composite", + task_type=CompositeTemplate.UPLOAD_CHMOD.value, + remote_host=remote_host, remote_user=remote_user, + description=(f"Upload '{_display_path(local_dir + '/' + filename, os_type)}' to " + f"{remote_user}@{remote_host}:/{remote_dir}/ " + f"and set its permissions to {target_perm}"), + goal_checks=[ + {"type": "file_exists", "machine": "remote", + "path": f"{remote_dir}/{filename}", "content": content}, + {"type": "permission_equals", "machine": "remote", + "path": f"{remote_dir}/{filename}", "value": target_perm}, + ], + local_fs_spec=local_fs, remote_fs_spec=remote_fs, + params={"filename": filename, "local_dir": local_dir, + "remote_dir": remote_dir, "target_perm": target_perm}, + ) + + +def _gen_upload_delete_source(rng, seed, os_type, remote_host, remote_user): + """Upload file to remote, then delete the local copy.""" + filename = rng.choice(FILENAMES) + content = _make_file_content(filename, seed) + local_dir = _pick_subdir(rng, os_type) + remote_dir = _pick_remote_dir(rng, remote_user) + + local_fs = [{"path": f"{local_dir}/{filename}", "content": content}] + _add_distractor_files(local_fs, local_dir, rng, seed) + remote_fs = [{"path": remote_dir, "is_dir": True}] + + return TerminalTask( + seed=seed, local_os=os_type, task_category="composite", + task_type=CompositeTemplate.UPLOAD_DELETE_SOURCE.value, + remote_host=remote_host, remote_user=remote_user, + description=(f"Upload '{_display_path(local_dir + '/' + filename, os_type)}' to " + f"{remote_user}@{remote_host}:/{remote_dir}/ " + f"and then delete the local copy"), + goal_checks=[ + {"type": "file_exists", "machine": "remote", + "path": f"{remote_dir}/{filename}", "content": content}, + {"type": "file_not_exists", "machine": "local", + "path": f"{local_dir}/{filename}"}, + ], + local_fs_spec=local_fs, remote_fs_spec=remote_fs, + params={"filename": filename, "local_dir": local_dir, "remote_dir": remote_dir}, + ) + + +def _gen_pack_upload_extract(rng, seed, os_type, remote_host, remote_user): + """Pack local dir, upload, extract on remote.""" + fmt = rng.choice(ARCHIVE_FORMATS) + if os_type == OSType.WINDOWS and fmt.startswith("tar"): + fmt = "zip" + + dir_name = rng.choice(DIRNAMES) + files = rng.sample(FILENAMES, rng.randint(2, 4)) + local_dir = _pick_subdir(rng, os_type) + remote_dir = _pick_remote_dir(rng, remote_user) + + src_dir = f"{local_dir}/{dir_name}" + local_fs = [] + expected_files = {} + for f in files: + content = _make_file_content(f, seed) + local_fs.append({"path": f"{src_dir}/{f}", "content": content}) + expected_files[f"{dir_name}/{f}"] = content + + remote_fs = [{"path": remote_dir, "is_dir": True}] + + # Goal: files extracted on remote + checks = [] + for rel, content in expected_files.items(): + checks.append({ + "type": "file_exists", "machine": "remote", + "path": f"{remote_dir}/{rel}", "content": content, + }) + + return TerminalTask( + seed=seed, local_os=os_type, task_category="composite", + task_type=CompositeTemplate.PACK_UPLOAD_EXTRACT.value, + remote_host=remote_host, remote_user=remote_user, + description=(f"Pack '{_display_path(src_dir, os_type)}' into a {fmt} archive, " + f"upload to {remote_user}@{remote_host}:/{remote_dir}/, " + f"and extract it there"), + goal_checks=checks, + local_fs_spec=local_fs, remote_fs_spec=remote_fs, + params={"fmt": fmt, "dir_name": dir_name, "local_dir": local_dir, + "remote_dir": remote_dir, "files": files}, + ) + + +def _gen_download_rename(rng, seed, os_type, remote_host, remote_user): + """Download file from remote, rename locally.""" + old_name = rng.choice(FILENAMES) + new_name = rng.choice(FILENAMES) + while new_name == old_name: + new_name = rng.choice(FILENAMES) + content = _make_file_content(old_name, seed) + local_dir = _pick_subdir(rng, os_type) + remote_dir = _pick_remote_dir(rng, remote_user) + + remote_fs = [{"path": f"{remote_dir}/{old_name}", "content": content}] + local_fs = [{"path": local_dir, "is_dir": True}] + + return TerminalTask( + seed=seed, local_os=os_type, task_category="composite", + task_type=CompositeTemplate.DOWNLOAD_RENAME.value, + remote_host=remote_host, remote_user=remote_user, + description=(f"Download '{old_name}' from " + f"{remote_user}@{remote_host}:/{remote_dir}/ " + f"and rename it to '{new_name}' in " + f"'{_display_path(local_dir, os_type)}'"), + goal_checks=[ + {"type": "file_exists", "machine": "local", + "path": f"{local_dir}/{new_name}", "content": content}, + {"type": "file_not_exists", "machine": "local", + "path": f"{local_dir}/{old_name}"}, + ], + local_fs_spec=local_fs, remote_fs_spec=remote_fs, + params={"old_name": old_name, "new_name": new_name, + "local_dir": local_dir, "remote_dir": remote_dir}, + ) + + +def _gen_backup_replace(rng, seed, os_type, remote_host, remote_user): + """Backup existing remote file to .bak, upload new version.""" + filename = rng.choice(FILENAMES) + old_content = _make_file_content(filename, seed) + new_content = _make_file_content(filename, seed + 99999) + local_dir = _pick_subdir(rng, os_type) + remote_dir = _pick_remote_dir(rng, remote_user) + + local_fs = [{"path": f"{local_dir}/{filename}", "content": new_content}] + remote_fs = [{"path": f"{remote_dir}/{filename}", "content": old_content}] + + bak_name = filename + ".bak" + + return TerminalTask( + seed=seed, local_os=os_type, task_category="composite", + task_type=CompositeTemplate.BACKUP_REPLACE.value, + remote_host=remote_host, remote_user=remote_user, + description=(f"On the remote server, rename '/{remote_dir}/{filename}' to " + f"'{bak_name}' as a backup, then upload the new version from " + f"'{_display_path(local_dir + '/' + filename, os_type)}'"), + goal_checks=[ + {"type": "file_exists", "machine": "remote", + "path": f"{remote_dir}/{bak_name}", "content": old_content}, + {"type": "file_exists", "machine": "remote", + "path": f"{remote_dir}/{filename}", "content": new_content}, + ], + local_fs_spec=local_fs, remote_fs_spec=remote_fs, + params={"filename": filename, "bak_name": bak_name, + "local_dir": local_dir, "remote_dir": remote_dir}, + ) + + +COMPOSITE_GENERATORS = { + CompositeTemplate.PACK_UPLOAD: _gen_pack_upload, + CompositeTemplate.DOWNLOAD_EXTRACT: _gen_download_extract, + CompositeTemplate.MKDIR_UPLOAD: _gen_mkdir_upload, + CompositeTemplate.UPLOAD_CHMOD: _gen_upload_chmod, + CompositeTemplate.UPLOAD_DELETE_SOURCE: _gen_upload_delete_source, + CompositeTemplate.PACK_UPLOAD_EXTRACT: _gen_pack_upload_extract, + CompositeTemplate.DOWNLOAD_RENAME: _gen_download_rename, + CompositeTemplate.BACKUP_REPLACE: _gen_backup_replace, +} + + +# --------------------------------------------------------------------------- +# Main entry point +# --------------------------------------------------------------------------- + + +def generate_task( + seed: int, + composite_ratio: float = 0.5, + forced_task_type: Optional[str] = None, +) -> TerminalTask: + """Generate a complete terminal task from a seed. + + Args: + seed: Random seed that fully determines the task. + composite_ratio: Probability of generating a composite task (0-1). + forced_task_type: If set, force this task type (TaskType or CompositeTemplate value). + """ + rng = random.Random(seed) + + # Pick OS + local_os = rng.choice(list(OSType)) + + # Pick remote server + remote_host = f"192.168.1.{rng.randint(10, 99)}" + remote_user = rng.choice(REMOTE_USERS) + + # Decide task type + if forced_task_type: + # Try single first, then composite + try: + tt = TaskType(forced_task_type) + return SINGLE_GENERATORS[tt](rng, seed, local_os, remote_host, remote_user) + except ValueError: + ct = CompositeTemplate(forced_task_type) + return COMPOSITE_GENERATORS[ct](rng, seed, local_os, remote_host, remote_user) + + is_composite = rng.random() < composite_ratio + if is_composite: + template = rng.choice(list(CompositeTemplate)) + return COMPOSITE_GENERATORS[template](rng, seed, local_os, remote_host, remote_user) + else: + task_type = rng.choice(list(TaskType)) + return SINGLE_GENERATORS[task_type](rng, seed, local_os, remote_host, remote_user) + + +# --------------------------------------------------------------------------- +# Build environment from task +# --------------------------------------------------------------------------- + + +def build_env_from_task(task: TerminalTask, max_steps: int = 15): + """Construct a TerminalEnv from a TerminalTask specification.""" + from trinity.common.workflows.connect_the_dots.terminal.commands import ( + build_command_registry, + ) + + os_type = task.local_os + + # Build local machine + home = _home_dir(os_type) + local_fs = VirtualFS() + local_fs.makedirs(home) + if os_type == OSType.WINDOWS: + env_vars = { + "USERPROFILE": _display_path(home, os_type), + "USERNAME": "user", + "CD": _display_path(home, os_type), + } + else: + env_vars = { + "HOME": _display_path(home, os_type), + "USER": "user", + "PWD": _display_path(home, os_type), + } + + for spec in task.local_fs_spec: + path = spec["path"] + if spec.get("is_dir"): + local_fs.makedirs(path) + else: + parent = "/".join(path.split("/")[:-1]) + if parent: + local_fs.makedirs(parent) + local_fs.writetext(path, spec.get("content", "")) + meta = local_fs.get_meta(path) + if "permissions" in spec: + meta.permissions = spec["permissions"] + if "archive_type" in spec: + meta.archive_type = spec["archive_type"] + meta.archive_entries = spec.get("archive_entries") + + local = MachineState( + os_type=os_type, hostname="localhost", username="user", + home_dir=home, cwd=home, fs=local_fs, env_vars=env_vars, + ) + + # Build remote machine + remote_home = f"home/{task.remote_user}" + remote_fs = VirtualFS() + remote_fs.makedirs(remote_home) + + for spec in task.remote_fs_spec: + path = spec["path"] + if spec.get("is_dir"): + remote_fs.makedirs(path) + else: + parent = "/".join(path.split("/")[:-1]) + if parent: + remote_fs.makedirs(parent) + remote_fs.writetext(path, spec.get("content", "")) + meta = remote_fs.get_meta(path) + if "permissions" in spec: + meta.permissions = spec["permissions"] + if "archive_type" in spec: + meta.archive_type = spec["archive_type"] + meta.archive_entries = spec.get("archive_entries") + + remote = MachineState( + os_type=OSType.LINUX, hostname=task.remote_host, username=task.remote_user, + home_dir=remote_home, cwd=remote_home, fs=remote_fs, + env_vars={ + "HOME": f"/home/{task.remote_user}", + "USER": task.remote_user, + "IP": task.remote_host, + }, + ) + + handlers = build_command_registry() + return TerminalEnv(local, remote, handlers, max_steps=max_steps) diff --git a/trinity/common/workflows/connect_the_dots/terminal/workflow.py b/trinity/common/workflows/connect_the_dots/terminal/workflow.py new file mode 100644 index 00000000000..dc1e9f0cbaa --- /dev/null +++ b/trinity/common/workflows/connect_the_dots/terminal/workflow.py @@ -0,0 +1,209 @@ +# -*- coding: utf-8 -*- +""" +CoD Terminal workflow for file operation tasks. + +Extends AsyncCoDMultiStepWorkflow following the same patterns as +CoDFrozenLakeWorkflow and CoDRandomAlchemyWorkflow. +""" + +from typing import List, Optional, Tuple + +from trinity.common.experience import Experience +from trinity.common.models.model import ModelWrapper +from trinity.common.workflows.connect_the_dots.base_workflow import ( + AsyncCoDMultiStepWorkflow, +) +from trinity.common.workflows.connect_the_dots.utils import extract_content_between_keys +from trinity.common.workflows.connect_the_dots.terminal.env import ( + OSType, + TerminalEnv, +) +from trinity.common.workflows.connect_the_dots.terminal.goal_check import ( + check_goal, +) +from trinity.common.workflows.connect_the_dots.terminal.prompts import ( + load_system_prompt, + load_user_prompt, +) +from trinity.common.workflows.connect_the_dots.terminal.task_gen import ( + TerminalTask, + build_env_from_task, + generate_task, +) +from trinity.common.workflows.workflow import Task + + +_OS_DISPLAY = { + OSType.WINDOWS: "Windows", + OSType.MAC: "macOS", + OSType.LINUX: "Linux", +} + + +def _parse_action(response: str) -> Optional[str]: + """Extract command from ... tags.""" + content, success = extract_content_between_keys(response, "", "") + if not success: + return None + cmd = content.strip() + return cmd if cmd else None + + +class CoDTerminalWorkflow(AsyncCoDMultiStepWorkflow): + """CoD workflow for simulated terminal file-operation tasks.""" + + is_async: bool = True + + def __init__( + self, + model: ModelWrapper, + task: Task, + auxiliary_models: Optional[List] = None, + use_openai_client: bool = False, + ): + super().__init__( + task=task, + model=model, + auxiliary_models=auxiliary_models, + use_openai_client=use_openai_client, + ) + + workflow_args = task.workflow_args if hasattr(task, "workflow_args") else {} + self.agent_max_steps: int = workflow_args.get("agent_max_steps", 15) + self.max_response_tokens_restraint = workflow_args.get( + "max_response_tokens_restraint", None + ) + composite_ratio: float = workflow_args.get("composite_ratio", 0.5) + + # Reconstruct task from seed + raw_task = task.raw_task if hasattr(task, "raw_task") else {} + self.seed: int = raw_task.get("seed", 42) + forced_task_type = raw_task.get("forced_task_type", None) + + self.terminal_task: TerminalTask = generate_task( + seed=self.seed, + composite_ratio=composite_ratio, + forced_task_type=forced_task_type, + ) + self.env: Optional[TerminalEnv] = None + + # State + self.last_output: str = "" + self.format_error: bool = False + self.early_completion: bool = False + + async def run_async(self) -> List[Experience]: + self.env = build_env_from_task(self.terminal_task, max_steps=self.agent_max_steps) + self.last_output = "" + self.format_error = False + self.early_completion = False + + os_type = self.terminal_task.local_os + sys_prompt = load_system_prompt( + os_type=os_type.value, + os_name=_OS_DISPLAY[os_type], + remote_user=self.terminal_task.remote_user, + remote_host=self.terminal_task.remote_host, + task_description=self.terminal_task.description, + max_steps=self.agent_max_steps, + ) + sys_prompt = self._augment_system_prompt(sys_prompt) + + self.memory.clear() + self.memory.append({"role": "system", "content": sys_prompt}) + + return await super().run_async() + + async def step_async(self, step_num: int) -> Tuple[bool, List[Experience]]: + # Build user prompt + terminal_prompt = self.env.get_prompt_string() + user_content = load_user_prompt( + current_step=step_num + 1, + max_steps=self.agent_max_steps, + command_output=self.last_output if step_num > 0 else None, + terminal_prompt=terminal_prompt, + ) + + if self.icl_examples and step_num == 0: + user_content = ( + f"{user_content}\n\n" + f"Here are some reference examples:\n\n{self.icl_examples}" + ) + + self.memory.append({"role": "user", "content": user_content}) + + if self.reply_prefix: + self.memory.append({"role": "assistant", "content": self.reply_prefix}) + + # Get model response + experiences = await self.model.chat_async(self.memory) + response_text = experiences[0].response_text + self.memory.append({"role": "assistant", "content": response_text}) + + # Store prompt info on experiences + sys_prompt = ( + self.memory[0]["content"] + if self.memory and self.memory[0]["role"] == "system" + else "" + ) + for exp in experiences: + exp.info["sys_prompt"] = sys_prompt + exp.info["user_prompt"] = user_content + + # Parse action + action = _parse_action(response_text) + if action is None: + self.format_error = True + self.last_output = ( + "ERROR: Could not parse command from ... tags. " + "Episode terminated." + ) + return False, experiences + + # Execute command via gym-style step + def _goal_fn(env): + return check_goal(self.terminal_task, env) + + observation, reward, done, info = self.env.step( + action, goal_check_fn=_goal_fn + ) + self.last_output = observation + + if info.get("early_completion"): + self.early_completion = True + self.last_output = "Task completed successfully." + + return not done, experiences + + async def reward_async(self, exps: List[Experience]) -> float: + if self.format_error: + reward = 0.0 + elif self.early_completion: + reward = 1.0 + else: + reward = check_goal(self.terminal_task, self.env) + + if exps: + # Feedback + if self.format_error: + feedback = "Failed: Episode terminated due to action format error." + elif reward > 0.5: + feedback = "Success! Task completed correctly." + else: + feedback = "Failed: Goal state not achieved." + + exps[-1].info["feedback"] = feedback + if exps[-1].metrics is None: + exps[-1].metrics = {} + exps[-1].metrics["format_error_termination"] = ( + 1.0 if self.format_error else 0.0 + ) + exps[-1].metrics["early_completion"] = ( + 1.0 if self.early_completion else 0.0 + ) + + return reward + + @property + def max_step_num(self) -> int: + return self.agent_max_steps diff --git a/trinity/common/workflows/connect_the_dots/update_context_workflow.py b/trinity/common/workflows/connect_the_dots/update_context_workflow.py new file mode 100644 index 00000000000..5ff76f77507 --- /dev/null +++ b/trinity/common/workflows/connect_the_dots/update_context_workflow.py @@ -0,0 +1,208 @@ +# -*- coding: utf-8 -*- +"""Update-context workflows for CoD packs. + +After a task in a pack is solved, the update-context episode reads the previous +context together with the solved trajectory and produces an updated context that +conditions the later tasks in the pack. + +``AsyncCoDUpdateContextWorkflow`` is the default implementation: it builds the +prompt, parses the response and scores it with a length-shaped reward, generating +through ``chat_async``. ``AsyncCoDUpdateContextAgentWorkflow`` inherits all of +that and replaces only the generation path with a single-shot AgentScope agent. +Override ``build_messages`` / ``parse_context`` / ``compute_reward`` for a +different context-update strategy. +""" + +from dataclasses import asdict +from typing import List, Optional, Tuple + +from trinity.common.experience import Experience +from trinity.common.models.model import ModelWrapper +from trinity.common.workflows.connect_the_dots.agentscope_utils import ( + build_agentscope_single_turn_agent, +) +from trinity.common.workflows.connect_the_dots.cod_utils import CoDPrompts +from trinity.common.workflows.workflow import Task, Workflow + + +class AsyncCoDUpdateContextWorkflow(Workflow): + """Async workflow for the CoD update-context episode. + + ``CoDWorkflow`` runs it once per task transition within a pack. The previous + context, solved trajectory, reward and feedback are injected via + ``set_context_inputs`` before each ``run_async``. The default strategy + iteratively refines a ``Hints:`` block and scores it with a length-shaped + reward; generation goes through ``chat_async``. + """ + + can_reset: bool = True + can_repeat: bool = False + is_async: bool = True + + def __init__( + self, + *, + task: Task, + model: ModelWrapper, + auxiliary_models: Optional[List[ModelWrapper]] = None, + ): + super().__init__(task=task, model=model, auxiliary_models=auxiliary_models) + # Inputs for the next context update, set via set_context_inputs. + self.prev_context: str = "" + self.trajectory: str = "" + self.reward: float = 0.0 + self.feedback: str = "" + self.reset(task) + + def reset(self, task: Task): + """Bind task-derived config and clear the per-call inputs.""" + self.task = task + self.prev_context = "" + self.trajectory = "" + self.reward = 0.0 + self.feedback = "" + + @property + def rollout_args(self): + return asdict(self.task.rollout_args) + + def set_context_inputs( + self, *, prev_context: str, trajectory: str, reward: float, feedback: str + ) -> None: + """Set the inputs for the next context update.""" + self.prev_context = prev_context + self.trajectory = trajectory + self.reward = reward + self.feedback = feedback + + def _append_token_limit(self, prompt: str) -> str: + """Append a response-length instruction when configured.""" + max_tokens = self.task.workflow_args.get("max_response_tokens_restraint") + if max_tokens: + prompt += f"\n\nPlease limit your response to {max_tokens} tokens." + return prompt + + def build_messages(self) -> List[dict]: + """Build the chat messages that ask the model to refine the hints.""" + hint_example = self.task.workflow_args.get("hint_example", False) + sys_prompt = CoDPrompts.sys_prompt_gen_hint_iteratively(hint_example) + sys_prompt = self._append_token_limit(sys_prompt) + user_prompt = CoDPrompts.user_prompt_gen_hint_iteratively( + prev_hint=self.prev_context, + trajectory=self.trajectory, + reward=self.reward, + feedback=self.feedback, + ) + return [ + {"role": "system", "content": sys_prompt}, + {"role": "user", "content": user_prompt}, + ] + + def parse_context(self, response: str) -> Tuple[str, bool]: + """Extract the updated context and a parse-success flag from the response.""" + return CoDPrompts.extract_hint(response) + + def compute_reward(self, exp: Experience, parse_success: bool) -> float: + """Length-shaped penalty for the generated context. + + Penalizes a parse failure, an over-short response, or an over-long one + ramped from ``len_zero_penalty`` to ``len_max_penalty``, scaled by + ``hint_penalty_coef``; returns 0 when these penalties are unset. Override + for a different reward. + """ + args = self.task.workflow_args + hint_penalty_coef = args.get("hint_penalty_coef", 0.0) + len_zero_penalty = args.get("len_zero_penalty", None) + len_max_penalty = args.get("len_max_penalty", None) + len_min_penalty = args.get("len_min_penalty", None) + + # Parse failure is the harshest signal: full penalty regardless of length. + if not parse_success: + return -1.0 * hint_penalty_coef + if (len_zero_penalty is None) or (len_max_penalty is None): + return 0.0 + assert ( + len_zero_penalty < len_max_penalty + ), "len_zero_penalty must be smaller than len_max_penalty." + + resp_len = len(exp.tokens) - exp.prompt_length + if len_min_penalty is not None and resp_len < len_min_penalty: + return -1.0 * hint_penalty_coef + ramp = min( + 1.0, + max( + 0.0, + (resp_len - len_zero_penalty) / (len_max_penalty - len_zero_penalty), + ), + ) + return -1.0 * hint_penalty_coef * ramp + + def _finalize(self, exp: Experience, messages: List[dict]) -> List[Experience]: + """Parse the response, score it, and attach the context to ``exp.info``.""" + new_context, parse_success = self.parse_context(exp.response_text or "") + exp.reward = self.compute_reward(exp, parse_success) + + if exp.metrics is None: + exp.metrics = {} + exp.metrics["hint_parse_success"] = 1.0 * parse_success + + exp.info["sys_prompt"] = messages[0]["content"] + exp.info["user_prompt"] = messages[-1]["content"] + exp.info["hint_parse_success"] = parse_success + exp.info["hint"] = new_context + exp.info["prev_hint"] = self.prev_context + return [exp] + + async def run_async(self) -> List[Experience]: + """Generate one updated context and return it as a single experience. + + The updated context is returned via ``exp.info["hint"]`` so callers do + not depend on the parsing details. + """ + messages = self.build_messages() + rollout_args = self.rollout_args + rollout_args["n"] = 1 + exps = await self.model.chat_async(messages, **rollout_args) + return self._finalize(exps[0], messages) + + +class AsyncCoDUpdateContextAgentWorkflow(AsyncCoDUpdateContextWorkflow): + """Update-context episode whose generation runs through an AgentScope agent. + + Prompt building, parsing and reward are inherited from the base workflow; + only the generation path is replaced with a single-shot AgentScope agent on + an isolated-history model clone. + """ + + requires_isolated_model_history: bool = True + + def __init__( + self, + *, + task: Task, + model: ModelWrapper, + auxiliary_models: Optional[List[ModelWrapper]] = None, + ): + super().__init__(task=task, model=model, auxiliary_models=auxiliary_models) + self._agent = None + self._agent_rollout_args = None + + async def _ensure_agent(self): + rollout_args = asdict(self.task.rollout_args) + if self._agent is None or self._agent_rollout_args != rollout_args: + self._agent = await build_agentscope_single_turn_agent( + name="cod_context_update", + model=self.model, + rollout_args=self.task.rollout_args, + ) + self._agent_rollout_args = rollout_args + return self._agent + + async def run_async(self) -> List[Experience]: + """Generate one updated context through the AgentScope agent.""" + self.model.history.clear() + messages = self.build_messages() + agent = await self._ensure_agent() + await agent(messages) + exp = self.model.extract_experience_from_history()[-1] + return self._finalize(exp, messages) diff --git a/trinity/common/workflows/connect_the_dots/utils.py b/trinity/common/workflows/connect_the_dots/utils.py new file mode 100644 index 00000000000..beb726a26e1 --- /dev/null +++ b/trinity/common/workflows/connect_the_dots/utils.py @@ -0,0 +1,85 @@ +"""General utils""" + +import hashlib +import json +import xml.etree.ElementTree as ET +from typing import Iterable, Optional, Tuple + + +def compute_stable_pack_seed(identities: Iterable[Tuple[int, int]]) -> int: + """Compute a deterministic seed from a pack's dataset identities.""" + payload = json.dumps(sorted(identities), separators=(",", ":")) + return int.from_bytes(hashlib.sha256(payload.encode()).digest()[:8], "big") + + +def extract_content_between_keys( + response: str, + key_start: str, + key_end: str, +) -> Tuple[str, bool]: + """Extract content in response between key_start and key_end. + + Strict success condition: + there is one and only one match for key_start / key_end, + and they should be in the correct order. + + Returns: + content (str): extracted content, "null" if failed. + success (bool): whether extraction is successful. + + TODO: consider requiring that key_end must appear at the end of response. + """ + + idx_start, idx_start_r = response.find(key_start), response.rfind(key_start) + idx_end, idx_end_r = response.find(key_end), response.rfind(key_end) + + if (idx_start == -1) or (idx_start != idx_start_r) or (idx_end == -1) or (idx_end != idx_end_r): + return "null", False + + if idx_start > idx_end: + return "null", False + + return response[(idx_start + len(key_start)) : idx_end], True + + +def parse_xml_answer( + response: str, + list_tags: Optional[set[str]] = None, +) -> Tuple[Optional[dict], str]: + """Parse one XML action wrapped in a unique ```` block.""" + content, success = extract_content_between_keys(response, "", "") + if not success: + return None, "expected_exactly_one_answer_tag" + + def element_value(element: ET.Element): + children = list(element) + if not children: + if element.attrib: + return dict(element.attrib) + return (element.text or "").strip() + if (element.text or "").strip() or any( + (child.tail or "").strip() for child in children + ): + raise ValueError("unexpected XML text") + if element.tag in (list_tags or set()): + return [element_value(child) for child in children] + value = dict(element.attrib) + for child in children: + child_value = element_value(child) + if child.tag in value: + current = value[child.tag] + value[child.tag] = ( + current + [child_value] + if isinstance(current, list) + else [current, child_value] + ) + else: + value[child.tag] = child_value + return value + + try: + action = ET.fromstring(content) + args = element_value(action) + return {"action": action.tag, "args": {} if args == "" else args}, "" + except (ET.ParseError, ValueError): + return None, "invalid_answer_xml" diff --git a/trinity/common/workflows/workflow.py b/trinity/common/workflows/workflow.py index 25853322fd9..c06cfe9804a 100644 --- a/trinity/common/workflows/workflow.py +++ b/trinity/common/workflows/workflow.py @@ -110,6 +110,16 @@ def __init__( self.run_id_base = 0 self.logger = get_logger(__name__) + # !!! Additional fields for CoD project !!! + # - hint: injected into system prompt to guide model's reasoning + # - max_response_tokens_restraint: injected into system prompt to guide model's response length + # - icl_examples: in-context learning examples, added only in first turn + # - reply_prefix: prefix for model's reply to guide model's generation + self.hint: Optional[str] = None + self.max_response_tokens_restraint: Optional[int] = None + self.icl_examples: Optional[str] = None + self.reply_prefix: Optional[str] = None # not really used in CoD + @property def resettable(self): """Deprecated, use cls.can_reset instead.""" @@ -132,6 +142,18 @@ def reset(self, task: Task): """Reset the workflow.""" raise NotImplementedError + def set_hint(self, hint: str): + """!!! For CoD: set the hint !!!""" + self.hint = hint + + def set_max_response_tokens_restraint(self, value: int): + """!!! For CoD: set the max response tokens restraint !!!""" + self.max_response_tokens_restraint = value + + def set_icl_examples(self, icl_examples: str): + """!!! For CoD: set ICL examples to be prepended to user message !!!""" + self.icl_examples = icl_examples + def set_repeat_times(self, repeat_times: int, run_id_base: int) -> None: """ Set the number of times to repeat the workflow. @@ -270,6 +292,11 @@ def reset(self, task: Task): else: raise ValueError("`reward_fn` must be a subclass of `RewardFn`") + # !!! Reset fields for CoD !!! + self.hint = None + self.max_response_tokens_restraint = None + self.icl_examples = None + def set_repeat_times(self, repeat_times, run_id_base): self.repeat_times = repeat_times self.task.rollout_args.n = repeat_times @@ -290,6 +317,21 @@ def format_messages(self): return messages +# util for cod +def log_sys_user_prompts_in_exp(messages, responses) -> None: + sys_prompt = "(not found in messages)" + user_prompt = "(not found in messages)" + for j, msg in enumerate(messages): + if j == 0 and msg["role"] == "system": + sys_prompt = msg["content"] + if msg["role"] == "user": + user_prompt = msg["content"] + break + for response in responses: + response.info["sys_prompt"] = sys_prompt + response.info["user_prompt"] = user_prompt + + class SimpleWorkflow(BaseSimpleWorkflow): """A workflow for simple single-round task.""" @@ -318,6 +360,11 @@ def run(self) -> List[Experience]: self.logger.debug( f"self.task_desc: {self.task_desc}, messages: {messages}, response: {response.response_text}, reward: {reward}" ) + + # !!! PATCH FOR COD START !!! + log_sys_user_prompts_in_exp(messages, responses) + # !!! PATCH FOR COD END !!! + return responses @@ -346,6 +393,11 @@ async def run_async(self) -> List[Experience]: self.logger.debug( f"self.task_desc: {self.task_desc}, messages: {messages}, response: {response.response_text}, reward: {reward}" ) + + # !!! PATCH FOR COD START !!! + log_sys_user_prompts_in_exp(messages, responses) + # !!! PATCH FOR COD END !!! + return responses diff --git a/trinity/explorer/explorer.py b/trinity/explorer/explorer.py index 0188b14ee93..70b1f9ec04c 100644 --- a/trinity/explorer/explorer.py +++ b/trinity/explorer/explorer.py @@ -317,6 +317,20 @@ async def explore_step(self) -> bool: self.explore_start_time = time.time() try: tasks = await self.taskset.read() + # !!! PATCH FOR COD START !!! + from trinity.common.workflows.connect_the_dots.cod_workflow import ( + pack_tasks, + ) + + pack_size = self.config.cod.task_pack_size + cod_workflow_args = self.config.cod.cod_workflow_args + tasks = pack_tasks( + tasks, + pack_size, + cod_workflow_args, + pack_strategy=self.config.cod.packing_strategy, + ) + # !!! PATCH FOR COD END !!! except StopAsyncIteration: self.logger.warning("No more tasks to explore. Stop exploring.") await self.finish_current_steps() @@ -418,7 +432,8 @@ async def eval(self): f"Use '{self.config.buffer.explorer_input.default_eval_workflow_type}' for evaluation." ) - for eval_taskset_config in self.config.buffer.explorer_input.eval_tasksets: + eval_taskset_configs = self.config.buffer.explorer_input.eval_tasksets + for eval_taskset_id, eval_taskset_config in enumerate(eval_taskset_configs): self.logger.info( f"Evaluation on {eval_taskset_config.name} at step {self.explore_step_num} started." ) @@ -428,7 +443,27 @@ async def eval(self): eval_tasks = [] while True: try: - eval_tasks.extend(await eval_taskset.read()) + data = await eval_taskset.read() + # !!! PATCH FOR COD START !!! + from trinity.common.workflows.connect_the_dots.cod_workflow import ( + pack_tasks, + ) + + if self.config.cod.eval_task_pack_size: + pack_size = self.config.cod.eval_task_pack_size + else: + pack_size = self.config.cod.task_pack_size + cod_workflow_args = self.config.cod.cod_workflow_args + for task in data: + task.index["taskset_id"] = eval_taskset_id + data = pack_tasks( + data, + pack_size, + cod_workflow_args, + pack_strategy=self.config.cod.packing_strategy, + ) + # !!! PATCH FOR COD END !!! + eval_tasks.extend(data) except StopAsyncIteration: break assert ( diff --git a/trinity/explorer/workflow_runner.py b/trinity/explorer/workflow_runner.py index dcdcd2fb5ea..20135be0d23 100644 --- a/trinity/explorer/workflow_runner.py +++ b/trinity/explorer/workflow_runner.py @@ -486,8 +486,28 @@ def __init__( async def debug(self) -> None: """Run the debug workflow.""" - tasks = await self.taskset.read(batch_size=1) + await self.prepare() + + # !!! ORIGINAL !!! + # tasks = await self.taskset.read_async(batch_size=1) + # task = tasks[0] + # !!! PATCH FOR COD START !!! + pack_size = self.config.cod.task_pack_size + tasks = await self.taskset.read(batch_size=pack_size) + print(f"!!! original number of read tasks: {len(tasks)} !!!") + from trinity.common.workflows.connect_the_dots.cod_workflow import pack_tasks + + cod_workflow_args = self.config.cod.cod_workflow_args + tasks = pack_tasks( + tasks, + pack_size, + cod_workflow_args, + pack_strategy=self.config.cod.packing_strategy, + ) task = tasks[0] + task.batch_id = 1 + # !!! PATCH FOR COD END !!! + self.logger.info(f"Start debugging task:\n{task.raw_task}") if not self.enable_profiling: status, exp_payload = await self.run_task( @@ -506,7 +526,7 @@ async def debug(self) -> None: self.logger.info( f"Debugging failed, extracting {len(experiences)} experiences from history." ) - await self.sqlite_writer.write(experiences) + # await self.sqlite_writer.write_async(experiences) # !!! disable for CoD !!! if status.ok: print(f"Task {task.task_id} completed successfully with metrics:\n{status.metrics}") else: diff --git a/trinity/trainer/verl/losses.py b/trinity/trainer/verl/losses.py index 2ab321fdc8b..4a9ebdb5b05 100644 --- a/trinity/trainer/verl/losses.py +++ b/trinity/trainer/verl/losses.py @@ -23,7 +23,11 @@ class TrinityPolicyLoss: workers via set_loss_fn(). """ - def __init__(self, algo_config: AlgorithmConfig): + def __init__( + self, + algo_config: AlgorithmConfig, + fix_actor_microbatch_loss_scale: bool = False, + ): self.policy_loss_fn = POLICY_LOSS_FN.get(algo_config.policy_loss_fn)( backend="verl", **algo_config.policy_loss_fn_args ) @@ -33,6 +37,7 @@ def __init__(self, algo_config: AlgorithmConfig): ) self.calculate_entropy = algo_config.entropy_loss_fn != "none" self.loss_agg_mode = algo_config.loss_agg_mode + self.fix_actor_microbatch_loss_scale = fix_actor_microbatch_loss_scale self.use_kl_loss = not isinstance(self.kl_loss_fn, DummyKLFn) def __call__( @@ -96,6 +101,18 @@ def __call__( # different scaling semantics. metrics["final_loss"] = policy_loss.detach().item() + # The custom loss bypasses veRL's global token normalization. + if ( + self.fix_actor_microbatch_loss_scale + and self.loss_agg_mode == "token-mean" + ): + policy_loss = ( + policy_loss + * response_mask.sum() + / data["batch_num_tokens"] + * data["dp_size"] + ) + return policy_loss, metrics def __repr__(self) -> str: @@ -106,6 +123,9 @@ def __repr__(self) -> str: ) -def build_trinity_loss(algo_config: AlgorithmConfig) -> TrinityPolicyLoss: +def build_trinity_loss( + algo_config: AlgorithmConfig, + fix_actor_microbatch_loss_scale: bool = False, +) -> TrinityPolicyLoss: """Build a TrinityPolicyLoss instance for veRL's engine API.""" - return TrinityPolicyLoss(algo_config) + return TrinityPolicyLoss(algo_config, fix_actor_microbatch_loss_scale) diff --git a/trinity/trainer/verl/trainer.py b/trinity/trainer/verl/trainer.py index 582786eebbc..3f18becbdd1 100644 --- a/trinity/trainer/verl/trainer.py +++ b/trinity/trainer/verl/trainer.py @@ -449,6 +449,9 @@ async def prepare(self): algo_config=self.algorithm_config, rollout_engine_type=self.global_config.explorer.rollout_model.engine_type, ray_namespace=self.global_config.synchronizer.ray_namespace, + fix_actor_microbatch_loss_scale=( + self.global_config.trainer.fix_actor_microbatch_loss_scale + ), ) self.global_steps = 0 self._load_checkpoint() diff --git a/trinity/trainer/verl/workers.py b/trinity/trainer/verl/workers.py index 6baa2967f0e..99e262dfbf3 100644 --- a/trinity/trainer/verl/workers.py +++ b/trinity/trainer/verl/workers.py @@ -263,19 +263,27 @@ def _cache_state_dict_meta(self): @register(dispatch_mode=Dispatch.ONE_TO_ALL) def set_trinity_config( - self, algo_config: AlgorithmConfig, rollout_engine_type: str, ray_namespace: str + self, + algo_config: AlgorithmConfig, + rollout_engine_type: str, + ray_namespace: str, + fix_actor_microbatch_loss_scale: bool, ): """Set Trinity-specific runtime config on the worker. This is called by VERLTrainer after worker initialization to inject: - The pluggable policy loss, KL loss, and entropy loss from Trinity's algorithm registry + - Token-based dynamic-microbatch loss scaling when enabled - The Ray namespace used to locate Synchronizer and CheckpointMonitor actors """ self._algo_config = algo_config self._ray_namespace = ray_namespace self._rollout_engine_type = rollout_engine_type if self.actor is not None: - loss_fn = build_trinity_loss(algo_config) + loss_fn = build_trinity_loss( + algo_config, + fix_actor_microbatch_loss_scale, + ) self.actor.set_loss_fn(loss_fn) @register(dispatch_mode=Dispatch.ONE_TO_ALL) diff --git a/trinity/trainer/verl_legacy/monkey_patch.py b/trinity/trainer/verl_legacy/monkey_patch.py index cb4c2a194dd..a2722acfc61 100644 --- a/trinity/trainer/verl_legacy/monkey_patch.py +++ b/trinity/trainer/verl_legacy/monkey_patch.py @@ -440,8 +440,10 @@ def state_dict(self, *args, **kwargs): from trinity.common.patch.qwen3_5 import ulysses_gate_delta_net_decorator - for layer in model.model.language_model.layers: - if layer.layer_type == "linear_attention": + language_model = model.model.language_model + for layer_idx, layer in enumerate(language_model.layers): + layer_type = language_model.config.layer_types[layer_idx] + if layer_type == "linear_attention": ulysses_gate_delta_net_decorator(layer.linear_attn, ulysses_sp_size) # Step 3: patch verl.utils.flops_counter diff --git a/trinity/trainer/verl_legacy/verl_config.py b/trinity/trainer/verl_legacy/verl_config.py index 2bf6c09ac01..463ecf8ed27 100644 --- a/trinity/trainer/verl_legacy/verl_config.py +++ b/trinity/trainer/verl_legacy/verl_config.py @@ -107,6 +107,7 @@ class FSDPConfig: wrap_policy: dict = field(default_factory=dict) fsdp_size: int = -1 forward_prefetch: bool = False + use_orig_params: bool = False model_dtype: Optional[str] = None dtype: str = "bfloat16" mixed_precision: dict = field(default_factory=dict)