@elizaos/training 2.0.0-alpha.11 → 2.0.0-alpha.26

package/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Shaw Walters and elizaOS Contributors
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@elizaos/training",
-  "version": "2.0.0-alpha.11",
+  "version": "2.0.0-alpha.26",
   "description": "ElizaOS RL training pipeline with benchmarking and model publishing support",
   "main": "./src/index.ts",
   "types": "./src/index.ts",
@@ -48,11 +48,11 @@
     "ethers": "^6.16.0",
     "uuid": "^11.1.0"
   },
-  "peerDependencies": {},
   "devDependencies": {
     "@types/node": "^24.10.0",
     "@types/uuid": "^10.0.0",
     "bun-types": "^1.3.2",
     "typescript": "^5.9.3"
-  }
+  },
+  "gitHead": "91dceb1d2e9762af27353dbc764e40e1a0599508"
 }

package/python/scripts/benchmark_should_respond.py

@@ -0,0 +1,190 @@
+
+import argparse
+import json
+import re
+import matplotlib
+matplotlib.use('Agg')
+import matplotlib.pyplot as plt
+import mlx.core as mx
+from mlx_lm import load, generate
+
+def extract_action(text):
+    match = re.search(r"<action>(.*?)</action>", text, re.DOTALL)
+    if match:
+        return match.group(1).strip().upper()
+    return "NONE"
+
+def is_valid_format(text):
+    return "<response>" in text and "</response>" in text and "<action>" in text and "</action>" in text
+
+def run_inference(model_path, adapter_path, prompts, ground_truths, label="Model", verbose=False, args=None):
+    print(f"Loading {label} from {model_path} (adapter: {adapter_path})...")
+    model, tokenizer = load(model_path, adapter_path=adapter_path)
+    
+    results = []
+    correct_count = 0
+    valid_format_count = 0
+    
+    for i, prompt in enumerate(prompts):
+        print(f"Generating {i+1}/{len(prompts)}...")
+        
+        # Apply chat template to match training
+        system_prompt = "You are a helpful assistant."
+        if hasattr(args, "system_prompt_file") and args.system_prompt_file:
+            with open(args.system_prompt_file, 'r') as f:
+                system_prompt = f.read().strip()
+
+        messages = [
+            {"role": "system", "content": system_prompt},
+            {"role": "user", "content": prompt}
+        ]
+        
+        try:
+             # apply_chat_template returns a string if tokenize=False
+             # We want the text prompt that the model sees
+             full_prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+             
+             text = generate(model, tokenizer, prompt=full_prompt, max_tokens=100, verbose=verbose)
+        except TypeError:
+             # Fallback if arguments differ in version
+             full_prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+             text = generate(model, tokenizer, prompt=full_prompt, max_tokens=100, verbose=verbose)
+        
+        pred_action = extract_action(text)
+        true_action = extract_action(ground_truths[i])
+        
+        is_correct = (pred_action == true_action)
+        is_valid = is_valid_format(text)
+        
+        if is_correct: correct_count += 1
+        if is_valid: valid_format_count += 1
+        
+        results.append({
+            "prompt": prompt,
+            "generated": text,
+            "predicted_action": pred_action,
+            "ground_truth_action": true_action,
+            "is_correct": is_correct,
+            "is_valid_format": is_valid
+        })
+        
+    accuracy = correct_count / len(prompts)
+    format_compliance = valid_format_count / len(prompts)
+    
+    return {
+        "model": label,
+        "accuracy": accuracy,
+        "format_compliance": format_compliance,
+        "details": results
+    }
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--base-model", type=str, default="mlx-community/Qwen2.5-1.5B-Instruct-4bit")
+    parser.add_argument("--sft-adapter", type=str, required=True)
+    parser.add_argument("--rl-adapter", type=str, required=False)
+    parser.add_argument("--pure-rl-adapter", type=str, required=False, help="Path to RL-only adapter")
+    parser.add_argument("--data", type=str, required=True)
+    parser.add_argument("--count", type=int, default=10)
+    parser.add_argument("--verbose", action="store_true", help="Enable verbose generation output")
+    parser.add_argument("--system-prompt-file", type=str, help="Path to custom system prompt file")
+    args = parser.parse_args()
+    
+    # Load prompts and ground truths
+    prompts = []
+    ground_truths = []
+    
+    print("Loading data...")
+    with open(args.data, 'r') as f:
+        for line in f:
+            if not line.strip(): continue
+            try:
+                item = json.loads(line)
+                if "messages" in item:
+                    # User message is prompt, Assistant message is ground truth
+                    user_msg = next((m["content"] for m in reversed(item["messages"]) if m["role"] == "user"), None)
+                    asst_msg = next((m["content"] for m in reversed(item["messages"]) if m["role"] == "assistant"), None)
+                    
+                    if user_msg and asst_msg:
+                        prompts.append(user_msg)
+                        ground_truths.append(asst_msg)
+            except:
+                pass
+    
+    # Select subset
+    subset_indices = range(min(args.count, len(prompts)))
+    selected_prompts = [prompts[i] for i in subset_indices]
+    selected_truths = [ground_truths[i] for i in subset_indices]
+
+    if not selected_prompts:
+        print("No prompts found.")
+        return
+
+    print(f"--- Running Benchmark on {len(selected_prompts)} samples ---")
+    
+    metrics = []
+    
+    # 1. Base Model
+    base_metrics = run_inference(args.base_model, None, selected_prompts, selected_truths, "Base Model", verbose=args.verbose, args=args)
+    metrics.append(base_metrics)
+    
+    # 2. SFT Model
+    sft_metrics = run_inference(args.base_model, args.sft_adapter, selected_prompts, selected_truths, "SFT Model", verbose=args.verbose, args=args)
+    metrics.append(sft_metrics)
+    
+    # 3. Pure RL Model (Base -> RL)
+    if args.pure_rl_adapter:
+        try:
+             pure_rl_metrics = run_inference(args.base_model, args.pure_rl_adapter, selected_prompts, selected_truths, "Pure RL Model", verbose=args.verbose, args=args)
+             metrics.append(pure_rl_metrics)
+        except Exception as e:
+             print(f"Failed to run Pure RL benchmark: {e}")
+
+    # 4. SFT+RL Model (SFT -> RL)
+    if args.rl_adapter:
+        rl_metrics = run_inference(args.base_model, args.rl_adapter, selected_prompts, selected_truths, "SFT+RL Model", verbose=args.verbose, args=args)
+        metrics.append(rl_metrics)
+    
+    # Save Results
+    with open("benchmark_results.json", "w") as f:
+        json.dump(metrics, f, indent=2)
+    print("\nSaved detailed results to benchmark_results.json")
+
+    # Print Summary
+    print("\n=== SUMMARY ===")
+    print(f"{'Model':<15} | {'Accuracy':<10} | {'Format':<10}")
+    print("-" * 40)
+    for m in metrics:
+        print(f"{m['model']:<15} | {m['accuracy']:.2%}     | {m['format_compliance']:.2%}")
+
+    # Plotting
+    try:
+        models = [m['model'] for m in metrics]
+        accuracies = [m['accuracy'] for m in metrics]
+        formats = [m['format_compliance'] for m in metrics]
+        
+        x = range(len(models))
+        width = 0.35
+        
+        fig, ax = plt.subplots(figsize=(10, 6))
+        rects1 = ax.bar([i - width/2 for i in x], accuracies, width, label='Action Accuracy')
+        rects2 = ax.bar([i + width/2 for i in x], formats, width, label='Format Compliance')
+        
+        ax.set_ylabel('Score')
+        ax.set_title('ShouldRespond Benchmark Results')
+        ax.set_xticks(x)
+        ax.set_xticklabels(models)
+        ax.legend()
+        
+        ax.bar_label(rects1, padding=3, fmt='%.2f')
+        ax.bar_label(rects2, padding=3, fmt='%.2f')
+        
+        plt.tight_layout()
+        plt.savefig("benchmark_chart.png")
+        print("Saved chart to benchmark_chart.png")
+        
+    except Exception as e:
+        print(f"Failed to create chart: {e}")
+
+if __name__ == "__main__":
+    main()

