refactor(train): update train.py for openenv flow

train.py

@@ -10,13 +10,10 @@ Usage:
     python train.py --model unsloth/Qwen2.5-VL-7B-Instruct --epochs 3
     python train.py --model OX-PIXL/SpatialThinker-Qwen2.5-VL-7B --epochs 3
     python train.py --dry_run  # test rewards without GPU
+    python train.py --no_semantic  # use coordinate-based prompts instead of semantic
 """
 import argparse
-import json
-import copy
 import random
-from pathlib import Path
-from typing import Optional
 
 # VL (vision-language) model identifiers — use FastVisionModel for these
 _VL_MODEL_PATTERNS = ['VL', 'vl', 'Vision', 'vision', 'SpatialThinker', 'SpaceThinker']
@@ -40,16 +37,21 @@ def parse_args():
     parser.add_argument('--max_folds', type=int, default=8)
     parser.add_argument('--output', default='origami-grpo')
     parser.add_argument('--level', type=int, default=1, help='Target difficulty level (1-3)')
+    parser.add_argument('--no_semantic', action='store_true',
+                        help='Disable semantic prompts; use coordinate-based target (default: semantic)')
     parser.add_argument('--dry_run', action='store_true', help='Test reward function without training')
     return parser.parse_args()
 
 
-def build_dataset(env, level: int = 1, max_folds: int = 8) -> list[dict]:
+def build_dataset(env, level: int = 1, max_folds: int = 8, semantic: bool = True) -> list[dict]:
     """
     Build a training dataset of prompts from available targets.
     Each item: {'prompt': str, 'target_name': str}
     Repeats each target multiple times to give enough training steps.
+    When semantic=True, uses get_semantic_description for task descriptions.
     """
+    from env.prompts import get_semantic_description, code_as_policy_prompt
+
     all_names = env.available_targets()
 
     # Filter by level; fall back to all targets if none match
@@ -62,8 +64,12 @@ def build_dataset(env, level: int = 1, max_folds: int = 8) -> list[dict]:
 
     items = []
     for name in level_names:
-        obs = env.reset(target_name=name)
-        prompt = obs['prompt']
+        target = env._targets[name]
+        if semantic:
+            desc = get_semantic_description(name, target)
+            prompt = code_as_policy_prompt(target, max_folds=max_folds, semantic_description=desc)
+        else:
+            prompt = code_as_policy_prompt(target, max_folds=max_folds, semantic_description=None)
         items.append({'prompt': prompt, 'target_name': name})
 
     # Repeat each target 10x; ensure at least 50 examples
@@ -143,8 +149,9 @@ def main():
     env = OrigamiEnvironment(mode='code_as_policy', max_steps=args.max_folds)
 
     # Build dataset
-    dataset_items = build_dataset(env, level=args.level, max_folds=args.max_folds)
-    print(f"Dataset: {len(dataset_items)} examples from level-{args.level} targets")
+    use_semantic = not args.no_semantic
+    dataset_items = build_dataset(env, level=args.level, max_folds=args.max_folds, semantic=use_semantic)
+    print(f"Dataset: {len(dataset_items)} examples from level-{args.level} targets (semantic={'on' if use_semantic else 'off'})")
     print(f"Targets: {env.available_targets()}")
 
     # Dry run: test reward function without loading model