feat(arena): add deterministic headless RL environments
This commit is contained in:
@@ -0,0 +1,316 @@
|
||||
import assert from "node:assert/strict";
|
||||
import { describe, it } from "node:test";
|
||||
import {
|
||||
ARENA_API_VERSION,
|
||||
ARENA_MANIFESTS,
|
||||
ARENA_SOURCE_HASHES,
|
||||
CALIFORNIA_FLIGHT_INACTION,
|
||||
CALIFORNIA_FLIGHT_SCENARIOS,
|
||||
CROW_NAV_INACTION,
|
||||
CROW_NAV_SCENARIOS,
|
||||
DRIVE_101_SCENARIOS,
|
||||
DRIVE_INACTION,
|
||||
OFFICE_NAV_INACTION,
|
||||
OFFICE_NAV_SCENARIOS,
|
||||
CaliforniaFlightEnvironment,
|
||||
CrowNavEnvironment,
|
||||
Drive101Environment,
|
||||
OfficeNavEnvironment,
|
||||
arenaChecksum,
|
||||
californiaFlightScriptedBaseline,
|
||||
crowNavScriptedBaseline,
|
||||
driveScriptedBaseline,
|
||||
officeNavScriptedBaseline,
|
||||
type ArenaEnvironment,
|
||||
type ArenaManifest,
|
||||
type ArenaScenarioRegistry,
|
||||
type ArenaStepResult,
|
||||
} from "../index.ts";
|
||||
|
||||
type NumericRewards = Record<string, number>;
|
||||
type AnyEnvironment = ArenaEnvironment<unknown, unknown, NumericRewards, unknown>;
|
||||
|
||||
interface EnvironmentCase {
|
||||
name: string;
|
||||
create(): AnyEnvironment;
|
||||
registry: ArenaScenarioRegistry<object>;
|
||||
inaction: unknown;
|
||||
scripted(observation: unknown): unknown;
|
||||
}
|
||||
|
||||
const CASES: EnvironmentCase[] = [
|
||||
{
|
||||
name: "drive",
|
||||
create: () => new Drive101Environment() as AnyEnvironment,
|
||||
registry: DRIVE_101_SCENARIOS as ArenaScenarioRegistry<object>,
|
||||
inaction: DRIVE_INACTION,
|
||||
scripted: () => driveScriptedBaseline(),
|
||||
},
|
||||
{
|
||||
name: "office",
|
||||
create: () => new OfficeNavEnvironment() as AnyEnvironment,
|
||||
registry: OFFICE_NAV_SCENARIOS as ArenaScenarioRegistry<object>,
|
||||
inaction: OFFICE_NAV_INACTION,
|
||||
scripted: (observation) => officeNavScriptedBaseline(
|
||||
observation as Parameters<typeof officeNavScriptedBaseline>[0],
|
||||
),
|
||||
},
|
||||
{
|
||||
name: "crow",
|
||||
create: () => new CrowNavEnvironment() as AnyEnvironment,
|
||||
registry: CROW_NAV_SCENARIOS as ArenaScenarioRegistry<object>,
|
||||
inaction: CROW_NAV_INACTION,
|
||||
scripted: (observation) => crowNavScriptedBaseline(
|
||||
observation as Parameters<typeof crowNavScriptedBaseline>[0],
|
||||
),
|
||||
},
|
||||
{
|
||||
name: "flight",
|
||||
create: () => new CaliforniaFlightEnvironment() as AnyEnvironment,
|
||||
registry: CALIFORNIA_FLIGHT_SCENARIOS as ArenaScenarioRegistry<object>,
|
||||
inaction: CALIFORNIA_FLIGHT_INACTION,
|
||||
scripted: (observation) => californiaFlightScriptedBaseline(
|
||||
observation as Parameters<typeof californiaFlightScriptedBaseline>[0],
|
||||
),
|
||||
},
|
||||
];
|
||||
|
||||
function run(
|
||||
entry: EnvironmentCase,
|
||||
scenarioId: string,
|
||||
seed: number,
|
||||
policy: (observation: unknown) => unknown,
|
||||
): { total: number; final: ArenaStepResult<unknown, NumericRewards> } {
|
||||
const environment = entry.create();
|
||||
let observation = environment.reset(seed, scenarioId).observation;
|
||||
let total = 0;
|
||||
let final: ArenaStepResult<unknown, NumericRewards> | undefined;
|
||||
for (let step = 0; step < environment.manifest.maxSteps; step += 1) {
|
||||
final = environment.step(policy(observation));
|
||||
observation = final.observation;
|
||||
total += final.reward;
|
||||
if (final.terminated || final.truncated) break;
|
||||
}
|
||||
if (!final) throw new Error("environment manifest must permit at least one step");
|
||||
return { total, final };
|
||||
}
|
||||
|
||||
describe("arena contract and manifests", () => {
|
||||
it("exports four versioned renderer-independent manifests with disjoint public splits", () => {
|
||||
assert.deepEqual(ARENA_MANIFESTS.map((manifest: ArenaManifest) => manifest.id), [
|
||||
"drive-101-v1", "office-nav-v1", "crow-nav-v1", "california-flight-v1",
|
||||
]);
|
||||
for (const manifest of ARENA_MANIFESTS) {
|
||||
assert.equal(manifest.apiVersion, ARENA_API_VERSION);
|
||||
assert.ok(manifest.version >= 1 && manifest.maxSteps > 0);
|
||||
assert.ok(manifest.safetyTerminals.length > 0);
|
||||
assert.ok(Object.keys(manifest.rewardComponents).length >= 6);
|
||||
assert.equal(
|
||||
manifest.scenarioIds.train.some((id) => manifest.scenarioIds.dev.includes(id)),
|
||||
false,
|
||||
);
|
||||
assert.match(ARENA_SOURCE_HASHES[manifest.id]!.environment, /^sha256:[0-9a-f]{64}$/);
|
||||
assert.match(ARENA_SOURCE_HASHES[manifest.id]!.simulator, /^sha256:[0-9a-f]{64}$/);
|
||||
}
|
||||
});
|
||||
|
||||
it("selects and materializes scenarios deterministically by split, id and seed", () => {
|
||||
for (const entry of CASES) {
|
||||
const a = entry.create().reset(0x1234abcd, { split: "train" });
|
||||
const b = entry.create().reset(0x1234abcd, { split: "train" });
|
||||
assert.deepEqual(a, b, entry.name);
|
||||
assert.equal(a.info.scenarioSplit, "train");
|
||||
assert.match(a.info.scenarioHash, /^fnv1a64:[0-9a-f]{16}$/);
|
||||
const different = entry.create().reset(0x1234abce, a.info.scenarioId);
|
||||
assert.notEqual(different.info.scenarioHash, a.info.scenarioHash);
|
||||
assert.throws(
|
||||
() => entry.create().reset(1, { split: "dev", id: entry.registry.ids("train")[0] }),
|
||||
/unknown/,
|
||||
);
|
||||
}
|
||||
});
|
||||
|
||||
it("sanitizes non-finite and out-of-range actions before state or reward hashing", () => {
|
||||
const malformed = {
|
||||
throttle: Number.POSITIVE_INFINITY,
|
||||
brake: Number.NaN,
|
||||
steering: -99,
|
||||
handbrake: false,
|
||||
};
|
||||
const first = new Drive101Environment();
|
||||
const second = new Drive101Environment();
|
||||
first.reset(4, "train-us101-ventura");
|
||||
second.reset(4, "train-us101-ventura");
|
||||
const a = first.step(malformed);
|
||||
const b = second.step({ throttle: 0, brake: 0, steering: -1, handbrake: false });
|
||||
assert.deepEqual(a, b);
|
||||
assert.ok(Number.isFinite(a.reward));
|
||||
});
|
||||
|
||||
it("uses terminated for outcomes, truncated only for max-step exhaustion", () => {
|
||||
for (const entry of CASES) {
|
||||
const id = entry.registry.ids("train")[0]!;
|
||||
const idle = run(entry, id, 5, () => entry.inaction).final;
|
||||
assert.equal(idle.terminated, false, entry.name);
|
||||
assert.equal(idle.truncated, true, entry.name);
|
||||
assert.equal(idle.info.terminalReason, "max-steps", entry.name);
|
||||
|
||||
const scripted = run(entry, id, 5, entry.scripted).final;
|
||||
assert.equal(scripted.terminated, true, entry.name);
|
||||
assert.equal(scripted.truncated, false, entry.name);
|
||||
assert.equal(scripted.info.terminalReason, "goal", entry.name);
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
describe("arena baseline proofs", () => {
|
||||
it("keeps inaction below zero and proves a positive scripted completion on every public scenario", () => {
|
||||
for (const entry of CASES) {
|
||||
for (const definition of entry.registry.definitions) {
|
||||
for (const seed of [1, 0xdecafbad]) {
|
||||
const idle = run(entry, definition.id, seed, () => entry.inaction);
|
||||
assert.ok(idle.total < 0, `${entry.name}/${definition.id}/${seed} inaction=${idle.total}`);
|
||||
assert.notEqual(idle.final.info.terminalReason, "goal");
|
||||
|
||||
const scripted = run(entry, definition.id, seed, entry.scripted);
|
||||
assert.equal(
|
||||
scripted.final.info.terminalReason,
|
||||
"goal",
|
||||
`${entry.name}/${definition.id}/${seed}`,
|
||||
);
|
||||
assert.ok(scripted.total > 0, `${entry.name}/${definition.id}/${seed}=${scripted.total}`);
|
||||
assert.ok(scripted.total > idle.total);
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
it("exposes reachable safety terminals rather than reward-only safety labels", () => {
|
||||
const drive = new Drive101Environment();
|
||||
drive.reset(2, "train-us101-ventura");
|
||||
let driveReason: string | null = null;
|
||||
for (let index = 0; index < drive.manifest.maxSteps; index += 1) {
|
||||
const result = drive.step({ throttle: 1, brake: 0, steering: 1, handbrake: false });
|
||||
driveReason = result.info.terminalReason;
|
||||
if (result.terminated) break;
|
||||
}
|
||||
assert.equal(driveReason, "guardrail-contact");
|
||||
|
||||
const office = new OfficeNavEnvironment();
|
||||
office.reset(2, "train-hangar-crossing");
|
||||
let officeReason: string | null = null;
|
||||
for (let index = 0; index < office.manifest.maxSteps; index += 1) {
|
||||
const result = office.step({ x: 0, z: 1 });
|
||||
officeReason = result.info.terminalReason;
|
||||
if (result.terminated) break;
|
||||
}
|
||||
assert.equal(officeReason, "collision-stall");
|
||||
|
||||
const crow = new CrowNavEnvironment();
|
||||
crow.reset(2, "train-east-crosswind");
|
||||
let crowReason: string | null = null;
|
||||
for (let index = 0; index < crow.manifest.maxSteps; index += 1) {
|
||||
const result = crow.step({ forward: 1, turn: 0, pitch: 0, climb: 1, glide: false });
|
||||
crowReason = result.info.terminalReason;
|
||||
if (result.terminated) break;
|
||||
}
|
||||
assert.equal(crowReason, "flight-envelope-contact");
|
||||
|
||||
const flight = new CaliforniaFlightEnvironment();
|
||||
flight.reset(2, "train-la-east-leg");
|
||||
let flightReason: string | null = null;
|
||||
for (let index = 0; index < flight.manifest.maxSteps; index += 1) {
|
||||
const result = flight.step({ throttle: 1, yaw: 0, pitch: 1, roll: 0 });
|
||||
flightReason = result.info.terminalReason;
|
||||
if (result.terminated) break;
|
||||
}
|
||||
assert.equal(flightReason, "flight-envelope-contact");
|
||||
});
|
||||
});
|
||||
|
||||
describe("arena snapshot, trace and replay", () => {
|
||||
it("restores each simulator bit-for-bit and preserves the next transition", () => {
|
||||
for (const entry of CASES) {
|
||||
const environment = entry.create();
|
||||
let observation = environment.reset(91, entry.registry.ids("dev")[0]!).observation;
|
||||
for (let index = 0; index < 12; index += 1) {
|
||||
observation = environment.step(entry.scripted(observation)).observation;
|
||||
}
|
||||
const checkpoint = environment.snapshot();
|
||||
const action = entry.scripted(observation);
|
||||
const expected = environment.step(action);
|
||||
const restored = environment.restore(checkpoint);
|
||||
assert.equal(restored.info.step, checkpoint.step);
|
||||
const actual = environment.step(action);
|
||||
assert.deepEqual(actual, expected, entry.name);
|
||||
}
|
||||
});
|
||||
|
||||
it("replays checksummed action traces to the exact final state", () => {
|
||||
for (const entry of CASES) {
|
||||
const environment = entry.create();
|
||||
let observation = environment.reset(144, entry.registry.ids("train")[0]!).observation;
|
||||
for (let index = 0; index < 30; index += 1) {
|
||||
const result = environment.step(entry.scripted(observation));
|
||||
observation = result.observation;
|
||||
if (result.terminated) break;
|
||||
}
|
||||
const trace = environment.trace();
|
||||
assert.equal(arenaChecksum({ ...trace, checksum: undefined }), trace.checksum);
|
||||
const replay = entry.create().replay(trace);
|
||||
assert.equal(replay.finalStateChecksum, trace.finalStateChecksum, entry.name);
|
||||
assert.equal(replay.cumulativeReward, trace.cumulativeReward, entry.name);
|
||||
assert.equal(replay.steps, trace.steps.length, entry.name);
|
||||
}
|
||||
});
|
||||
|
||||
it("rejects tampered snapshots and traces before applying them", () => {
|
||||
const environment = new Drive101Environment();
|
||||
environment.reset(7, "train-us101-ventura");
|
||||
environment.step(driveScriptedBaseline());
|
||||
const snapshot = environment.snapshot();
|
||||
assert.throws(
|
||||
() => environment.restore({ ...snapshot, cumulativeReward: snapshot.cumulativeReward + 1 }),
|
||||
/checksum mismatch/,
|
||||
);
|
||||
const trace = environment.trace();
|
||||
assert.throws(
|
||||
() => new Drive101Environment().replay({ ...trace, cumulativeReward: trace.cumulativeReward + 1 }),
|
||||
/checksum mismatch/,
|
||||
);
|
||||
|
||||
const invalidCore = { ...snapshot, step: -1 };
|
||||
const { checksum: _oldChecksum, ...invalidWithoutChecksum } = invalidCore;
|
||||
assert.throws(
|
||||
() => environment.restore({
|
||||
...invalidWithoutChecksum,
|
||||
checksum: arenaChecksum(invalidWithoutChecksum),
|
||||
}),
|
||||
/episode state is invalid/,
|
||||
);
|
||||
|
||||
const invalidFrames = trace.steps.map((frame, index) => ({
|
||||
...frame,
|
||||
index: index === 0 ? 2 : frame.index,
|
||||
}));
|
||||
const { checksum: _traceChecksum, ...invalidTraceCore } = { ...trace, steps: invalidFrames };
|
||||
assert.throws(
|
||||
() => new Drive101Environment().replay({
|
||||
...invalidTraceCore,
|
||||
checksum: arenaChecksum(invalidTraceCore),
|
||||
}),
|
||||
/frame 1 is invalid/,
|
||||
);
|
||||
});
|
||||
|
||||
it("refuses post-terminal stepping until reset", () => {
|
||||
const environment = new OfficeNavEnvironment();
|
||||
let observation = environment.reset(5, "train-galley-aisle").observation;
|
||||
for (;;) {
|
||||
const result = environment.step(officeNavScriptedBaseline(observation));
|
||||
observation = result.observation;
|
||||
if (result.terminated) break;
|
||||
}
|
||||
assert.throws(() => environment.step(OFFICE_NAV_INACTION), /episode is complete/);
|
||||
});
|
||||
});
|
||||
Reference in New Issue
Block a user