Skip to content
angadjosanPublic

About

Environment state cache for RL agent training: fork and reuse sandbox state.

Resources

Stars

0 stars

Watchers

0 watching

Forks

Repository files navigation

TrajCache

Naming environment state, so it can be reused, branched, and shared.

The design argument is in idea.md. This is the implementation of it: a Rust workspace that gives every reachable environment state a stable name, reuses states across runs and across machines, branches them without copying, and refuses to serve a name it cannot stand behind.

$ trajcache --annotations examples/agent-tools.json run examples/setup.json --twice
run 1        Executed in 781.958µs
run 2        CacheHit in 39.5µs
speedup      20x

What is here

crate what it does
tool-annotations The companion spec from §10. Five declarations — mutates, idempotent, commutes_with, inverse_of, deterministic — over a tri-state that keeps undeclared distinguishable from declared false. Zero dependencies, enforced by a test that reads the manifest.
trajcache-core The naming layer. Environment fingerprints, the prefix-free canonical encoding, canonicalization, and the two names of §4 as two distinct Rust types.
trajcache-store Content-addressed blobs, the copy-on-write state graph, deferred materialization, eviction by recompute cost.
trajcache-index The learned provenance→content mapping, divergence detection, shadow verification, tool quarantine.
trajcache-sandbox Filesystem ↔ manifest, tool runtimes (including real /bin/sh), the §7 determinism envelope, and the entropy carve-out. Also the adapter seam: SandboxBackend + FsProvider, and a conformance suite anyone can run against their own — see PROVIDERS.md.
trajcache-engine Single-flight realization, branch trees, the ordered event bus.
trajcache-server The §6 registry: content-addressed distribution with signed provenance attestation.
trajcache-harness Phase 0. Canonicalization density, branch divergence, per-rule contribution.
trajcache-cli trajcache name / lookup / run / branch / measure / annotations / serve.

Design decisions are logged in decisions.md; the things they cost are in tradeoffs.md. Both are written to be argued with. PROVIDERS.md is how to put this on top of your own sandboxes.

The idea in sixty seconds

A state gets two names, doing two different jobs:

  • A provenance name — the environment fingerprint plus the canonicalized action sequence. Available before the state exists, which is what makes "has anyone already run these commands?" a question you can ask cheaply.
  • A content name — a Merkle hash over the resulting state. Available only after. This is the identity: it deduplicates states reached by different paths, and it is what you verify when accepting a state from someone else.

They are separate Rust types over the same 32 bytes, so the compiler refuses to let a prediction stand in for a confirmation. The mapping between them is learned, not assumed: first execution records it, later lookups consult it, a sampled fraction re-executes to check it, and a mismatch quarantines the tool that claimed to be deterministic and was not.

What it does, measured

Every number below comes from a test in this repository, on a debug build. cargo test reproduces all of them.

Cross-run reuse. A real /bin/sh workload: cold 30.6 ms, warm 14.4 µs — a 2100× speedup, and the second run performs no setup at all.

Canonicalization. Two runs that install the same three packages in different orders reach one name and one execution; 3! orderings collapse. A sequence of five actions with two reads and a repeat canonicalizes to one.

