Trajectory-Based Self-Improvement

What Is a Trajectory?

Recording a Trajectory

from dataclasses import dataclass, field
from datetime import datetime
from typing import Any

@dataclass
class TrajectoryStep:
    step_index: int
    state_summary: str    # short description of world state
    action_name: str      # tool or reasoning step name
    action_params: dict
    result: Any
    timestamp: str = ''

    def __post_init__(self):
        if not self.timestamp:
            self.timestamp = datetime.utcnow().isoformat()

@dataclass
class Trajectory:
    trajectory_id: str
    task: str
    steps: list = field(default_factory=list)
    outcome: str = 'unknown'   # 'success', 'failure', 'partial'
    final_score: float = 0.0

    def add_step(self, step: TrajectoryStep):
        self.steps.append(step)

    def mark_success(self, score: float = 1.0):
        self.outcome = 'success'
        self.final_score = score

    def mark_failure(self, reason: str = ''):
        self.outcome = 'failure'
        self.final_score = 0.0

if __name__ == '__main__':
    traj = Trajectory(trajectory_id='t-1', task='Book a flight to Tokyo')
    traj.add_step(TrajectoryStep(
        step_index=0, state_summary='searching flights',
        action_name='search_flights', action_params={'dest': 'NRT'}, result='5 options found'
    ))
    traj.mark_success(score=0.95)
    print(f'Trajectory {traj.trajectory_id}: outcome={traj.outcome}, score={traj.final_score}')
    print('Steps recorded:', len(traj.steps))

Storing Trajectories

import json
import os
from dataclasses import asdict

TRAJECTORY_DIR = 'trajectories'

def save_trajectory(traj: Trajectory):
    os.makedirs(TRAJECTORY_DIR, exist_ok=True)
    filename = f'{TRAJECTORY_DIR}/{traj.trajectory_id}_{traj.outcome}.json'
    data = asdict(traj)
    with open(filename, 'w') as f:
        json.dump(data, f, indent=2)
    print(f'Saved trajectory: {filename}')

def load_successful_trajectories(task_type: str, n: int = 5) -> list:
    results = []
    for fname in os.listdir(TRAJECTORY_DIR):
        if '_success.json' not in fname:
            continue
        with open(os.path.join(TRAJECTORY_DIR, fname)) as f:
            traj = json.load(f)
        if task_type.lower() in traj['task'].lower():
            results.append(traj)
    results.sort(key=lambda t: t['final_score'], reverse=True)
    return results[:n]

Using Successful Trajectories as Few-Shot Examples

def trajectory_to_few_shot(traj: dict) -> str:
    lines = [f'Example task: {traj["task"]}', 'Steps taken:']
    for step in traj['steps']:
        lines.append(
            f'  [{step["step_index"]}] {step["action_name"]}'
            f'({json.dumps(step["action_params"])}) -> {str(step["result"])[:80]}'
        )
    lines.append(f'Outcome: {traj["outcome"]} (score={traj["final_score"]:.2f})')
    return '\n'.join(lines)

def build_few_shot_system_prompt(task_type: str) -> str:
    successful = load_successful_trajectories(task_type, n=2)
    if not successful:
        return 'Complete the following task step by step.'
    examples = '\n\n---\n\n'.join(
        trajectory_to_few_shot(t) for t in successful
    )
    return (
        'Here are examples of successfully completed similar tasks:\n\n'
        + examples
        + '\n\n---\n\nNow complete the new task using the same approach.'
    )

Analysing Failed Trajectories

def analyze_failure(traj: dict, client) -> dict:
    steps_str = json.dumps(traj['steps'], indent=2)
    prompt = (
        f'Task: {traj["task"]}\n\n'
        f'Agent trajectory (failed):\n{steps_str}\n\n'
        'Identify the failure mode. Return JSON:\n'
        '{"failure_step": 0, "failure_mode": "", "root_cause": "", '
        '"prevention": ""}'
    )
    import anthropic
    client_obj = anthropic.Anthropic(api_key='YOUR_API_KEY')
    result = client_obj.messages.create(
        model='claude-opus-4-5',
        max_tokens=512,
        messages=[{'role': 'user', 'content': prompt}]
    )
    import json
    return json.loads(result.content[0].text)

Failure Mode Taxonomy

from collections import Counter

def build_failure_taxonomy(failed_trajectories: list, client) -> dict:
    failure_modes = []
    for traj in failed_trajectories:
        analysis = analyze_failure(traj, client)
        failure_modes.append(analysis['failure_mode'])

    counts = Counter(failure_modes)
    total = len(failure_modes)

    taxonomy = [
        {
            'failure_mode': mode,
            'count': count,
            'percentage': round(count / total * 100, 1)
        }
        for mode, count in counts.most_common()
    ]

    print('Failure Mode Taxonomy:')
    for entry in taxonomy:
        print(f'  {entry["failure_mode"]}: {entry["count"]} ({entry["percentage"]}%)')

    return {'taxonomy': taxonomy, 'total_failures': total}

Creating Training Pairs from Trajectories

def trajectory_to_training_pair(
    failed_traj: dict,
    failure_analysis: dict
) -> dict:
    """
    Creates an SFT-ready training pair:
    input = task + context at failure step
    output = what the agent should have done
    """
    fail_step_idx = failure_analysis['failure_step']
    steps = failed_traj['steps']

    # Context up to (but not including) the failure step
    context_steps = steps[:fail_step_idx]
    context_str = '\n'.join(
        f'Step {s["step_index"]}: {s["action_name"]}({s["action_params"]})'
        for s in context_steps
    )
    return {
        'input': f'Task: {failed_traj["task"]}\n\nPrevious steps:\n{context_str}\n\nNext action:',
        'output': failure_analysis['prevention'],  # correct action
        'source': 'failure_trajectory',
        'trajectory_id': failed_traj.get('trajectory_id', 'unknown')
    }

