diff --git a/README.md b/README.md index a760485..00f79c5 100644 --- a/README.md +++ b/README.md @@ -132,6 +132,12 @@ Run the IgnitionRAG evaluation bridge prototype: bun run --filter './examples/ignitionrag-evaluation-bridge' dev ``` +Run a deterministic RAG environment episode and trajectory report: + +```bash +bun run --filter './examples/rag-environment-episode' dev +``` + --- ## Minimal usage diff --git a/bun.lock b/bun.lock index 0022eab..ded9729 100644 --- a/bun.lock +++ b/bun.lock @@ -73,6 +73,14 @@ "@ignitionai/agent-trainer-preset-rag": "workspace:*", }, }, + "examples/rag-environment-episode": { + "name": "rag-environment-episode", + "version": "0.0.0", + "dependencies": { + "@ignitionai/agent-trainer-environment": "workspace:*", + "@ignitionai/agent-trainer-rl": "workspace:*", + }, + }, "packages/adapter-callable": { "name": "@ignitionai/agent-trainer-adapter-callable", "version": "0.1.0-alpha.1", @@ -644,6 +652,8 @@ "postcss-load-config": ["postcss-load-config@6.0.1", "", { "dependencies": { "lilconfig": "^3.1.1" }, "peerDependencies": { "jiti": ">=1.21.0", "postcss": ">=8.0.9", "tsx": "^4.8.1", "yaml": "^2.4.2" }, "optionalPeers": ["jiti", "postcss", "tsx", "yaml"] }, "sha512-oPtTM4oerL+UXmx+93ytZVN82RrlY/wPUV8IeDxFrzIjXOLF1pN+EmKPLbubvKHT2HC20xXsCAH2Z+CKV6Oz/g=="], + "rag-environment-episode": ["rag-environment-episode@workspace:examples/rag-environment-episode"], + "readdirp": ["readdirp@4.1.2", "", {}, "sha512-GDhwkLfywWL2s6vEjyhri+eXmfH6j1L7JE27WhqLeYzoh/A3DBaYGEj2H/HFZCn/kMfim73FXxEJTw06WtxQwg=="], "resolve-from": ["resolve-from@5.0.0", "", {}, "sha512-qYg9KP24dD5qka9J47d0aVky0N+b4fTU89LN9iDnjB5waksiC49rvMB0PrUJQGoTmH50XPiqOvAjDfaijGxYZw=="], diff --git a/docs/ALPHA_READINESS.md b/docs/ALPHA_READINESS.md index 71df207..df03060 100644 --- a/docs/ALPHA_READINESS.md +++ b/docs/ALPHA_READINESS.md @@ -29,7 +29,7 @@ All packages declare `license: MIT`, matching the root `LICENSE` file. | `@ignitionai/agent-trainer-adapter-vercel-ai` | ready | ready | ready | ready | partial | Structural adapter only; no streaming, tools or live provider calls. | | `@ignitionai/agent-trainer-cli` | ready | ready | ready | ready | partial | Runs typed experiments and writes standalone reports or timestamped bundles; no history/baseline/regression flags yet. | | `@ignitionai/agent-trainer-core` | ready | ready | missing | ready | partial | Foundational types and helpers need dedicated tests before alpha-stable status. | -| `@ignitionai/agent-trainer-environment` | ready | ready | missing | ready | prototype | Early environment loop only; no recorder, policy evaluation or tests yet. | +| `@ignitionai/agent-trainer-environment` | ready | ready | ready | ready | partial | Tested episode runner with safety guards and a deterministic RAG episode example; no production runtime or optimization loop. | | `@ignitionai/agent-trainer-evals` | ready | ready | ready | ready | partial | Current rewards are tested; RAG presets and richer scoring are still missing. | | `@ignitionai/agent-trainer-experiments` | ready | ready | ready | ready | ready | Local runner, definitions, gates and JSONL history are tested and documented. | | `@ignitionai/agent-trainer-exporters` | ready | ready | ready | ready | ready | JSON/Markdown report shape and local report bundles are tested. | @@ -55,10 +55,10 @@ All packages declare `license: MIT`, matching the root `LICENSE` file. ## Known Work After Internal Alpha - Add dedicated tests for `@ignitionai/agent-trainer-core`. -- Add dedicated tests for `@ignitionai/agent-trainer-environment` or keep it explicitly prototype-only. - Add CLI history/baseline/regression ergonomics after report bundles and CI examples. - Decide package publishing policy before any npm publication. - Add deeper examples for ecosystem adapters. +- Add a lightweight policy optimization loop after real dogfood produces trajectory data. ## Explicit Non-goals diff --git a/docs/BACKLOG.md b/docs/BACKLOG.md index a84c1a2..9031334 100644 --- a/docs/BACKLOG.md +++ b/docs/BACKLOG.md @@ -2596,7 +2596,7 @@ Next PR: Status: -- current +- completed Branch: @@ -2658,6 +2658,72 @@ Definition of done: - docs show manual alpha publish order and external smoke test, - no npm publish happens in this PR. +Next PR: + +- PR #43 - `feat: add tested environment episodes and trajectory rewards` + +### PR #43 - `feat: add tested environment episodes and trajectory rewards` + +Status: + +- current + +Branch: + +```txt +feat/tested-environment-episodes +``` + +Goal: + +Make `@ignitionai/agent-trainer-environment` usable as a tested RL foundation before adding policy optimization or PPO. + +Scope: + +- enrich `runEpisode` with options, final state, average reward, metadata and safety guards, +- add dedicated environment tests, +- convert environment episodes to RL trajectories without making `environment` depend on `rl`, +- export deterministic trajectory JSON/Markdown reports, +- prove compatibility with offline policy evaluation, +- add a deterministic RAG environment episode example. + +Out of scope: + +- PPO implementation, +- GRPO training, +- fine-tuning, +- real provider calls, +- production routing, +- CLI optimization loop, +- IgnitionRAG app integration. + +Required APIs / files: + +- `packages/environment`, +- `packages/rl`, +- `examples/rag-environment-episode`, +- readiness and audit docs. + +Acceptance: + +```bash +bun run lint +bun run typecheck +bun run test +bun run build +bun run pack:check +bun run --filter './examples/rag-environment-episode' dev +``` + +Definition of done: + +- `runEpisode(environment, policy)` remains compatible, +- environment episodes have dedicated tests, +- episode trajectories can become offline policy records, +- trajectory report output is deterministic, +- example documents `search -> rerank -> verify -> answer`, +- docs state that PPO and training remain out of scope. + ## Dogfood phase - IgnitionRAG Next work should happen in IgnitionRAG, not by adding more framework abstractions here. diff --git a/docs/MILESTONES.md b/docs/MILESTONES.md index e7bc83f..778a20a 100644 --- a/docs/MILESTONES.md +++ b/docs/MILESTONES.md @@ -296,13 +296,13 @@ docs/ALPHA_VALIDATION_PLAN.md Current status: ```txt -PR #36 through PR #41 are complete. The internal v0.1.0-alpha.0 tag is pushed on the merged bridge prototype. PR #42 is the current npm alpha publishing readiness branch. +PR #36 through PR #42 are complete. The internal v0.1.0-alpha.0 tag and npm v0.1.0-alpha.1 publication are complete. ``` Next phase: ```txt -Prepare npm alpha publishing, tag v0.1.0-alpha.1 after PR #42 merges, publish with dist-tag alpha, then dogfood the alpha inside IgnitionRAG. +Dogfood the alpha inside IgnitionRAG, then return to this repository for concrete framework gaps such as tested environment episodes, trajectory reporting and policy-evaluation ergonomics. ``` Explicit non-goals: diff --git a/docs/PROJECT_AUDIT.md b/docs/PROJECT_AUDIT.md index f339106..592e9ba 100644 --- a/docs/PROJECT_AUDIT.md +++ b/docs/PROJECT_AUDIT.md @@ -94,11 +94,11 @@ If a package exists but is intentionally narrow, minimal or untested, it is part ### `@ignitionai/agent-trainer-environment` - Purpose: early state/action/reward/policy environment loop primitives. -- Main exports: `runEpisode`, `AgentEnvironment`, `EnvironmentState`, `EnvironmentAction`, `Policy`, `EpisodeResult`. -- Stability level: prototype. -- Tests present: no. -- Example present: no. -- Known limitations: no tests, no rollout recorder, no production policy evaluation. +- Main exports: `runEpisode`, `RunEpisodeOptions`, `AgentEnvironment`, `EnvironmentState`, `EnvironmentAction`, `Policy`, `EpisodeStep`, `EpisodeResult`. +- Stability level: partial. +- Tests present: yes. +- Example present: yes, `examples/rag-environment-episode`. +- Known limitations: no production environment implementation, no durable rollout store, no automatic policy optimization loop and no CLI integration for environment episodes. ### `@ignitionai/agent-trainer-evals` @@ -148,10 +148,10 @@ If a package exists but is intentionally narrow, minimal or untested, it is part ### `@ignitionai/agent-trainer-rl` - Purpose: experimental RL-inspired utilities, deterministic policy selection helpers, fixed-strategy/contextual bandit prototypes, offline policy evaluation, GRPO-style candidate selection and PPO interface skeletons. -- Main exports: `Policy`, `PolicyContext`, `PolicyDecision`, `createStaticPolicy`, `createScoreBasedPolicy`, `Trajectory`, `TrajectoryStep`, `recordTrajectory`, `summarizeTrajectory`, `EpsilonGreedyBandit`, `RandomPolicy`, `ExperimentalBanditStrategySelector`, `ContextualBanditStrategySelector`, `ContextFeatures`, `scoreContextMatch`, `evaluatePolicyOffline`, `PolicyEvaluationResult`, `selectGroupRelativeBest`, `rankCandidateGroup`, `GroupRelativeSelectionResult`, `PPOConfig`, `PPOTrainer`, `PPOTrainingBatch`, `PPOTrainingResult`, `UnimplementedPPOTrainer`, `createUnimplementedPPOTrainer`. +- Main exports: `Policy`, `PolicyContext`, `PolicyDecision`, `createStaticPolicy`, `createScoreBasedPolicy`, `Trajectory`, `TrajectoryStep`, `recordTrajectory`, `summarizeTrajectory`, `recordEpisodeTrajectory`, `exportTrajectoryReport`, `toMarkdownTrajectoryReport`, `EpsilonGreedyBandit`, `RandomPolicy`, `ExperimentalBanditStrategySelector`, `ContextualBanditStrategySelector`, `ContextFeatures`, `scoreContextMatch`, `evaluatePolicyOffline`, `PolicyEvaluationResult`, `selectGroupRelativeBest`, `rankCandidateGroup`, `GroupRelativeSelectionResult`, `PPOConfig`, `PPOTrainer`, `PPOTrainingBatch`, `PPOTrainingResult`, `UnimplementedPPOTrainer`, `createUnimplementedPPOTrainer`. - Stability level: prototype. -- Tests present: yes for policy helpers, trajectories, fixed-strategy bandit, contextual bandit, offline policy evaluation, GRPO-style selection and PPO skeletons. -- Example present: no. +- Tests present: yes for policy helpers, trajectories, environment episode trajectory conversion, fixed-strategy bandit, contextual bandit, offline policy evaluation, GRPO-style selection and PPO skeletons. +- Example present: yes, `examples/rag-environment-episode`. - Known limitations: no PPO optimization, no GRPO model training, no production routing. ### `@ignitionai/agent-trainer` @@ -210,6 +210,13 @@ If a package exists but is intentionally narrow, minimal or untested, it is part - Mocked or live: deterministic mocked records and adapters. - Product concept: bridge prototype proving `Dataset` and `AgentVariant` mapping without IgnitionRAG app code. +### `examples/rag-environment-episode` + +- Demonstrates: modeling a RAG workflow as an environment episode with `search`, `rerank`, `verify` and `answer` actions. +- Command: `bun run --filter './examples/rag-environment-episode' dev`. +- Mocked or live: deterministic mocked environment. +- Product concept: episode rewards, trajectory reporting and offline policy record creation before any training loop. + ## Current Capabilities Matrix | Capability | Status | Package/File | Stable? | Notes | @@ -235,7 +242,9 @@ If a package exists but is intentionally narrow, minimal or untested, it is part | IgnitionRAG evaluation bridge prototype | prototype | `examples/ignitionrag-evaluation-bridge` | No | Deterministic record mapping only; no database, hosted worker, auth or real provider calls. | | file-based history | done | `packages/experiments/src/history.ts` | Yes | JSONL local history helpers, no CLI flag yet. | | policy abstraction | partial | `packages/rl/src/policy.ts` | No | Deterministic static and score-based selection only; no training loop. | +| environment episodes | partial | `packages/environment` | No | Tested episode runner with max-step safety, seed forwarding, final state and metadata. | | trajectory recorder | partial | `packages/rl/src/trajectory.ts` | No | Local state/action/reward/outcome records with deterministic summaries. | +| trajectory reports | partial | `packages/rl/src/episode-trajectory.ts` | No | Environment episodes can become trajectories, Markdown reports and offline policy records. | | bandit prototype | prototype | `packages/rl/src/strategy-bandit.ts` | No | Clearly experimental, fixed arms only, no PPO. | | contextual bandit prototype | prototype | `packages/rl/src/contextual-bandit.ts` | No | Deterministic fixed-feature scoring over task type, citation need, cost sensitivity, latency sensitivity and risk level. | | offline policy evaluation | prototype | `packages/rl/src/offline-policy-evaluation.ts` | No | Deterministic replay over local records or observed trajectory steps; no live traffic path. | @@ -250,7 +259,8 @@ If a package exists but is intentionally narrow, minimal or untested, it is part - No real LLM calls by default. - No hosted IgnitionRAG integration code. - IgnitionRAG integration is design-only. -- Core and environment need dedicated tests before alpha-stable status. +- Core needs dedicated tests before alpha-stable status. +- Environment episodes still need real dogfood data and CLI ergonomics before stable status. - Ecosystem adapters are minimal and structural. - CLI does not yet support history, baseline selection or fail-on-regression flags. - Bandit support is prototype-only. diff --git a/examples/rag-environment-episode/README.md b/examples/rag-environment-episode/README.md new file mode 100644 index 0000000..37c4654 --- /dev/null +++ b/examples/rag-environment-episode/README.md @@ -0,0 +1,19 @@ +# RAG Environment Episode + +This example models a deterministic RAG workflow as a lightweight environment episode. + +The scripted policy chooses: + +```txt +search -> rerank -> verify -> answer +``` + +Each action returns a reward. The completed episode is converted into a trajectory, summarized, exported as Markdown and converted into offline policy evaluation records. + +## Run + +```bash +bun run --filter './examples/rag-environment-episode' dev +``` + +The example is fully mocked. It does not call LLM providers, vector databases or external tools. diff --git a/examples/rag-environment-episode/package.json b/examples/rag-environment-episode/package.json new file mode 100644 index 0000000..2dfa307 --- /dev/null +++ b/examples/rag-environment-episode/package.json @@ -0,0 +1,13 @@ +{ + "name": "rag-environment-episode", + "version": "0.0.0", + "private": true, + "type": "module", + "scripts": { + "dev": "bun src/index.ts" + }, + "dependencies": { + "@ignitionai/agent-trainer-environment": "workspace:*", + "@ignitionai/agent-trainer-rl": "workspace:*" + } +} diff --git a/examples/rag-environment-episode/src/index.ts b/examples/rag-environment-episode/src/index.ts new file mode 100644 index 0000000..624581c --- /dev/null +++ b/examples/rag-environment-episode/src/index.ts @@ -0,0 +1,162 @@ +import { + type AgentEnvironment, + type EnvironmentAction, + type EnvironmentState, + type EnvironmentStepResult, + type Policy, + runEpisode, +} from "@ignitionai/agent-trainer-environment"; +import { + createOfflinePolicyRecordsFromTrajectories, + exportTrajectoryReport, + recordEpisodeTrajectory, + summarizeTrajectory, + toMarkdownTrajectoryReport, +} from "@ignitionai/agent-trainer-rl"; + +type RagActionName = "search" | "rerank" | "verify" | "answer"; + +class ScriptedRagPolicy implements Policy { + private readonly sequence: RagActionName[] = ["search", "rerank", "verify", "answer"]; + private index = 0; + + async chooseAction( + _state: EnvironmentState, + actions: EnvironmentAction[], + ): Promise { + const expected = this.sequence[this.index] ?? "answer"; + this.index += 1; + const action = actions.find((candidate) => candidate.name === expected); + if (action === undefined) throw new Error(`Expected action is unavailable: ${expected}`); + return action; + } +} + +class MockRagEnvironment implements AgentEnvironment { + private stage: RagActionName | "start" | "done" = "start"; + + async reset(seed?: number): Promise { + this.stage = "start"; + return state("start", { + query: "What is the annual subscription refund policy?", + seed: seed ?? 0, + }); + } + + async actions(): Promise { + if (this.stage === "start") return [{ name: "search", input: { topK: 5 } }]; + if (this.stage === "search") return [{ name: "rerank", input: { keep: 2 } }]; + if (this.stage === "rerank") return [{ name: "verify", input: { requireCitation: true } }]; + if (this.stage === "verify") return [{ name: "answer" }]; + return []; + } + + async step(_state: EnvironmentState, action: EnvironmentAction): Promise { + if (action.name === "search") { + this.stage = "search"; + return { + state: state("searched", { retrievedDocuments: 5 }), + reward: { + name: "retrieval_relevance", + score: 0.7, + weight: 0.2, + reason: "Retrieved enough potentially relevant documents.", + }, + done: false, + }; + } + + if (action.name === "rerank") { + this.stage = "rerank"; + return { + state: state("reranked", { retainedDocuments: 2, bestDocument: "billing-handbook" }), + reward: { + name: "rerank_precision", + score: 0.8, + weight: 0.2, + reason: "Kept the most relevant billing policy document.", + }, + done: false, + }; + } + + if (action.name === "verify") { + this.stage = "verify"; + return { + state: state("verified", { citation: "billing-handbook.md#refund-window" }), + reward: { + name: "citation_verified", + score: 1, + weight: 0.2, + reason: "Verified the answer has a policy citation.", + }, + done: false, + }; + } + + this.stage = "done"; + return { + state: state( + "answered", + { + answer: "Annual subscriptions have a 30 day refund window.", + citation: "billing-handbook.md#refund-window", + }, + true, + ), + reward: { + name: "final_answer_quality", + score: 0.9, + weight: 0.4, + reason: "Answered correctly with the verified citation.", + }, + done: true, + metadata: { terminal: true }, + }; + } +} + +const episode = await runEpisode(new MockRagEnvironment(), new ScriptedRagPolicy(), { + seed: 7, + policyId: "scripted-rag-policy", + metadata: { example: "rag-environment-episode" }, +}); + +const trajectory = recordEpisodeTrajectory(episode, { + id: "rag-environment-episode", + startedAt: "2026-06-30T00:00:00.000Z", + endedAt: "2026-06-30T00:00:04.000Z", +}); +const summary = summarizeTrajectory(trajectory); +const offlineRecords = createOfflinePolicyRecordsFromTrajectories([trajectory], { + experimentName: "rag-environment-episode", +}); +const report = exportTrajectoryReport(trajectory, { + generatedAt: "2026-06-30T00:00:05.000Z", +}); + +console.log("RAG environment episode"); +console.table( + episode.steps.map((step) => ({ + action: step.action.name, + reward: (step.reward.score * step.reward.weight).toFixed(3), + done: step.done, + })), +); +console.log(`Total reward: ${episode.totalReward.toFixed(3)}`); +console.log(`Trajectory steps: ${summary.stepCount}`); +console.log(`Offline records: ${offlineRecords.length}`); +console.log(""); +console.log(toMarkdownTrajectoryReport(report)); + +function state( + id: string, + observation: Record, + done = false, +): EnvironmentState { + return { + id, + observation, + ...(done ? { done: true } : {}), + }; +} diff --git a/packages/environment/README.md b/packages/environment/README.md index 38bc30e..6c3213b 100644 --- a/packages/environment/README.md +++ b/packages/environment/README.md @@ -1,8 +1,8 @@ # @ignitionai/agent-trainer-environment -Prototype state/action/reward environment loop primitives for future policy work. +State/action/reward environment loop primitives for future policy work. -Use this package only when experimenting with agent strategy environments. It is not a production RL runtime. +Use this package when modeling deterministic agent strategy episodes before deeper RL or policy optimization. It is not a production RL runtime and it does not train models. ## Current API @@ -17,17 +17,42 @@ Main exports: - episode types: `EpisodeStep`, `EpisodeResult`, - runner helper: `runEpisode`. +`runEpisode(environment, policy, options)` supports: + +- `seed` passed to `environment.reset(seed)`, +- `maxSteps` safety guard, +- `policyId` and `metadata` copied onto the episode result. + +Episode steps include the previous state, action, next state, reward, done flag and optional step metadata. The result includes total reward, average reward and final state. + +## Example + +```ts +const episode = await runEpisode(environment, policy, { + seed: 7, + maxSteps: 10, + policyId: "scripted-rag-policy", +}); +``` + +See `examples/rag-environment-episode` for a full deterministic sequence: + +```txt +search -> rerank -> verify -> answer +``` + +That example records an episode trajectory through `@ignitionai/agent-trainer-rl` and converts it into offline policy records. + ## Alpha Readiness Status -This package is prototype-level. +This package is partial alpha-level. Known gaps: -- no dedicated tests, -- no example, -- no trajectory recorder, -- no policy evaluation tooling, -- no production environment implementation. +- no production environment implementation, +- no durable rollout store, +- no automatic policy optimization loop, +- no CLI integration for environment episodes. ## Non-goals diff --git a/packages/environment/package.json b/packages/environment/package.json index f82bd85..4e531e4 100644 --- a/packages/environment/package.json +++ b/packages/environment/package.json @@ -1,7 +1,7 @@ { "name": "@ignitionai/agent-trainer-environment", "version": "0.1.0-alpha.1", - "description": "Prototype state/action/reward environment primitives for Ignition Agent Trainer.", + "description": "Tested state/action/reward environment episode primitives for Ignition Agent Trainer.", "license": "MIT", "type": "module", "main": "./dist/index.js", diff --git a/packages/environment/src/environment.test.ts b/packages/environment/src/environment.test.ts new file mode 100644 index 0000000..01b507e --- /dev/null +++ b/packages/environment/src/environment.test.ts @@ -0,0 +1,181 @@ +import { describe, expect, it } from "vitest"; +import { + type AgentEnvironment, + type EnvironmentAction, + type EnvironmentState, + type Policy, + runEpisode, +} from "./environment"; + +describe("runEpisode", () => { + it("runs a complete weighted episode and preserves state transitions", async () => { + const environment = new ScriptedEnvironment([ + { + action: "search", + score: 0.5, + weight: 0.4, + nextState: "searched", + }, + { + action: "answer", + score: 1, + weight: 0.6, + nextState: "answered", + done: true, + metadata: { accepted: true }, + }, + ]); + + const result = await runEpisode(environment, new OrderedPolicy(), { + policyId: "ordered-rag", + metadata: { source: "unit-test" }, + }); + + expect(result).toMatchObject({ + totalReward: 0.8, + averageReward: 0.4, + done: true, + policyId: "ordered-rag", + metadata: { source: "unit-test" }, + finalState: { id: "answered", done: true }, + }); + expect(result.steps.map((step) => step.action.name)).toEqual(["search", "answer"]); + expect(result.steps[0]).toMatchObject({ + state: { id: "initial" }, + nextState: { id: "searched" }, + done: false, + }); + expect(result.steps[1]?.metadata).toEqual({ accepted: true }); + }); + + it("passes seed to reset", async () => { + const environment = new ScriptedEnvironment([ + { + action: "answer", + score: 1, + weight: 1, + nextState: "answered", + done: true, + }, + ]); + + await runEpisode(environment, new OrderedPolicy(), { seed: 42 }); + + expect(environment.seed).toBe(42); + }); + + it("stops unsafe infinite episodes with maxSteps", async () => { + const environment = new ScriptedEnvironment([ + { + action: "search", + score: 0.2, + weight: 1, + nextState: "loop", + }, + { + action: "search", + score: 0.2, + weight: 1, + nextState: "loop", + }, + ]); + + await expect(runEpisode(environment, new OrderedPolicy(), { maxSteps: 1 })).rejects.toThrow( + "Episode exceeded maxSteps (1).", + ); + }); + + it("rejects empty action lists", async () => { + const environment: AgentEnvironment = { + async reset() { + return state("initial"); + }, + async actions() { + return []; + }, + async step() { + throw new Error("Unexpected step."); + }, + }; + + await expect(runEpisode(environment, new OrderedPolicy())).rejects.toThrow( + "Environment returned no actions for state initial.", + ); + }); + + it("rejects non-finite reward values", async () => { + const environment = new ScriptedEnvironment([ + { + action: "answer", + score: Number.NaN, + weight: 1, + nextState: "answered", + done: true, + }, + ]); + + await expect(runEpisode(environment, new OrderedPolicy())).rejects.toThrow( + "Environment reward score must be finite for action answer.", + ); + }); +}); + +interface ScriptedTransition { + action: string; + score: number; + weight: number; + nextState: string; + done?: boolean; + metadata?: Record; +} + +class ScriptedEnvironment implements AgentEnvironment { + seed: number | undefined; + private index = 0; + + constructor(private readonly transitions: ScriptedTransition[]) {} + + async reset(seed?: number): Promise { + this.seed = seed; + this.index = 0; + return state("initial"); + } + + async actions(): Promise { + const transition = this.transitions[this.index] ?? this.transitions.at(-1); + return transition === undefined ? [] : [{ name: transition.action }]; + } + + async step() { + const transition = this.transitions[this.index] ?? this.transitions.at(-1); + if (transition === undefined) throw new Error("Missing scripted transition."); + this.index += 1; + + return { + state: state(transition.nextState, transition.done), + reward: { + name: transition.action, + score: transition.score, + weight: transition.weight, + }, + done: transition.done ?? false, + ...(transition.metadata !== undefined ? { metadata: transition.metadata } : {}), + }; + } +} + +class OrderedPolicy implements Policy { + async chooseAction(_state: EnvironmentState, actions: EnvironmentAction[]) { + const action = actions[0]; + if (action === undefined) throw new Error("No action available."); + return action; + } +} + +function state(id: string, done = false): EnvironmentState { + return { + id, + observation: { id }, + ...(done ? { done: true } : {}), + }; +} diff --git a/packages/environment/src/environment.ts b/packages/environment/src/environment.ts index 3c5bb20..ad2f41d 100644 --- a/packages/environment/src/environment.ts +++ b/packages/environment/src/environment.ts @@ -1,4 +1,4 @@ -import type { JsonRecord, JsonValue, RewardResult } from "@ignitionai/agent-trainer-core"; +import type { JsonRecord, JsonValue, Metadata, RewardResult } from "@ignitionai/agent-trainer-core"; export interface EnvironmentState { id: string; @@ -15,7 +15,7 @@ export interface EnvironmentStepResult { state: EnvironmentState; reward: RewardResult; done: boolean; - metadata?: JsonRecord; + metadata?: Metadata; } export interface AgentEnvironment { @@ -31,32 +31,90 @@ export interface Policy { export interface EpisodeStep { state: EnvironmentState; action: EnvironmentAction; + nextState: EnvironmentState; reward: RewardResult; + done: boolean; + metadata?: Metadata; } export interface EpisodeResult { steps: EpisodeStep[]; totalReward: number; + averageReward: number; + finalState: EnvironmentState; + done: boolean; + policyId?: string; + metadata?: Metadata; +} + +export interface RunEpisodeOptions { + seed?: number; + maxSteps?: number; + policyId?: string; + metadata?: Metadata; } export async function runEpisode( environment: AgentEnvironment, policy: Policy, + options: RunEpisodeOptions = {}, ): Promise { - let state = await environment.reset(); + const maxSteps = options.maxSteps ?? 100; + validateMaxSteps(maxSteps); + + let state = await environment.reset(options.seed); const steps: EpisodeStep[] = []; while (!state.done) { + if (steps.length >= maxSteps) { + throw new Error(`Episode exceeded maxSteps (${maxSteps}).`); + } + const actions = await environment.actions(state); + if (actions.length === 0) { + throw new Error(`Environment returned no actions for state ${state.id}.`); + } + const action = await policy.chooseAction(state, actions); const result = await environment.step(state, action); - steps.push({ state, action, reward: result.reward }); + validateReward(result.reward, action); + + steps.push({ + state, + action, + nextState: result.state, + reward: result.reward, + done: result.done, + ...(result.metadata !== undefined ? { metadata: result.metadata } : {}), + }); state = result.state; if (result.done) break; } + const totalReward = steps.reduce((sum, step) => sum + step.reward.score * step.reward.weight, 0); + return { steps, - totalReward: steps.reduce((sum, step) => sum + step.reward.score * step.reward.weight, 0), + totalReward, + averageReward: steps.length === 0 ? 0 : totalReward / steps.length, + finalState: state, + done: state.done === true || steps.at(-1)?.done === true, + ...(options.policyId !== undefined ? { policyId: options.policyId } : {}), + ...(options.metadata !== undefined ? { metadata: options.metadata } : {}), }; } + +function validateMaxSteps(maxSteps: number): void { + if (!Number.isInteger(maxSteps) || maxSteps <= 0) { + throw new Error("Episode maxSteps must be a positive integer."); + } +} + +function validateReward(reward: RewardResult, action: EnvironmentAction): void { + if (!Number.isFinite(reward.score)) { + throw new Error(`Environment reward score must be finite for action ${action.name}.`); + } + if (!Number.isFinite(reward.weight)) { + throw new Error(`Environment reward weight must be finite for action ${action.name}.`); + } +} diff --git a/packages/rl/README.md b/packages/rl/README.md index 18368d0..89735f9 100644 --- a/packages/rl/README.md +++ b/packages/rl/README.md @@ -47,6 +47,28 @@ const summary = summarizeTrajectory(trajectory); Trajectories are plain JSON-compatible records. They are not an external tracing service and they do not imply a training loop. +Environment episodes from `@ignitionai/agent-trainer-environment` can be converted directly: + +```ts +import { + exportTrajectoryReport, + recordEpisodeTrajectory, + toMarkdownTrajectoryReport, +} from "@ignitionai/agent-trainer-rl"; + +const trajectory = recordEpisodeTrajectory(episode, { + id: "rag-environment-episode", +}); +const report = exportTrajectoryReport(trajectory); +const markdown = toMarkdownTrajectoryReport(report); +``` + +See `examples/rag-environment-episode` for a complete mocked RAG episode that runs: + +```txt +search -> rerank -> verify -> answer +``` + ## Experimental Fixed-Strategy Bandit Use `ExperimentalBanditStrategySelector` to choose among fixed, developer-supplied strategies and update their rewards from observed experiment outcomes. @@ -154,6 +176,15 @@ const result = await evaluatePolicyOffline(createScoreBasedPolicy(), [ `createOfflinePolicyRecordsFromTrajectories()` can convert recorded trajectories into observed-action records. Those records only know the reward for the logged action unless you provide richer offline records yourself. +The same helper works with trajectories created from environment episodes: + +```ts +const trajectory = recordEpisodeTrajectory(episode); +const records = createOfflinePolicyRecordsFromTrajectories([trajectory], { + experimentName: "rag-environment", +}); +``` + ## GRPO-Style Candidate Selection Use `selectGroupRelativeBest()` when prompt, workflow or strategy candidates should be ranked by their advantage relative to other candidates in the same group. diff --git a/packages/rl/src/episode-trajectory.test.ts b/packages/rl/src/episode-trajectory.test.ts new file mode 100644 index 0000000..f2873a7 --- /dev/null +++ b/packages/rl/src/episode-trajectory.test.ts @@ -0,0 +1,176 @@ +import { + type AgentEnvironment, + type EnvironmentAction, + type Policy as EnvironmentPolicy, + type EnvironmentState, + runEpisode, +} from "@ignitionai/agent-trainer-environment"; +import { describe, expect, it } from "vitest"; +import { + exportTrajectoryReport, + recordEpisodeTrajectory, + toMarkdownTrajectoryReport, +} from "./episode-trajectory"; +import { + createOfflinePolicyRecordsFromTrajectories, + evaluatePolicyOffline, +} from "./offline-policy-evaluation"; +import { createStaticPolicy } from "./policy"; + +describe("episode trajectory bridge", () => { + it("records environment episodes as deterministic trajectories", async () => { + const episode = await runEpisode(new RerankEnvironment(), new ScriptedPolicy(), { + policyId: "scripted-rag", + metadata: { caseId: "case-1" }, + }); + + const trajectory = recordEpisodeTrajectory(episode, { + id: "episode-trajectory", + startedAt: "2026-06-30T10:00:00.000Z", + endedAt: "2026-06-30T10:00:03.000Z", + metadata: { source: "environment" }, + }); + + expect(trajectory).toMatchObject({ + id: "episode-trajectory", + policyId: "scripted-rag", + startedAt: "2026-06-30T10:00:00.000Z", + endedAt: "2026-06-30T10:00:03.000Z", + metadata: { source: "environment", caseId: "case-1" }, + }); + expect(trajectory.steps.map((step) => step.index)).toEqual([0, 1, 2]); + expect(trajectory.steps.map((step) => step.action)).toEqual([ + { id: "search", name: "search" }, + { id: "rerank", name: "rerank" }, + { id: "answer", name: "answer" }, + ]); + expect(trajectory.steps.map((step) => step.reward)).toEqual([0.25, 0.3, 0.45]); + expect(trajectory.steps.at(-1)?.outcome).toEqual({ + nextStateId: "answered", + done: true, + }); + }); + + it("supports offline policy evaluation from episode trajectories", async () => { + const episode = await runEpisode(new RerankEnvironment(), new ScriptedPolicy(), { + policyId: "scripted-rag", + }); + const trajectory = recordEpisodeTrajectory(episode, { id: "episode-trajectory" }); + const records = createOfflinePolicyRecordsFromTrajectories([trajectory], { + experimentName: "rag-environment", + }); + + expect(records.map((record) => record.id)).toEqual([ + "episode-trajectory:0", + "episode-trajectory:1", + "episode-trajectory:2", + ]); + + const answerRecord = records.find((record) => record.expectedCandidateId === "answer"); + if (answerRecord === undefined) throw new Error("Expected answer offline record."); + + const result = await evaluatePolicyOffline(createStaticPolicy("answer"), [answerRecord], { + policyId: "static-answer", + }); + + expect(result).toMatchObject({ + policyId: "static-answer", + recordCount: 1, + totalReward: 0.45, + averageReward: 0.45, + accuracy: 1, + }); + }); + + it("exports stable JSON and Markdown trajectory reports", async () => { + const episode = await runEpisode(new RerankEnvironment(), new ScriptedPolicy(), { + policyId: "scripted-rag", + }); + const trajectory = recordEpisodeTrajectory(episode, { id: "episode-trajectory" }); + const report = exportTrajectoryReport(trajectory, { + generatedAt: "2026-06-30T12:00:00.000Z", + metadata: { suite: "rl" }, + }); + + expect(report).toMatchObject({ + schemaVersion: "ignition.trajectory-report.v1", + generatedAt: "2026-06-30T12:00:00.000Z", + trajectory: { + id: "episode-trajectory", + policyId: "scripted-rag", + }, + summary: { + trajectoryId: "episode-trajectory", + stepCount: 3, + totalReward: 1, + }, + metadata: { suite: "rl" }, + }); + + const markdown = toMarkdownTrajectoryReport(report); + + expect(markdown).toContain("# Trajectory report: episode-trajectory"); + expect(markdown).toContain("Policy: scripted-rag"); + expect(markdown).toContain("Total reward: 1.000"); + expect(markdown).toContain("| 2 | answer | 0.450 | yes |"); + }); +}); + +class RerankEnvironment implements AgentEnvironment { + private index = 0; + + async reset(): Promise { + this.index = 0; + return state("start", { query: "refund policy" }); + } + + async actions(): Promise { + return [{ name: ["search", "rerank", "answer"][this.index] ?? "answer" }]; + } + + async step(_state: EnvironmentState, action: EnvironmentAction) { + this.index += 1; + + if (action.name === "search") { + return { + state: state("searched", { docs: 3 }), + reward: { name: "retrieval", score: 0.5, weight: 0.5 }, + done: false, + }; + } + + if (action.name === "rerank") { + return { + state: state("reranked", { docs: 2 }), + reward: { name: "rerank", score: 0.6, weight: 0.5 }, + done: false, + }; + } + + return { + state: state("answered", { cited: true }, true), + reward: { name: "answer", score: 0.9, weight: 0.5 }, + done: true, + }; + } +} + +class ScriptedPolicy implements EnvironmentPolicy { + async chooseAction(_state: EnvironmentState, actions: EnvironmentAction[]) { + const action = actions[0]; + if (action === undefined) throw new Error("Expected action."); + return action; + } +} + +function state( + id: string, + observation: Record, + done = false, +): EnvironmentState { + return { + id, + observation, + ...(done ? { done: true } : {}), + }; +} diff --git a/packages/rl/src/episode-trajectory.ts b/packages/rl/src/episode-trajectory.ts new file mode 100644 index 0000000..2bdbb1c --- /dev/null +++ b/packages/rl/src/episode-trajectory.ts @@ -0,0 +1,188 @@ +import type { JsonRecord, JsonValue, Metadata } from "@ignitionai/agent-trainer-core"; +import type { EnvironmentAction, EpisodeResult } from "@ignitionai/agent-trainer-environment"; +import { + recordTrajectory, + summarizeTrajectory, + type Trajectory, + type TrajectorySummary, +} from "./trajectory"; + +export interface RecordEpisodeTrajectoryOptions { + id?: string; + startedAt?: string | Date; + endedAt?: string | Date; + metadata?: Metadata; +} + +export interface TrajectoryReportOptions { + generatedAt?: string | Date; + metadata?: Metadata; +} + +export interface TrajectoryReportStep { + index: number; + state: JsonRecord; + action: JsonValue; + reward: number; + outcome?: JsonValue; + metadata?: Metadata; +} + +export interface TrajectoryReport { + schemaVersion: "ignition.trajectory-report.v1"; + generatedAt: string; + trajectory: { + id: string; + policyId?: string; + startedAt?: string; + endedAt?: string; + metadata?: Metadata; + }; + summary: TrajectorySummary; + steps: TrajectoryReportStep[]; + metadata?: Metadata; +} + +export function recordEpisodeTrajectory( + episode: EpisodeResult, + options: RecordEpisodeTrajectoryOptions = {}, +): Trajectory { + return recordTrajectory( + episode.steps.map((step) => ({ + state: { + stateId: step.state.id, + observation: step.state.observation, + done: step.state.done === true, + }, + action: actionToTrajectoryAction(step.action), + reward: step.reward.score * step.reward.weight, + outcome: { + nextStateId: step.nextState.id, + done: step.done, + }, + metadata: { + rewardName: step.reward.name, + rewardScore: step.reward.score, + rewardWeight: step.reward.weight, + nextObservation: step.nextState.observation, + ...(step.metadata ?? {}), + }, + })), + { + id: options.id ?? "episode-trajectory", + ...(episode.policyId !== undefined ? { policyId: episode.policyId } : {}), + ...(options.startedAt !== undefined ? { startedAt: options.startedAt } : {}), + ...(options.endedAt !== undefined ? { endedAt: options.endedAt } : {}), + metadata: { + ...(episode.metadata ?? {}), + ...(options.metadata ?? {}), + }, + }, + ); +} + +export function exportTrajectoryReport( + trajectory: Trajectory, + options: TrajectoryReportOptions = {}, +): TrajectoryReport { + return { + schemaVersion: "ignition.trajectory-report.v1", + generatedAt: normalizeTimestamp(options.generatedAt), + trajectory: { + id: trajectory.id, + ...(trajectory.policyId !== undefined ? { policyId: trajectory.policyId } : {}), + ...(trajectory.startedAt !== undefined ? { startedAt: trajectory.startedAt } : {}), + ...(trajectory.endedAt !== undefined ? { endedAt: trajectory.endedAt } : {}), + ...(trajectory.metadata !== undefined ? { metadata: trajectory.metadata } : {}), + }, + summary: summarizeTrajectory(trajectory), + steps: trajectory.steps.map((step) => ({ + index: step.index, + state: step.state, + action: step.action, + reward: step.reward, + ...(step.outcome !== undefined ? { outcome: step.outcome } : {}), + ...(step.metadata !== undefined ? { metadata: step.metadata } : {}), + })), + ...(options.metadata !== undefined ? { metadata: options.metadata } : {}), + }; +} + +export function toMarkdownTrajectoryReport(report: TrajectoryReport): string { + const sections = [ + `# Trajectory report: ${report.trajectory.id}`, + [ + `Generated: ${report.generatedAt}`, + `Policy: ${report.trajectory.policyId ?? "n/a"}`, + `Steps: ${report.summary.stepCount}`, + `Total reward: ${formatReward(report.summary.totalReward)}`, + `Average reward: ${formatReward(report.summary.averageReward)}`, + ].join("\n"), + "## Actions", + actionTable(report.summary.actionCounts), + "## Steps", + stepTable(report.steps), + ]; + + return `${sections.join("\n\n")}\n`; +} + +function actionToTrajectoryAction(action: EnvironmentAction): JsonRecord { + return { + id: action.name, + name: action.name, + ...(action.input !== undefined ? { input: action.input } : {}), + }; +} + +function actionTable(actionCounts: Record): string { + const rows = Object.entries(actionCounts); + if (rows.length === 0) return "No actions recorded."; + + return [ + "| Action | Count |", + "|---|---:|", + ...rows.map(([action, count]) => `| ${action} | ${count} |`), + ].join("\n"); +} + +function stepTable(steps: readonly TrajectoryReportStep[]): string { + if (steps.length === 0) return "No steps recorded."; + + return [ + "| Step | Action | Reward | Done |", + "|---:|---|---:|---|", + ...steps.map( + (step) => + `| ${step.index} | ${actionName(step.action)} | ${formatReward(step.reward)} | ${isDone(step.outcome) ? "yes" : "no"} |`, + ), + ].join("\n"); +} + +function actionName(action: JsonValue): string { + if (typeof action === "string") return action; + if (isRecord(action)) { + const name = action.name; + if (typeof name === "string") return name; + const id = action.id; + if (typeof id === "string") return id; + } + return JSON.stringify(action); +} + +function isDone(outcome: JsonValue | undefined): boolean { + return isRecord(outcome) && outcome.done === true; +} + +function formatReward(value: number): string { + return value.toFixed(3); +} + +function normalizeTimestamp(value: string | Date | undefined): string { + if (value === undefined) return new Date().toISOString(); + return value instanceof Date ? value.toISOString() : value; +} + +function isRecord(value: JsonValue | undefined): value is JsonRecord { + return typeof value === "object" && value !== null && !Array.isArray(value); +} diff --git a/packages/rl/src/index.ts b/packages/rl/src/index.ts index 76c1557..8e97db7 100644 --- a/packages/rl/src/index.ts +++ b/packages/rl/src/index.ts @@ -1,5 +1,6 @@ export * from "./bandit"; export * from "./contextual-bandit"; +export * from "./episode-trajectory"; export * from "./group-relative-selection"; export * from "./offline-policy-evaluation"; export * from "./policy";