package/python/scripts/debug_inference.py

@@ -0,0 +1,62 @@
+import mlx.core as mx
+from mlx_lm import load, generate
+
+model_name = "mlx-community/Qwen2.5-1.5B-Instruct-4bit"
+print(f"Loading {model_name}...")
+model, tokenizer = load(model_name)
+
+# Original Prompt
+prompt_orig = """<task>Decide on behalf of Eliza whether they should respond to the message, ignore it or stop the conversation.</task>
+
+<providers>
+[RECENT_MESSAGES]
+User: Hey @Eliza, what's up?
+
+</providers>
+
+<instructions>Decide if Eliza should respond to or interact with the conversation.
+
+IMPORTANT RULES FOR RESPONDING:
+- If YOUR name (Eliza) is directly mentioned -> RESPOND
+- If someone uses a DIFFERENT name (not Eliza) -> IGNORE (they're talking to someone else)
+- If you're actively participating in a conversation and the message continues that thread -> RESPOND
+- If someone tells you to stop or be quiet -> STOP
+- Otherwise -> IGNORE
+
+The key distinction is:
+- "Talking TO Eliza" (your name mentioned, replies to you, continuing your conversation) -> RESPOND
+- "Talking ABOUT Eliza" or to someone else -> IGNORE
+</instructions>
+
+<output>
+Do NOT include any thinking, reasoning, or <think> sections in your response.
+Go directly to the XML response format without any preamble or explanation.
+
+Respond using XML format like this:
+<response>
+  <name>Eliza</name>
+  <reasoning>Your reasoning here</reasoning>
+  <action>RESPOND | IGNORE | STOP</action>
+</response>
+
+IMPORTANT: Your response must ONLY contain the <response></response> XML block above. Do not include any text, thinking, or reasoning before or after this XML block. Start your response immediately with <response> and end with </response>.
+</output>"""
+
+prompts = [
+    ("Original", prompt_orig),
+    ("No Providers Block", prompt_orig.replace("<providers>\n[RECENT_MESSAGES]", "Messages:\n").replace("\n\n</providers>", "")),
+    ("Simplified", """User: Hey @Eliza, what's up?
+Instructions: You are Eliza. Provide an XML response deciding to RESPOND, IGNORE, or STOP.
+Criteria: Respond if addressed directly (@Eliza).
+Format: <response><name>Eliza</name><reasoning>...</reasoning><action>...</action></response>""")
+]
+
+for name, p in prompts:
+    print(f"\n--- Testing {name} ---")
+    messages = [
+        {"role": "system", "content": "You are a helpful assistant."},
+        {"role": "user", "content": p}
+    ]
+    text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    output = generate(model, tokenizer, prompt=text, max_tokens=100, verbose=False)
+    print(output)

