jmparejaz/DeepResearch-RL-Trainer

DeepResearch-RL (tearsheet-agent-service)

Containerized pi (@earendil-works/pi-coding-agent@1.0.4) deep-research agent service for generating multi-section company tearsheets with grounded [N] citations, deterministic section/fact completeness auditing, tool-execution policy guards, and OpenTelemetry GenAI trajectory + reasoning-token export to W&B Weave (forge-ai-dev/Tearsheet-Pi-Agent).

Built under epic FND-516 as the harnessed agent foundation for Reinforcement Learning (RL / GRPO with Prime Intellect): frontier trajectories (xai/grok-4.7) and deterministic completeness audits (section_parser) serve as reference rollouts, reward signals, and evaluation benchmarks for training compact open-weight reasoning models (such as LiquidAI/LFM2.5-8B-A1B served via vLLM).

---

Architecture & Layout

• agent_runtime/ — PI agent runtime configuration (baked into /app/agent_runtime_template in the container and rendered at startup by docker/entrypoint.sh):
  • AGENTS.md — 10-section tearsheet output contract, inline [N] citation rules, and structured ```facts / ```retractions YAML block specification.
  • required_facts.yaml — Per-section required-fact inventory with stable identities and intentional blank allow-lists (e.g., target-company valuation/multiple cells in the comparables table).
  • skills/section-completeness/SKILL.md — Meta-harness skill embedding the canonical fact inventory and anti-gaming rules.
  • extensions/websearch_fetch.ts — PI extension registering custom websearch (Perplexity /search API) and fetch tools backed by Python clients and SourceRegistry.
  • extensions/tool_call_guard.ts — Pre-tool execution hook enforcing ToolPolicy ({read, websearch, fetch} allow-list and protected config/registry paths).
  • models.json — Multi-provider configuration supporting both:
    • perplexity-agent (api: "openai-responses", default model xai/grok-4.7, reasoning: true)
    • hf-vllm (api: "openai-completions", default model LiquidAI/LFM2.5-8B-A1B, contextWindow: 128000, reasoning: true, compatible with any vLLM / OpenAI-compatible server).
• src/tearsheet_agent_service/
  • cli.py — Terminal CLI entrypoint (tearsheet-research).
  • pi_rpc_client.py — Asyncio JSONL RPC client for pi --mode rpc.
  • run_manager.py — run_research_stream() async generator and run_research() wrapper.
  • source_registry.py — Persistent citation ledger ([N] ↔ URL, title, retrieved excerpt, status).
  • agent_runtime_support/ — section_parser.py (deterministic 10-section & fact-identity auditor) and tool_policy.py (tool/path authorization).
  • api/ — FastAPI HTTP/SSE service (app.py, run_manager.py, schemas.py, auth.py, trace_bridge.py).
  • web_tools/ — SearchClient, FetchClient, and subprocess CLI bridge.
  • rl/ — Prime Intellect verifiers-compatible Stage 2 RL / GRPO training & evaluation suite:
    • rubric.py — 33-failure-mode (EA03-*) TearsheetRubric (15 Tier-A deterministic Python verifiers + 18 Tier-B closed-world NLI Judge specs), Multiplicative Anti-Hacking Gate (MiMo-V2.6 Eq. 2), Groupwise Advantage Redistribution (GAR, Eq. 3), and Group-Relative Length Penalty (Eq. 4).
    • census_meta_eval.py — Offline meta-verification runner evaluating the 33-mode rubric against the 555-claim historical census (146 not_supported + 409 supported controls).
    • env.py — Stateful multi-turn TearsheetToolEnv with isolated per-rollout SourceRegistry, HybridReplayCache, and full_tearsheet prompt contract enforcement (Decision D3).
    • sampler.py — DynamicDiscrepancySampler ($\max R - \min R \ge 0.05$, Decision D1) and multi-provider JudgeClient (perplexity xai/grok-4.7 primary + mistral mistral-large-latest fallback, Decision D2).
    • weave_dataset.py — W&B Weave teacher trajectory harvester (forge-ai-dev/Tearsheet-Pi-Agent) and Teacher-Prefix On-Policy Distillation (Teacher-Prefix OPD, MiMo-V2.6 §5.6) dataset builder.
    • train_sft_warmup.py — Stage 1 100-step SFT / Prefix-OPD warm-up trainer with strict assistant-only loss masking (labels = -100 on system/user/tool turns) and frozen MoE routing gates.
    • train_grpo.py — Stage 2 Split-GPU A100x2 MoE GRPO trainer with freeze_moe_router (MiMo-V2.6 §5.4 Fig. 11) and prompt_mean token loss aggregation (MiMo-V2.6 §5.2).
• spaces/DeepResearch-RL-Trainer/ — Ready-to-deploy Hugging Face Space bundle (Dockerfile, rl.toml, start_a100x2_grpo.sh, README.md) targeting a100x2 (2× NVIDIA A100 80GB = 160GB VRAM) with GPU 0 dedicated to vLLM rollout generation and GPU 1 dedicated to SFT/GRPO training.
• scripts/
  • build_and_push_ecr.py — Automated cloud Docker build & push to AWS ECR via AWS CodeBuild (no local Docker daemon required).
  • deploy_hf_lfm25_endpoint.sh — Deploys LiquidAI/LFM2.5-8B-A1B on Hugging Face Inference Endpoints using vllm/vllm-openai:v0.29.0 with FP8 quantization (--quantization fp8 --kv-cache-dtype fp8), --reasoning-parser qwen3, and --tool-call-parser lfm2 on nvidia-l40s (48 GB VRAM).

---

Docker & AWS ECR Usage

1. Pre-built Image in AWS ECR

• ECR Repository URI: 526489132647.dkr.ecr.us-west-2.amazonaws.com/tearsheet-agent-service-poc:latest
• Current Image Digest: sha256:d71f274bd32ca71615828b7754febea01da00ddaf69c931c2cfd22521e01a659
• Region / Account: us-west-2 / 526489132647 (sandbox profile)

2. Rebuilding & Pushing to ECR (via AWS CodeBuild)

When local Docker is unavailable, build and push directly in AWS using:

aws sso login --profile sandbox
python3 scripts/build_and_push_ecr.py --cleanup-builder

(Packages the repo, runs buildspec.yml inside CodeBuild including container smoke tests, pushes :latest to ECR, and cleans up the temporary CodeBuild project/IAM role/S3 zip).

If building on a machine with Docker installed:

docker build -t tearsheet-agent-service:latest .

3. Running the Container as an HTTP/SSE Microservice (Default CMD)

By default, the container starts uvicorn tearsheet_agent_service.api.app:create_default_app --factory --host 0.0.0.0 --port 8000:

# Authenticate to ECR (if pulling from AWS)
aws ecr get-login-password --region us-west-2 --profile sandbox \
  | docker login --username AWS --password-stdin 526489132647.dkr.ecr.us-west-2.amazonaws.com

docker run --rm -p 8000:8000 \
  -e PERPLEXITY_API_KEY="$PERPLEXITY_API_KEY" \
  -e WANDB_API_KEY="$WANDB_API_KEY" \
  -e TEARSHEET_SERVICE_API_KEY="dev-local-api-key" \
  -e TEARSHEET_AGENT_PROVIDER="perplexity-agent" \
  -e TEARSHEET_AGENT_MODEL_ID="xai/grok-4.7" \
  -e HF_TOKEN="${HF_TOKEN:-}" \
  -e HF_INFERENCE_ENDPOINT_URL="https://6ad3b6e6d2d0ab8296e5999d.endpoints.huggingface.cloud/v1" \
  -e OTEL_EXPORTER_OTLP_ENDPOINT="https://trace.wandb.ai/agents/otel" \
  -e OTEL_EXPORTER_OTLP_PROTOCOL="http/protobuf" \
  -e OTEL_RESOURCE_ATTRIBUTES="wandb.entity=forge-ai-dev,wandb.project=Tearsheet-Pi-Agent" \
  -e OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT="true" \
  526489132647.dkr.ecr.us-west-2.amazonaws.com/tearsheet-agent-service-poc:latest

