src/controllers.py

import json
import os


ALLOWED = {
    "learning_rate_multiplier": {0.5, 0.75, 1.0, 1.25, 1.5},
    "beta_multiplier": {0.5, 1.0, 2.0},
    "temperature_delta": {-0.1, 0.0, 0.1},
    "max_completion_length_delta": {0},
    "num_generations_delta": {0},
    "early_stop": {False},
    "rollback_to_best_checkpoint": {False, True},
}


def no_change_decision(reason="Fixed schedule"):
    return {
        "learning_rate_multiplier": 1.0,
        "beta_multiplier": 1.0,
        "temperature_delta": 0.0,
        "max_completion_length_delta": 0,
        "num_generations_delta": 0,
        "early_stop": False,
        "rollback_to_best_checkpoint": False,
        "reason": reason,
    }


def rule_decision(config, metrics, history):
    decision = no_change_decision("No trigger fired")
    reasons = []
    recent = [item.get("eval_accuracy", 0.0) for item in history[-3:]]
    if len(recent) >= 3 and max(recent[1:]) <= recent[0] + 1e-6:
        decision["temperature_delta"] = 0.1
        reasons.append("eval accuracy stalled")
    if (metrics.get("kl_mean") or 0.0) > 0.15:
        decision["learning_rate_multiplier"] = 0.5
        decision["beta_multiplier"] = 2.0
        reasons.append("KL is high")
    if (metrics.get("completion_length_clip_ratio") or 0.0) > 0.2:
        decision["max_completion_length_delta"] = 64
        reasons.append("completions are frequently clipped")
    if (metrics.get("train_reward_std") or 0.0) < 0.02 and (metrics.get("train_reward_mean") or 0.0) < 0.8:
        decision["temperature_delta"] = max(decision["temperature_delta"], 0.1)
        reasons.append("reward diversity is near zero")
    decision["reason"] = "; ".join(reasons) if reasons else decision["reason"]
    return decision


def validate_decision(decision):
    for key, choices in ALLOWED.items():
        if key not in decision or decision[key] not in choices:
            raise ValueError(f"invalid action for {key}: {decision.get(key)}")
    decision["reason"] = str(decision.get("reason", ""))[:500]
    return decision


def parse_llm_decision(text):
    marker = "FINAL_JSON:"
    marker_index = text.rfind(marker)
    if marker_index < 0:
        raise ValueError("LLM response is missing FINAL_JSON marker")
    analysis = text[:marker_index].strip()
    payload = text[marker_index + len(marker):].strip()
    if payload.startswith("```json"):
        payload = payload[len("```json"):].strip()
    elif payload.startswith("```"):
        payload = payload[len("```"):].strip()
    if payload.endswith("```"):
        payload = payload[:-3].strip()
    decision = json.loads(payload)
    decision["controller_analysis"] = analysis[:4000]
    return validate_decision(decision)


CONTROLLER_METRIC_KEYS = [
    "stage",
    "global_train_steps",
    "train_reward_mean",
    "train_reward_std",
    "eval_accuracy",
    "eval_greedy_accuracy",
    "eval_sampled_pass_at_1",
    "eval_sampled_pass_at_4",
    "kl_mean",
    "entropy_mean",
    "grad_norm",
    "last_loss",
    "end_learning_rate",
    "avg_completion_length",
    "completion_length_clip_ratio",
    "wall_clock_seconds",
]


def compact_metrics(metrics):
    return {key: metrics.get(key) for key in CONTROLLER_METRIC_KEYS if key in metrics}