package/python/scripts/optimize_prompt_grpo.py

@@ -0,0 +1,269 @@
+import argparse
+import json
+import re
+import random
+import os
+import mlx.core as mx
+from mlx_lm import load, generate
+from dotenv import load_dotenv
+
+# Load environment variables
+load_dotenv()
+
+# -------------------------------------------------------------------------
+# Reward Function (Copied from train_grpo.py)
+# -------------------------------------------------------------------------
+ACTION_RE = re.compile(r"<action>\s*(.*?)\s*</action>", re.DOTALL)
+
+def compute_rewards(prompts, completions):
+    """
+    Reward function for shouldRespond task.
+    Returns list of scores.
+    """
+    rewards = []
+    
+    for prompt, text in zip(prompts, completions):
+        score = 0.0
+        
+        # Parse action
+        action_match = ACTION_RE.search(text)
+        action = action_match.group(1).strip().upper() if action_match else "NONE"
+        
+        # Heuristics
+        last_user_msg = prompt.split("User:")[-1] if "User:" in prompt else prompt
+        is_direct_mention = ("@Eliza" in last_user_msg or "Eliza" in last_user_msg)
+        is_stop = any(w in last_user_msg.lower() for w in ["stop", "shut up", "quiet", "be quiet"])
+        is_continuation = "Eliza:" in prompt
+        is_ambiguous = any(w in last_user_msg.lower() for w in ["anyone", "anybody", "help", "assist", "question", "somebody"])
+        should_respond = is_direct_mention or is_continuation or is_ambiguous
+        
+        # Scoring
+        if is_stop:
+            if action == "STOP": score += 1.0
+            elif action == "IGNORE": score += 0.3
+            else: score -= 0.3
+        elif should_respond:
+            if action == "RESPOND": score += 1.0
+            else: score -= 0.3
+        else: # should ignore
+            if action == "IGNORE": score += 1.0
+            else: score -= 0.3
+            
+        # Format bonus
+        if "<response>" in text and "</response>" in text:
+            score += 0.2
+        if action == "NONE":
+            score -= 0.5
+            
+        rewards.append(score)
+        
+    return rewards
+
+# -------------------------------------------------------------------------
+# Meta-Optimizer
+# -------------------------------------------------------------------------
+META_PROMPT_TEMPLATE = """You are an expert Prompt Engineer optimizing a system prompt for an AI agent.
+
+Current System Prompt:
+"{current_prompt}"
+
+Task:
+The AI needs to decide whether to REPLY, IGNORE, or STOP based on a chat history.
+It MUST output valid XML: <response><name>Eliza</name><reason>...</reason><action>...</action></response>.
+
+Here are some examples of the AI's performance with the current prompt:
+
+[SUCCESSFUL EXAMPLE]
+Input: {good_input}
+Output: {good_output}
+Reward: {good_score} (This was good!)
+
+[FAILED EXAMPLE]
+Input: {bad_input}
+Output: {bad_output}
+Reward: {bad_score} (This was bad!)
+
+[INSTRUCTIONS]
+Analyze why the successful example worked and the failed one didn't.
+Rewrite the System Prompt to fix the failure case while maintaining the success case.
+The new prompt should be concise, clear, and emphasize the XML format and the decision logic.
+Return ONLY the new system prompt text. Do not add explanations.
+"""
+
+def call_external_model(provider, model_name, prompt):
+    """
+    Call external API (Groq or Anthropic) to optimize the prompt.
+    """
+    try:
+        if provider == "groq":
+            from groq import Groq
+            client = Groq(api_key=os.environ.get("GROQ_API_KEY"))
+            chat_completion = client.chat.completions.create(
+                messages=[
+                    {"role": "system", "content": "You are an expert prompt engineer."},
+                    {"role": "user", "content": prompt}
+                ],
+                model=model_name,
+                temperature=0.7,
+                max_tokens=1024,
+            )
+            return chat_completion.choices[0].message.content
+            
+        elif provider == "anthropic":
+            import anthropic
+            client = anthropic.Anthropic(api_key=os.environ.get("ANTHROPIC_API_KEY"))
+            message = client.messages.create(
+                model=model_name,
+                max_tokens=1024,
+                messages=[
+                    {"role": "user", "content": prompt}
+                ]
+            )
+            return message.content[0].text
+            
+        elif provider == "openai":
+            from openai import OpenAI
+            client = OpenAI(api_key=os.environ.get("OPENAI_API_KEY"))
+            completion = client.chat.completions.create(
+                model=model_name,
+                messages=[
+                    {"role": "system", "content": "You are an expert prompt engineer."},
+                    {"role": "user", "content": prompt}
+                ]
+            )
+            return completion.choices[0].message.content
+            
+        else:
+            raise ValueError(f"Unknown provider: {provider}")
+            
+    except Exception as e:
+        print(f"Error calling {provider}: {e}")
+        return None
+
+def optimize_prompt(provider, model_name, current_prompt, good_example, bad_example):
+    meta_prompt = META_PROMPT_TEMPLATE.format(
+        current_prompt=current_prompt,
+        good_input=good_example['prompt'],
+        good_output=good_example['output'],
+        good_score=good_example['score'],
+        bad_input=bad_example['prompt'],
+        bad_output=bad_example['output'],
+        bad_score=bad_example['score']
+    )
+    
+    new_prompt = call_external_model(provider, model_name, meta_prompt)
+    
+    if new_prompt:
+        # Simple cleanup: remove quotes if the model added them
+        new_prompt = new_prompt.strip()
+        if new_prompt.startswith('"') and new_prompt.endswith('"'):
+            new_prompt = new_prompt[1:-1]
+        return new_prompt
+    else:
+        return current_prompt
+
+# -------------------------------------------------------------------------
+# Main Loop
+# -------------------------------------------------------------------------
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model", type=str, default="mlx-community/Qwen2.5-1.5B-Instruct-4bit")
+    parser.add_argument("--data", type=str, required=True)
+    parser.add_argument("--iter", type=int, default=10)
+    parser.add_argument("--batch-size", type=int, default=4)
+    parser.add_argument("--save-path", type=str, default="optimized_prompt.txt")
+    parser.add_argument("--optimizer-provider", type=str, default="groq", choices=["groq", "anthropic", "openai"])
+    parser.add_argument("--optimizer-model", type=str, default="llama-3.3-70b-versatile")
+    args = parser.parse_args()
+    
+    print(f"Loading Evaluator Model: {args.model}...")
+    model, tokenizer = load(args.model)
+    
+    print(f"Using Meta-Optimizer: {args.optimizer_provider} ({args.optimizer_model})")
+    
+    # Load Data
+    prompts = []
+    with open(args.data, 'r') as f:
+        for line in f:
+            if line.strip():
+                item = json.loads(line)
+                user_msg = next((m["content"] for m in reversed(item["messages"]) if m["role"] == "user"), None)
+                if user_msg: prompts.append(user_msg)
+                
+    random.shuffle(prompts)
+    
+    # Initialize Prompt
+    # Initialize Prompt
+    current_system_prompt = (
+        "You are an AI assistant named Eliza. "
+        "Decide whether to REPLY, IGNORE, or STOP. "
+        "You MUST respond in valid XML format: "
+        "<response><name>Eliza</name><reason>...</reason><action>REPLY/IGNORE/STOP</action></response>."
+    )
+    best_overall_avg_score = -float('inf')
+    
+    print(f"Starting Training-Free GRPO for {args.iter} iterations...")
+    
+    for i in range(args.iter):
+        print(f"\n=== Iteration {i+1}/{args.iter} ===")
+        
+        # 1. Sample Batch
+        batch_prompts = random.sample(prompts, args.batch_size)
+        
+        batch_results = []
+        
+        # 2. Evaluate Current Prompt on Batch
+        print(f"Evaluating current prompt on {len(batch_prompts)} examples...")
+        current_scores = []
+        
+        for p in batch_prompts:
+            # Construct full prompt with system instruction
+            messages = [
+                {"role": "system", "content": current_system_prompt},
+                {"role": "user", "content": p}
+            ]
+            full_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+            output = generate(model, tokenizer, prompt=full_text, max_tokens=100, verbose=False)
+            
+            # Score
+            score = compute_rewards([p], [output])[0]
+            current_scores.append(score)
+            
+            batch_results.append({
+                "prompt": p,
+                "output": output,
+                "score": score
+            })
+            
+        avg_score = sum(current_scores) / len(current_scores)
+        print(f"Avg Score: {avg_score:.2f}")
+        
+        if avg_score > best_overall_avg_score:
+            best_overall_avg_score = avg_score
+            print(f"New Best Score! Saving prompt to {args.save_path}...")
+            with open(args.save_path, "w") as f:
+                f.write(current_system_prompt)
+        
+        # 3. Select Good and Bad Examples
+        # Sort by score
+        batch_results.sort(key=lambda x: x['score'], reverse=True)
+        good = batch_results[0]
+        bad = batch_results[-1]
+        
+        if good['score'] > bad['score']:
+            # Variance exists, we can learn
+            print(f"Optimizing: Good ({good['score']}) vs Bad ({bad['score']})")
+            print(f"Good Output: {good['output'][:50]}...")
+            print(f"Bad Output: {bad['output'][:50]}...")
+            
+            new_prompt = optimize_prompt(args.optimizer_provider, args.optimizer_model, current_system_prompt, good, bad)
+            print(f"\nNew Prompt Proposed:\n{new_prompt}\n")
+            current_system_prompt = new_prompt
+        else:
+            print("No variance in batch (all good or all bad). Skipping optimization step.")
+            
+    print("\noptimization Complete.")
+    print(f"Final Prompt saved to {args.save_path}")
+
+if __name__ == "__main__":
+    main()