Branch storage (§12's "the number that decides whether this scales"), from a_tree_structured_rollout_costs_the_prefix_once:

tree width   2:     200 B of divergence (100 B per sibling) vs    66368 B if copied — 332x smaller
tree width   8:     800 B                                   vs   265472 B           — 332x
tree width  32:    3254 B                                   vs  1061888 B           — 326x
tree width  64:    6518 B                                   vs  2123776 B           — 326x

Per-sibling cost is flat in tree width and flat in the size of the state each branch starts from. Forking 4096 branches of a 131 KB state takes 60 µs total — 0.01 µs per branch — because a fork is a name plus a pin, and nothing is copied or materialized until something reads it.

Latency, methodology stated in crates/trajcache-engine/tests/latency.rs:

operation concurrency p50 p99 p99.9
name a 7-action trajectory 1 41 µs 47 µs 52 µs
lookup (miss) 1 5.8 µs 7.3 µs 8.9 µs
realize (warm hit) 1 42 µs 51 µs 58 µs
realize (warm hit) 64 46 µs 128 µs 962 µs
realize (warm hit) 256 46 µs 128 µs 473 µs

Single-flight. 1000 rollouts starting the same setup at the same instant produce exactly one execution.

Phase 0 (trajcache measure), over synthetic workloads shaped like terminal-bench, SWE-Gym, SkyRL-SQL and a browser session:

environment                               len↓  log₁₀ord   block  commuting  names/32  annotated
opaque-shell (terminal-bench-like)        1.00       0.0       0       0.0%       1.0        20%
repository-agent (SWE-Gym-like)           0.69      24.0      20      43.0%       1.0        58%
independent-queries (SkyRL-SQL-like)      0.16       9.3      32       6.2%       1.0        56%
browser-session                           0.67       0.0       0       0.0%       1.0        40%

Read that table as the shape of the answer, not as a measurement of those benchmarks — the workloads are modelled on them, not taken from them, and every report carries synthetic: true so a number from here can never be quoted as something it is not. What it does say is that rule 1 (elide) carries most of the weight, that commutation is worth real density where it exists at all, and that an undescribed environment gets exactly nothing — which is the behaviour §13 promises.

Three things worth reading the code for

Rule 3 is a Foata normal form, not a greedy sort (D6). The obvious implementation — walk the sequence, sort maximal runs of commuting calls — is not confluent when the commutation relation is non-transitive, which it always is for real tools. Two sequences that reach the same state get different names, and every miss is a rebuild nobody can explain. tests/confluence.rs enumerates every sequence over a 4-letter alphabet up to length 6, computes the true equivalence classes by exhaustive closure, and asserts the canonical form is constant within each class and distinct across classes.

The entropy carve-out is the §7 tension, resolved and then corrected (D21). Content naming wants byte-reproducible states; snapshot security wants states that differ on every restore. Declared entropy regions are excluded from the content hash and reseeded per branch, so eight siblings of one snapshot share a name and no two share a machine ID. The first implementation excluded a region's bytes but kept its existence in the hash — and a runtime invariant check inside reseed caught that this breaks ReseedPolicy::Remove, which is the correct handling for an SSH host key.

A working cache hides nondeterminism (D24). Run a nondeterministic workload twice and it does not diverge — the second call is a cache hit and never re-executes. Divergence is only observable when the work runs again, which makes sampled shadow re-execution not an optional audit but the only detector of a lying determinism claim in a system whose cache works.

Concurrency

The properties, and where each is tested:

property test
1000 simultaneous rollouts → 1 execution engine/tests/concurrency.rs
A failing leader promotes a follower; the work still runs once same
A cancelled leader does too — the claim is released by Drop same
A slow event subscriber drops events instead of stalling the engine same
Events are totally ordered across subscribers same
A pinned state always resolves, under a concurrent eviction storm store/tests/concurrency.rs
Concurrent inserts of one state store it once and leak no references same
16×CPU readers of a 60-deep chain reconstruct each node exactly once same
A failed materialization wakes its waiters and stays retryable same
Contradictory concurrent records all survive; no lost divergence index/tests/concurrency.rs
Ambiguity is monotone — a diverged name never becomes clean again same
Every core naming at once agrees with one core core/tests/saturation.rs

Two bugs were found this way and are documented in decisions.md rather than quietly fixed: a starvation window where eviction could exhaust a healthy caller's retry budget (D25), and a Mutex on the hit path whose replacement with an RwLock cut p99.9 at 256-way concurrency from 10 ms to 473 µs (D26).

How this was stress tested

Six adversarial agents were run against the finished system in parallel, each with a narrow brief and no permission to touch src/ — they had to report bugs, not quietly fix them, so every finding came with a reproducer.

sweep what it attacked outcome
store + engine races pins, eviction, single-flight, refcounts 19 tests, 4 bugs
index + event bus races lost divergences, ordering, lost wakeups 31 tests, 3 bugs
correctness adversary forging a wrong cache hit 45 tests, 3 bugs
failure injection provider/backend/materialization faults, cancellation 13 tests, 2 bugs (one a rediscovery)
latency round 1 engine hit path to 2048-way concurrency, 30 s soak, open-loop 5 benchmarks
latency round 2 store and naming in isolation 9 benchmarks

Eleven distinct bugs were found and all eleven are fixed, each with its former reproducer left behind as a regression test. No correctness test is #[ignore]d; the 20 ignored tests are the benchmarks, which are excluded from the gate on purpose (see below). The full list is in decisions.md D35–D42; four of them (Delta losing entropy declarations, Delta conflating "memory unchanged" with "memory removed", a leadership claim being released twice, and a check-then-act gap between lookup and admission) were genuine soundness failures that could serve state B under the name of state A, or run one piece of work twice.

Three things the campaign taught that are worth more than the individual fixes:

Drop is where the concurrency bugs concentrated. Three of the ten were destructor bugs: mem::forget suppressing more than intended, a destructor running a second time after an explicit completion path, and cleanup placed after an .await that cancellation never reaches. RAII is what makes this system's concurrency tractable, and it is also the part that runs on paths nobody writes down.

Simultaneity is the easy case. A barrier-released fan-out of 1000 callers produces exactly one execution and looks like proof — every participant takes the same branch, which is the interleaving least likely to find a handoff bug. Staggered callers found two independent duplicate-execution bugs (D40, D42) with one observable symptom. Both stagger sources are worth having in a suite: induced failure, and simply running the tests under load. The second one is free.

Benchmarks must not gate a correctness suite. A test asserting a latency bound or a drift ratio measures the machine, and under cargo test --workspace the machine is busy running the rest of the suite. The soak test failed exactly that way, its p99 buckets swinging 10× because five other test binaries were saturating fourteen cores — which says nothing about this system. All 20 benchmarks are #[ignore]d and run serially via make bench.

A delta type is a claim. Delta silently stopped being able to describe every difference between two StateManifests when fields were added to the manifest and not to it. Nothing in the type system objected, and the store reconstructs from deltas whenever a memoized manifest has been dropped.

What survived is as informative as what broke. The naming layer itself — the part every other claim rests on — did not yield: 1M random encodings, 41k exhaustively enumerated canonicalization sequences checked against a real interpreter under a deliberately non-transitive commutation relation, 300k random fingerprints, and 300k random manifests produced zero collisions.

Running it

cargo test --workspace                 # 378 correctness tests, 0 failures
make bench                             # the 20 benchmarks, serially, in release

cargo run -p trajcache-cli -- measure  # Phase 0 report
cargo run -p trajcache-cli -- --annotations examples/agent-tools.json \
    branch examples/setup.json examples/turn-0.json examples/turn-1.json examples/turn-2.json

cargo run -p trajcache-cli -- doctor   # is this sandbox provider trustworthy?
make adapter                           # the same suite, against a *remote* sandbox

Running it on your own sandboxes

TrajCache does not fork sandboxes, and §11 says it should not: E2B, Daytona and AgentENV already do that well, and the contribution is the naming layer above them. So the whole story is an adapter, and PROVIDERS.md is how to write one.

You implement a backend — create a blank box, run a command, list files, read one, write one — and FsProvider derives the real provider from it, owning the four rules that have nothing to do with sandboxes and everything to do with this project's storage layer: staging references, releasing them on a partial failure, ordering directories before symlinks, and refusing a hostile path. Getting any of those wrong produces a state that is named wrong rather than an error.

Which is why the conformance suite matters more than the trait does:

$ cargo run -p trajcache-cli -- doctor
conformance: local — 14/14 checks passed
  pass  blank_sandbox_is_empty
  pass  empty_directories_survive
  pass  symlinks_survive_including_dangling
  ...

An adapter that is 95% right does not fail loudly. If your restore drops empty directories, the naming layer cannot tell — it hashes the manifest it is handed, that manifest is internally consistent, and the name it produces is a perfectly valid name for a state that is missing a directory. Then it is cached, and every rollout that reaches that name is served the broken state, quickly and confidently, forever. Fourteen checks, each one a property the rest of the system assumes without asking, and each failure names the path that changed rather than printing two hashes.

make adapter runs the same suite against an adapter over an in-process "remote service" with no local filesystem anywhere in it (examples/remote_backend.rs), which is the claim worth making: the seam is not local-shaped. It is also the template to copy.

What is deliberately not here

  • Model-side (KV cache) checkpointing. §9 leaves it out and so does this. A KV artifact is invalidated by a model update, a quantization change, a tokenizer change, or an engine patch release; an environment state survives all of them.
  • A vendor adapter. §11 is explicit: "adapters, not implementations, at the sandbox layer". What is here is the seam, a reference implementation over a real directory, a worked adapter over a service with no local filesystem, and the conformance suite that says whether a third one is correct — but no E2B or Firecracker client, because that is a dependency and a set of credentials rather than a design question. See PROVIDERS.md.
  • Real memory snapshotting. MemoryImage and RestoreClass are modelled and named — filesystem state restores anywhere, memory state restores only within a compatible fleet, and the compatibility check is asymmetric in the right direction — but nothing here drives CRIU.

About

Environment state cache for RL agent training: fork and reuse sandbox state.

Resources

Stars

0 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages