diff --git a/openclaw-tinker/README.md b/openclaw-tinker/README.md index 91c21c6..f6b9584 100644 --- a/openclaw-tinker/README.md +++ b/openclaw-tinker/README.md @@ -14,7 +14,7 @@ Unified training framework for OpenClaw on [Tinker](https://tinker.build) cloud export TINKER_API_KEY="your-tinker-api-key" # Combined method -python run.py --method combine --model-name Qwen/Qwen3-8B --prm-m 1 --batch-size 16 --w-opd 1.0 --w-rl 1.0 +python run.py --method combine --model-name Qwen/Qwen3-8B --prm-m 1 --batch-size 16 --w-opd 1.0 --w-rl 1.0 --train-epochs 2 # RL method python run.py --method rl --model-name Qwen/Qwen3-8B --prm-m 3 --batch-size 16 @@ -77,6 +77,7 @@ All parameters can be set via CLI flags or environment variables: |------|---------|---------|--------|-------------| | `--w-opd` | `OPENCLAW_COMBINE_W_OPD` | `1.0` | combine | OPD advantage weight | | `--w-rl` | `OPENCLAW_COMBINE_W_RL` | `1.0` | combine | RL advantage weight | +| `--train-epochs` | `TRAIN_EPOCHS` | `1` | all | Duplicate samples N times per rollout batch (combine typically uses 2) | | `--eval-mode` | `EVAL_MODE` | `false` | opd | Enable PRM eval scoring alongside OPD | ### PRM / Hint Judge @@ -113,7 +114,7 @@ On-Policy Distillation using hindsight hints and teacher knowledge: 1. Policy model generates responses; environment provides next_state observations 2. Hint judge extracts key information from next_state into a concise hint 3. Teacher model scores the response (with hint context) to get token-level log-probs -4. Advantage = reverse KL from teacher: `-kl_coef * (student_lp - teacher_lp)` +4. Advantage = per-token distillation: `teacher_lp - student_lp` 5. All samples get reward = 1.0 (no explicit reward signal) 6. Optional `--eval-mode`: also compute PRM eval scores for monitoring @@ -122,12 +123,12 @@ On-Policy Distillation using hindsight hints and teacher knowledge: Weighted combination with three-way sample dispatch: - **OPD+RL samples** (have both next_state and reward): get both advantage components -- **OPD-only samples** (next_state but no reward): only teacher KL advantage +- **OPD-only samples** (next_state but no reward): only teacher distillation advantage - **RL-only samples** (reward but no next_state): only scalar reward advantage Combined advantage per token: ``` -combined_adv_i = w_opd * (-kl_coef * (student_lp_i - teacher_lp_i)) + w_rl * reward +combined_adv_i = w_opd * (teacher_lp_i - student_lp_i) + w_rl * reward ``` ## Tinker Integration diff --git a/openclaw-tinker/config.py b/openclaw-tinker/config.py index 80fc423..a530d84 100644 --- a/openclaw-tinker/config.py +++ b/openclaw-tinker/config.py @@ -42,6 +42,7 @@ class TinkerConfig: # -- Combined method: advantage weights -- w_opd: float = 1.0 w_rl: float = 1.0 + train_epochs: int = 1 # Duplicate samples N times per rollout batch (combine default: 2) # -- OPD: optional eval-mode (compute PRM eval scores alongside OPD) -- eval_mode: bool = False diff --git a/openclaw-tinker/data_formatter.py b/openclaw-tinker/data_formatter.py index 603ac54..0665335 100644 --- a/openclaw-tinker/data_formatter.py +++ b/openclaw-tinker/data_formatter.py @@ -3,12 +3,12 @@ Supports all three methods: RL / OPD (sample_to_datum): - advantage = scalar GRPO reward, optionally + KL penalty from teacher logprobs - Used via: batch_to_datums(batch, advantages, kl_penalty_coef) + advantage = scalar GRPO reward + (teacher_lp - student_lp) if teacher logprobs present + Used via: batch_to_datums(batch, advantages) Combined (sample_to_datum_combined): - combined_adv = w_opd * (-kl_coef * (student_lp - teacher_lp)) + w_rl * reward - Used via: batch_to_datums_combined(batch, w_opd, w_rl, kl_penalty_coef) + combined_adv = w_opd * (teacher_lp - student_lp) + w_rl * reward + Used via: batch_to_datums_combined(batch, w_opd, w_rl) Tinker Datum convention: model_input - input tokens (all but the last token of the full sequence) @@ -103,34 +103,37 @@ def _build_datum(all_tokens: list[int], logprobs: list[float], advantages: list[ # RL / OPD datum conversion # --------------------------------------------------------------------------- -def sample_to_datum(sample: TrainingSample, advantage: float, kl_penalty_coef: float = 0.0): - """Convert one sample + scalar advantage into a Tinker Datum (RL / OPD).""" +def sample_to_datum(sample: TrainingSample, advantage: float): + """Convert one sample + scalar advantage into a Tinker Datum (RL / OPD). + + For OPD samples with teacher_logprobs, the advantage is augmented with + per-token distillation signal: (teacher_lp - student_lp). + This matches Slime's --advantage-estimator on_policy_distillation where + advantage = teacher_logp - old_logp (raw, no coefficient). + """ prompt_len = len(sample.prompt_tokens) all_tokens = sample.prompt_tokens + sample.response_tokens logprobs = [0.0] * (prompt_len - 1) + list(sample.response_logprobs) resp_advantages = [advantage * float(m) for m in sample.loss_mask] - # OPD: add reverse-KL penalty to response advantages - if sample.teacher_logprobs is not None and kl_penalty_coef > 0: + # OPD: add per-token distillation advantage (teacher_lp - student_lp) + if sample.teacher_logprobs is not None: for i in range(min(len(resp_advantages), len(sample.teacher_logprobs))): student_lp = sample.response_logprobs[i] if i < len(sample.response_logprobs) else 0.0 teacher_lp = sample.teacher_logprobs[i] - kl_i = student_lp - teacher_lp - resp_advantages[i] += -kl_penalty_coef * kl_i * float(sample.loss_mask[i]) + resp_advantages[i] += (teacher_lp - student_lp) * float(sample.loss_mask[i]) advantages = [0.0] * (prompt_len - 1) + resp_advantages return _build_datum(all_tokens, logprobs, advantages, sample.session_id, sample.turn_num) -def batch_to_datums( - batch: list[TrainingSample], advantages: list[float], kl_penalty_coef: float = 0.0, -) -> list: +def batch_to_datums(batch: list[TrainingSample], advantages: list[float]) -> list: """Convert a batch of samples + per-sample scalar advantages to Tinker Datums.""" datums = [] for sample, adv in zip(batch, advantages): try: - datums.append(sample_to_datum(sample, adv, kl_penalty_coef=kl_penalty_coef)) + datums.append(sample_to_datum(sample, adv)) except Exception as e: logger.error( "[DataFormatter] FAILED to convert session=%s turn=%d: %s", @@ -147,11 +150,15 @@ def sample_to_datum_combined( sample: TrainingSample, w_opd: float = 1.0, w_rl: float = 1.0, - kl_penalty_coef: float = 0.0, ): """Convert one sample into a Tinker Datum with combined OPD+RL advantages. - combined_adv_i = w_opd * (-kl_coef * (student_lp_i - teacher_lp_i)) + w_rl * reward + combined_adv_i = w_opd * (teacher_lp_i - student_lp_i) + w_rl * reward + + Matches Slime's combine_loss.py: + combined_advantages = w_opd * teacher_advantages + w_rl * grpo_advantages + where teacher_advantages = teacher_logp - old_logp (token-level, raw) + and grpo_advantages = reward broadcast (scalar) """ prompt_len = len(sample.prompt_tokens) all_tokens = sample.prompt_tokens + sample.response_tokens @@ -165,12 +172,12 @@ def sample_to_datum_combined( # RL component: broadcast scalar reward rl_adv = w_rl * sample.reward * mask - # OPD component: reverse-KL from teacher + # OPD component: per-token (teacher_lp - student_lp) opd_adv = 0.0 - if sample.teacher_logprobs is not None and kl_penalty_coef > 0 and i < len(sample.teacher_logprobs): + if sample.teacher_logprobs is not None and i < len(sample.teacher_logprobs): student_lp = sample.response_logprobs[i] if i < len(sample.response_logprobs) else 0.0 teacher_lp = sample.teacher_logprobs[i] - opd_adv = w_opd * (-kl_penalty_coef * (student_lp - teacher_lp)) * mask + opd_adv = w_opd * (teacher_lp - student_lp) * mask resp_advantages.append(rl_adv + opd_adv) @@ -182,14 +189,13 @@ def batch_to_datums_combined( batch: list[TrainingSample], w_opd: float = 1.0, w_rl: float = 1.0, - kl_penalty_coef: float = 0.0, ) -> list: """Convert a batch of samples to Tinker Datums with combined advantages.""" datums = [] for sample in batch: try: datums.append(sample_to_datum_combined( - sample, w_opd=w_opd, w_rl=w_rl, kl_penalty_coef=kl_penalty_coef, + sample, w_opd=w_opd, w_rl=w_rl, )) except Exception as e: logger.error( diff --git a/openclaw-tinker/rollout.py b/openclaw-tinker/rollout.py index 00aa980..79f7ac4 100644 --- a/openclaw-tinker/rollout.py +++ b/openclaw-tinker/rollout.py @@ -129,5 +129,17 @@ async def drain_output_queue( if len(data) < batch_size: await asyncio.sleep(0.05) + # Duplicate samples for multiple training epochs (matches Slime's TRAIN_EPOCHS). + train_epochs = worker.config.train_epochs + if train_epochs > 1: + original = list(data) + for _ in range(train_epochs - 1): + data.extend(original) + print( + f"[Rollout] duplicated {len(original)} groups x{train_epochs} " + f"= {len(data)} groups for training", + flush=True, + ) + print(f"[Rollout] drained {len(data)} groups in {time.time() - start:.2f}s", flush=True) return data diff --git a/openclaw-tinker/run.py b/openclaw-tinker/run.py index 9f3d7de..6a254a7 100644 --- a/openclaw-tinker/run.py +++ b/openclaw-tinker/run.py @@ -50,6 +50,8 @@ def parse_args() -> TinkerConfig: help="OPD advantage weight (combine method only)") parser.add_argument("--w-rl", type=float, default=float(os.getenv("OPENCLAW_COMBINE_W_RL", "1.0")), help="RL advantage weight (combine method only)") + parser.add_argument("--train-epochs", type=int, default=int(os.getenv("TRAIN_EPOCHS", "1")), + help="Duplicate samples N times per rollout batch (combine default: 2)") # OPD eval mode parser.add_argument("--eval-mode", action="store_true", @@ -87,6 +89,7 @@ def parse_args() -> TinkerConfig: resume_from_ckpt=args.resume_from_ckpt, w_opd=args.w_opd, w_rl=args.w_rl, + train_epochs=args.train_epochs, eval_mode=args.eval_mode, prm_m=args.prm_m, prm_temperature=args.prm_temperature, diff --git a/openclaw-tinker/scorers.py b/openclaw-tinker/scorers.py index 3d3b205..5ff0545 100644 --- a/openclaw-tinker/scorers.py +++ b/openclaw-tinker/scorers.py @@ -250,6 +250,15 @@ async def _tinker_teacher_logprobs( prompt_logprobs = response.prompt_logprobs or [] prompt_token_count = len(tokenizer.encode(enhanced_prompt, add_special_tokens=False)) + + # Detect tokenizer drift: prompt_logprobs should cover full_ids + if len(prompt_logprobs) != len(full_ids): + logger.warning( + "[Scorer] tokenizer drift: prompt_logprobs len=%d vs full_ids len=%d " + "(session=%s turn=%d). Logprob alignment may be off.", + len(prompt_logprobs), len(full_ids), session_id, turn_num, + ) + teacher_lps = [ float(lp) if lp is not None else 0.0 for lp in prompt_logprobs[prompt_token_count:] diff --git a/openclaw-tinker/trainer.py b/openclaw-tinker/trainer.py index c7bc405..d596160 100644 --- a/openclaw-tinker/trainer.py +++ b/openclaw-tinker/trainer.py @@ -139,12 +139,11 @@ class Trainer: batch, w_opd=self.config.w_opd, w_rl=self.config.w_rl, - kl_penalty_coef=self.config.kl_loss_coef, ) else: # RL and OPD both use scalar GRPO advantages advantages = compute_grpo_advantages(batch) - datums = batch_to_datums(batch, advantages, kl_penalty_coef=self.config.kl_loss_coef) + datums = batch_to_datums(batch, advantages) if not datums: logger.error("[Trainer] EMPTY batch at step %d — all %d samples failed datum conversion, skipping", step, len(batch))