package/python/scripts/test_generation.py

@@ -0,0 +1,29 @@
+
+import mlx.core as mx
+from mlx_lm import load, generate
+import argparse
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model", type=str, default="mlx-community/Qwen2.5-1.5B-Instruct-4bit")
+    parser.add_argument("--adapter-path", type=str, default="trained_models/should_respond_sft/adapters")
+    parser.add_argument("--temp", type=float, default=1.0)
+    args = parser.parse_args()
+
+    print(f"Loading {args.model} with {args.adapter_path}")
+    model, tokenizer = load(args.model, adapter_path=args.adapter_path)
+
+    prompt = "<task>Decide on behalf of Eliza whether they should respond to the message, ignore it or stop the conversation.</task>\n\n<providers>\n[RECENT_MESSAGES]\nUser: I heard Eliza is helping\n</providers>\n\n<instructions>Decide if Eliza should respond to or interact with the conversation.\n\nIMPORTANT RULES FOR RESPONDING:\n- If YOUR name (Eliza) is directly mentioned → RESPOND\n- If someone uses a DIFFERENT name (not Eliza) → IGNORE (they're talking to someone else)\n- If you're actively participating in a conversation and the message continues that thread → RESPOND\n- If someone tells you to stop or be quiet → STOP\n- Otherwise → IGNORE\n\nThe key distinction is:\n- \"Talking TO Eliza\" (your name mentioned, replies to you, continuing your conversation) → RESPOND\n- \"Talking ABOUT Eliza\" or to someone else → IGNORE\n</instructions>\n\n<output>\nDo NOT include any thinking, reasoning, or <think> sections in your response.\nGo directly to the XML response format without any preamble or explanation.\n\nRespond using XML format like this:\n<response>\n  <name>Eliza</name>\n  <reasoning>Your reasoning here</reasoning>\n  <action>RESPOND | IGNORE | STOP</action>\n</response>\n\nIMPORTANT: Your response must ONLY contain the <response></response> XML block above. Do not include any text, thinking, or reasoning before or after this XML block. Start your response immediately with <response> and end with </response>.\n</output>"
+
+    print("\n--- Gen 1 (temp={}) ---".format(args.temp))
+    from mlx_lm.sample_utils import make_sampler
+    sampler = make_sampler(temp=args.temp)
+    
+    print(generate(model, tokenizer, prompt=prompt, max_tokens=50, verbose=True, sampler=sampler))
+
+    print("\n--- Gen 2 (temp={}) ---".format(args.temp))
+    sampler2 = make_sampler(temp=args.temp)
+    print(generate(model, tokenizer, prompt=prompt, max_tokens=50, verbose=True, sampler=sampler2))
+
+if __name__ == "__main__":
+    main()