4. Running the Container in One-Shot CLI Mode

Override the container command with tearsheet-research (or pi --list-models):

# Verify configured providers & models inside the container
docker run --rm \
  -e PERPLEXITY_API_KEY="$PERPLEXITY_API_KEY" \
  -e HF_TOKEN="$HF_TOKEN" \
  526489132647.dkr.ecr.us-west-2.amazonaws.com/tearsheet-agent-service-poc:latest \
  pi --list-models

# Run a tearsheet research job via CLI inside the container
docker run --rm \
  -e PERPLEXITY_API_KEY="$PERPLEXITY_API_KEY" \
  -e WANDB_API_KEY="$WANDB_API_KEY" \
  -e HF_TOKEN="${HF_TOKEN:-}" \
  526489132647.dkr.ecr.us-west-2.amazonaws.com/tearsheet-agent-service-poc:latest \
  tearsheet-research --company "Ayar Labs"

5. Container Environment Variables Reference

Variable	Default	Description
PERPLEXITY_API_KEY	(required)	API key used for Perplexity /search (websearch tool) and perplexity-agent provider.
WANDB_API_KEY	(optional)	W&B API key for exporting OTel GenAI trajectories to Weave.
TEARSHEET_SERVICE_API_KEY	dev-local-api-key	Expected X-API-Key header for all FastAPI endpoints (/runs).
TEARSHEET_AGENT_PROVIDER	perplexity-agent	Default PI provider (perplexity-agent or hf-vllm).
TEARSHEET_AGENT_MODEL_ID	xai/grok-4.7	Default model ID rendered into models.json at container startup.
HF_TOKEN	(optional)	Hugging Face token for authenticating against hf-vllm endpoints.
HF_INFERENCE_ENDPOINT_URL	https://6ad3b6e6d2d0ab8296e5999d.endpoints.huggingface.cloud/v1	OpenAI-compatible /v1 base URL for hf-vllm (HF Inference Endpoint or any vLLM server).
TEARSHEET_HF_MODEL_ID	LiquidAI/LFM2.5-8B-A1B	Additional model ID dynamically registered under hf-vllm in models.json.
TEARSHEET_WANDB_ENTITY	forge-ai-dev	Target W&B entity for Weave traces.
TEARSHEET_WANDB_PROJECT	Tearsheet-Pi-Agent	Target W&B Weave project for agent trajectories.

---

Local Usage (CLI & API)

Install Dependencies

uv sync

Mode A: Terminal CLI (tearsheet-research)

Only --company is required. Runs PI with live turn/tool/reasoning telemetry, prints the 10-section section_parser audit table, and exports the trajectory to forge-ai-dev/Tearsheet-Pi-Agent:

# 1. Frontier baseline (Perplexity Grok 4.7 with reasoning token tracking)
uv run tearsheet-research --company "Ayar Labs"

# 2. Quick single-section demo
uv run tearsheet-research \
  --company "FloQast" \
  --part "ONLY the single section ## Business Overview"

# 3. Open-weight model on HF Inference Endpoint / vLLM (LiquidAI/LFM2.5-8B-A1B FP8)
uv run tearsheet-research \
  --provider hf-vllm \
  --model "LiquidAI/LFM2.5-8B-A1B" \
  --hf-endpoint-url "https://6ad3b6e6d2d0ab8296e5999d.endpoints.huggingface.cloud/v1" \
  --company "Ayar Labs"

Mode B: Local FastAPI Server (runs-api)