if __name__ == '__main__':
    failed_traj = {
        'task': 'Cancel subscription',
        'trajectory_id': 'traj-7',
        'steps': [
            {'step_index': 0, 'action_name': 'find_account', 'action_params': {'user': 'u1'}},
            {'step_index': 1, 'action_name': 'delete_account', 'action_params': {'user': 'u1'}},
        ]
    }
    failure_analysis = {'failure_step': 1, 'prevention': 'call cancel_subscription(user="u1") instead'}
    pair = trajectory_to_training_pair(failed_traj, failure_analysis)
    print('Training input:')
    print(pair['input'])
    print('Expected output:', pair['output'])

Fine-Tuning from Trajectories

import json

def export_fine_tuning_jsonl(
    training_pairs: list,
    output_file: str,
    system_prompt: str = 'You are an efficient AI agent.'
):
    with open(output_file, 'w') as f:
        for pair in training_pairs:
            record = {
                'messages': [
                    {'role': 'system', 'content': system_prompt},
                    {'role': 'user', 'content': pair['input']},
                    {'role': 'assistant', 'content': pair['output']}
                ]
            }
            f.write(json.dumps(record) + '\n')
    print(f'Exported {len(training_pairs)} training pairs to {output_file}')

# Upload via OpenAI API (pseudocode):
# client.files.create(file=open('train.jsonl','rb'), purpose='fine-tune')
# client.fine_tuning.jobs.create(training_file='file-id', model='gpt-4o-mini')

if __name__ == '__main__':
    import tempfile, os
    pairs = [{'input': 'Task: Cancel subscription\n\nNext action:', 'output': 'cancel_subscription(user="u1")'}]
    out_path = os.path.join(tempfile.gettempdir(), 'demo_training.jsonl')
    export_fine_tuning_jsonl(pairs, out_path)

Quality Filtering Trajectories

def filter_high_quality_trajectories(
    trajectories: list,
    max_steps: int = 8,
    min_score: float = 0.85
) -> list:
    high_quality = []
    for traj in trajectories:
        if traj['outcome'] != 'success':
            continue
        if traj['final_score'] < min_score:
            continue
        if len(traj['steps']) > max_steps:
            continue
        high_quality.append(traj)

    # Sort by (score DESC, steps ASC)
    high_quality.sort(
        key=lambda t: (-t['final_score'], len(t['steps']))
    )
    print(f'High-quality trajectories: {len(high_quality)} / {len(trajectories)}')
    return high_quality

if __name__ == '__main__':
    trajectories = [
        {'outcome': 'success', 'final_score': 0.92, 'steps': [1, 2, 3]},
        {'outcome': 'failure', 'final_score': 0.10, 'steps': [1]},
        {'outcome': 'success', 'final_score': 0.60, 'steps': [1, 2]},
    ]
    filter_high_quality_trajectories(trajectories)

Trajectory Comparison for Insight

def compare_trajectories(success_traj: dict, failure_traj: dict) -> dict:
    s_steps = {s['step_index']: s for s in success_traj['steps']}
    f_steps = {s['step_index']: s for s in failure_traj['steps']}

    divergences = []
    for idx in sorted(set(s_steps) & set(f_steps)):
        s_action = s_steps[idx]['action_name']
        f_action = f_steps[idx]['action_name']
        if s_action != f_action:
            divergences.append({
                'step': idx,
                'success_action': s_action,
                'failure_action': f_action
            })
            break  # First divergence is most important

    return {
        'first_divergence': divergences[0] if divergences else None,
        'success_steps': len(s_steps),
        'failure_steps': len(f_steps)
    }

# Usage:
# comparison = compare_trajectories(good_traj, bad_traj)
# print('First divergence:', comparison['first_divergence'])

if __name__ == '__main__':
    good_traj = {'steps': [{'step_index': 0, 'action_name': 'search'}, {'step_index': 1, 'action_name': 'summarize'}]}
    bad_traj = {'steps': [{'step_index': 0, 'action_name': 'search'}, {'step_index': 1, 'action_name': 'delete'}]}
    comparison = compare_trajectories(good_traj, bad_traj)
    print('First divergence:', comparison['first_divergence'])

Trajectory-Based Improvement Flywheel

class TrajectorySelfImprovement:
    def __init__(self, agent_id: str):
        self.agent_id = agent_id
        self.trajectory_lib = []

    def run_and_record(self, task: str, agent_fn) -> Trajectory:
        traj = Trajectory(
            trajectory_id=f'{self.agent_id}_{len(self.trajectory_lib)}',
            task=task
        )
        result = agent_fn(task, traj)  # agent_fn appends steps to traj
        # Evaluate result (e.g., via reflection or user rating)
        score = evaluate_result(result)
        if score >= 0.7:
            traj.mark_success(score)
        else:
            traj.mark_failure('Low quality score')
        save_trajectory(traj)
        self.trajectory_lib.append(traj)
        return traj

    def improvement_cycle(self, client):
        failed = [t for t in self.trajectory_lib if t.outcome == 'failure']
        if len(failed) >= 10:
            taxonomy = build_failure_taxonomy([vars(t) for t in failed], client)
            print('Improvement cycle complete:', taxonomy)

def evaluate_result(result: str) -> float:
    return 0.8  # placeholder

Knowledge Check

Recap: Trajectory-Based Self-Improvement