def llm_decision(config, metrics, history):
    fallback = rule_decision(config, metrics, history)
    if config.get("llm_controller_mode", "mock") != "api":
        fallback["reason"] = "Mock LLM controller: " + fallback["reason"]
        return fallback
    try:
        from openai import OpenAI

        client = OpenAI(
            api_key=os.environ["OPENROUTER_API_KEY"],
            base_url="https://openrouter.ai/api/v1",
            default_headers={
                "HTTP-Referer": "https://github.com/llm-zero-lite",
                "X-OpenRouter-Title": "LLMZero-Lite Experiment",
            },
        )
        payload = {
            "current_schedule": {
                key: config[key]
                for key in [
                    "learning_rate",
                    "beta",
                    "temperature",
                    "max_completion_length",
                    "num_generations",
                    "per_device_train_batch_size",
                    "gradient_accumulation_steps",
                ]
            },
            "latest_metrics": compact_metrics(metrics),
            "recent_history": [compact_metrics(item) for item in history[-3:]],
            "allowed_actions": {key: sorted(values) for key, values in ALLOWED.items()},
        }
        messages = [
                {
                    "role": "system",
                    "content": (
                        "You are controlling a staged GRPO experiment. Analyze the learning dynamics, "
                        "including reward trend and variance, evaluation accuracy, KL, entropy, gradient norm, "
                        "completion length, and previous actions. Explain the evidence and tradeoffs in roughly "
                        "150-400 words. Then end with FINAL_JSON: followed by exactly one JSON object and no "
                        "text after it. Every action value must come from allowed_actions. Keep early_stop false "
                        "so all methods receive equal compute. Include a concise reason field in the JSON."
                    ),
                },
                {"role": "user", "content": json.dumps(payload, indent=2)},
        ]
        errors = []
        attempt_messages = messages
        for attempt in range(config.get("llm_controller_max_retries", 3)):
            try:
                response = client.chat.completions.create(
                    model=config["llm_controller_model"],
                    temperature=0.2 if attempt == 0 else 0,
                    max_tokens=config.get("llm_controller_max_tokens", 1200),
                    messages=attempt_messages,
                    timeout=90,
                )
                content = response.choices[0].message.content or ""
                decision = parse_llm_decision(content)
                decision["controller_attempts"] = attempt + 1
                return decision
            except Exception as exc:
                errors.append(f"attempt {attempt + 1}: {type(exc).__name__}: {exc}")
                attempt_messages = messages + [
                    {
                        "role": "user",
                        "content": (
                            "The previous response failed validation. Return only FINAL_JSON: followed by one "
                            "valid JSON object using allowed_actions. Put no text after the JSON."
                        ),
                    },
                ]
        raise RuntimeError("; ".join(errors))
    except Exception as exc:
        if not config.get("llm_controller_fail_open", True):
            raise RuntimeError(f"LLM controller API failed: {exc}") from exc
        decision = no_change_decision(f"GLM failed after retries; no-change fallback: {exc}")
        decision["controller_failed"] = True
        decision["controller_analysis"] = "Controller unavailable or response invalid; preserved the current schedule."
        return decision


def choose_decision(method, config, metrics, history):
    if method == "fixed_grpo":
        return no_change_decision()
    if method == "rule_controller":
        return rule_decision(config, metrics, history)
    if method == "llm_controller":
        return llm_decision(config, metrics, history)
    raise ValueError(f"unknown method: {method}")


def apply_decision(config, decision):
    updated = dict(config)
    updated["learning_rate"] = min(5e-6, max(1e-7, config["learning_rate"] * decision["learning_rate_multiplier"]))
    updated["beta"] = min(0.2, max(0.0, config["beta"] * decision["beta_multiplier"]))
    updated["temperature"] = min(1.1, max(0.5, config["temperature"] + decision["temperature_delta"]))
    updated["max_completion_length"] = min(512, max(64, config["max_completion_length"] + decision["max_completion_length_delta"]))
    candidate_generations = min(16, max(2, config["num_generations"] + decision["num_generations_delta"]))
    generation_batch = config["per_device_train_batch_size"] * config["gradient_accumulation_steps"]
    updated["num_generations"] = candidate_generations if generation_batch % candidate_generations == 0 else config["num_generations"]
    return updated