uv run uvicorn tearsheet_agent_service.api.app:create_default_app --factory --host 127.0.0.1 --port 8090

Trigger a run (only "company" is required; provider, model, part, website, start_year, issuer_slug are optional):

# Start a run
RUN_ID=$(curl -s -X POST http://127.0.0.1:8090/runs \
  -H "X-API-Key: dev-local-api-key" \
  -H "Content-Type: application/json" \
  -d '{"company": "Ayar Labs", "provider": "perplexity-agent", "model": "xai/grok-4.7"}' \
  | jq -r .run_id) && echo "RUN_ID=$RUN_ID"

# Stream live SSE events
curl -N -H "X-API-Key: dev-local-api-key" "http://127.0.0.1:8090/runs/$RUN_ID/events"

# Fetch status, resolved citations, reasoning tokens, and W&B Weave trace URL
curl -s -H "X-API-Key: dev-local-api-key" "http://127.0.0.1:8090/runs/$RUN_ID" \
  | jq '{status, cost, trace_url, sources_count: (.result.sources|length)}'

---

Managing the Hugging Face vLLM Endpoint (LiquidAI/LFM2.5-8B-A1B)

The endpoint lfm2-5-8b-a1b-xcj runs vllm/vllm-openai:v0.29.0 on 1× NVIDIA L40S (48 GB VRAM) with: -tp 1 --reasoning-parser qwen3 --enable-auto-tool-choice --tool-call-parser lfm2 --quantization fp8 --kv-cache-dtype fp8 --max-model-len 128000

# Resume when testing
hf endpoints resume lfm2-5-8b-a1b-xcj

# Check status & URL
hf endpoints describe lfm2-5-8b-a1b-xcj

# Pause when done ($0/hr)
hf endpoints pause lfm2-5-8b-a1b-xcj

# Or deploy a fresh endpoint from scratch
./scripts/deploy_hf_lfm25_endpoint.sh

---

Stage 2 RL / GRPO Training Suite (tearsheet_agent_service.rl & spaces/DeepResearch-RL-Trainer/)

1. 33-Failure-Mode (EA03-*) Rubric & 555-Claim Census Meta-Evaluator

• Exact 33-Mode Bijection (rl/rubric.py): Partitions the 33 EA03-* failure modes from ea03-taxonomy-v1/taxonomy.json into two disjoint tiers:
  • 15 Tier-A Deterministic Python Verifiers (TIER_A_FAILURE_MODES, <5ms/rollout on CPU): EA03-SUBJECT, EA03-PROFILE, EA03-TOPIC, EA03-LIST, EA03-MARKET, EA03-VALUE, EA03-DERIVE, EA03-TABLE, EA03-MISS, EA03-MIX, EA03-CITE, EA03-ACCESS, EA03-ASSOC, EA03-FETCH, and EA03-CROSS.
  • 18 Tier-B Closed-World NLI Judge Specs (TIER_B_FAILURE_MODES): EA03-REL, EA03-QTY, EA03-NUM, EA03-SCOPE, EA03-ROUND, EA03-INSTRUMENT, EA03-EVENT, EA03-INCUMBENT, EA03-ADJROLE, EA03-HOMONYM, EA03-BRIDGE, EA03-RANK, EA03-DIRECTION, EA03-FLOW, EA03-VOICE, EA03-PRODUCT, EA03-PRODUCT_TYPE, and EA03-DELTA.
• Offline 555-Claim Census Meta-Evaluation (rl/census_meta_eval.py): Replays all 555 parent claims (146 not_supported in population.jsonl + 409 supported controls in controls_index.jsonl, 207 mode memberships, and 10,149 evidence rows) to verify per-mode recall and control false-rejection rates:

  uv run python -m tearsheet_agent_service.rl.census_meta_eval
  

2. Xiaomi MiMo-V2.6 RL Stabilization Mechanisms & Locked Decisions (D1–D4)

• Multiplicative Anti-Reward-Hacking Gate (MiMo-V2.6 Eq. 2): Composite episode reward $R_i = R^{\text{gate}}_i \cdot S^{\text{sol}}_i \cdot S^{\text{beh}}_i$ caps rollouts with zero websearch calls or prompt-template contamination (SampleCorp, Jane Doe, Mina Patel, example.com) at $\le 0.05$, and penalizes malformed ```markdown\nfacts fences or premature 6/10-section Part-1 stops below $0.35$.
• Dynamic Discrepancy Filtering (MiMo-V2.6 §5.3, Decision D1): DynamicDiscrepancySampler (rl/sampler.py) requires $\max_i R^{\text{Tier-A}}_i - \min_i R^{\text{Tier-A}}_i \ge 0.05$ across a rollout group ($G=8$) before invoking JudgeClient or emitting the group to the GPU 1 Trainer—saving 100% of external Judge API calls on all-fail and all-pass groups.
• Closed-World NLI Judge & Groupwise Advantage Redistribution (MiMo-V2.6 Eq. 3, Decision D2): JudgeClient defaults to provider="perplexity", model="xai/grok-4.7" with secondary fallback provider="mistral", model="mistral-large-latest" (and configurable deepseek, openai, glm support), followed by zero-mean, unit-variance compute_gar_advantages() and apply_group_relative_length_penalty() (MiMo-V2.6 Eq. 4).
• Full-Tearsheet Curriculum Only (Decision D3): All environment prompts (rl/env.py), W&B Weave trajectory harvesters (rl/weave_dataset.py), SFT warm-up configs (rl/train_sft_warmup.py), and GRPO configs (rl/train_grpo.py) enforce curriculum_mode = "full_tearsheet" (all 10 sections from required_facts.yaml).
• MoE Router Freezing (MiMo-V2.6 §5.4 Fig. 11) & Prompt-Mean Loss (MiMo-V2.6 §5.2): freeze_moe_router(model) freezes all MoE routing gate parameters (*.mlp.gate.*, *.block_sparse_moe.gate.*, *.moe.gate.*, *.router.*, *.expert_gate.*) on LiquidAI/LFM2.5-8B-A1B (Lfm2MoeForCausalLM) while preserving SwiGLU gate_proj/up_proj/down_proj/w1..w3 and attention projections as trainable under LoRA, and compute_prompt_mean_grpo_loss() normalizes per-rollout token losses by each rollout's active assistant length before group averaging.
• Teacher-Prefix On-Policy Distillation (MiMo-V2.6 §5.6): weave_dataset.py harvests 10/10-complete xai/grok-4.7 teacher traces from W&B Weave (forge-ai-dev/Tearsheet-Pi-Agent), extracts reference_facts and replay_snapshots, and slices intermediate tool-observation prefix states via build_teacher_prefix_opd_examples().

3. Split-GPU A100x2 (160GB VRAM) Hugging Face Space (spaces/DeepResearch-RL-Trainer/)

The spaces/DeepResearch-RL-Trainer/ bundle packages the complete Split-GPU training stack (GPU 0 co-located vLLM rollout server + GPU 1 Stage 1 SFT Warm-Up & Stage 2 GRPO Trainer) while keeping the standalone HF Inference Endpoint lfm2-5-8b-a1b-xcj paused at $0/hr for live boss demos (Decision D4):

# Validate shell syntax and run the offline Split-GPU dry-run check
bash -n spaces/DeepResearch-RL-Trainer/start_a100x2_grpo.sh
./spaces/DeepResearch-RL-Trainer/start_a100x2_grpo.sh --dry-run

# Launch full Stage 1 (100-step SFT warm-up) + Stage 2 (GRPO) on an A100x2 Space
./spaces/DeepResearch-RL-Trainer/start_a100x2_grpo.sh

---

Running Tests & Linter

uv run pytest tests/ -q        # Full unit + RL + census meta-eval test suite
uv run ruff check .            # Ruff lint check