Visual Causal-chain Bookmarking

[MSc Dissertation] Retrieval-weighted credit assignment for reinforcement learning from delayed reward. VCBM concentrates learning signal on the causally relevant decision points in a trajectory instead of spreading it uniformly, via two instantiations: statistical causal discovery with Welford value estimation in a tabular MDP (beating a RUDDER return-redistribution baseline — 95.07% correct-decision rate in a mean of 1,532 vs. RUDDER's 3,327 episodes across 20 seeds), and retrieval-weighted advantage blending for an LLM agent trained with LOOP inside the AppWorld benchmark.

RLDissertation

# features

Key Features

Core technologies and system features.

Causal Bookmark Detection

Automatically identifies causally significant steps via HTTP state-change events (POST/PUT/DELETE/PATCH with status < 400). Mathematically equivalent to full environment state-diff detection under REST semantics — zero accuracy loss.

VCC Credit Redistribution

Loss weighting: L_VCC = -(Σ_{t∈B} L_t^PPO + α·Σ_{t∉B} L_t^PPO) / (|B| + α(T−|B|)). Setting α=1 recovers standard LOOP exactly. α=0.1 used in experiments. Implemented on top of Apple ML Research's LOOP framework with Qwen2.5-1.5B-Instruct + LoRA rank-8.

AppWorld Benchmark

Evaluated on AppWorld (ACL 2024) — a multi-app digital assistant benchmark with 750+ tasks across Spotify, Gmail, Calendar and 10+ APIs. Agent generates Python code to interact with a sandboxed app environment. Tasks require 5–20+ API calls with sparse rewards.

Hardware & Scale

Training on NVIDIA RTX 4080 (15.57 GiB VRAM), 1.55B parameter model (Qwen2.5-1.5B-Instruct), LoRA fine-tuning (9.23M trainable params, 0.59%). vLLM V1 inference engine, FSDP2 single-GPU training. 100-iteration VCBM run completed; scaling to 3.5B on UoM CSF cluster (ticket RITM0104892 approved).

Timing Breakthrough

Root cause analysis revealed rollout collection accounts for 98% of VCBM's overhead vs LOOP (1500s vs 88s per 50 steps). Fix 1 — HTTP bookmark detector added inline to execute_with_bookmark() — projects to reduce get_rollouts from 1500s → ~75–100s, making VCBM ~8% faster than LOOP overall.

# graphs

Performance Graphs

Visualizations of model performance and results across experiments.

Trained on a single RTX 4090 (24 GB VRAM)

Gradient norm and gradient norm mean over training

VCBM vs LOOP phase-by-phase timing comparison

Peak GPU memory per training phase on RTX 4090

avg_return — Learning Curve

Average episode return across 100 training iterations on AppWorld. Starts ~0.10, trends upward with high variance (0.05–0.45). Mean ≈ 0.27 — consistent with sparse-reward long-horizon task difficulty on Qwen2.5-1.5B + LoRA rank-8.

# source

Project Source Code

Explore the primary logical modules.

EXPLORER

patch_appworld_server.py

srcpatch_appworld_server.py

1from pathlib import Path
2
3
4def patch_server():
5    filepath = "appworld-env/lib/python3.12/site-packages/appworld/serve/environment.py"
6    if not Path(filepath).exists():
7        print(f"Error: {filepath} not found. Please make sure appworld-env is set up.")
8        return False
9
10    with open(filepath) as f:
11        content = f.read()
12
13    if "execute_with_bookmark" in content:
14        print("Server is already patched.")
15        return True
16
17    execute_marker = '@app.post("/execute")'
18    idx = content.find(execute_marker)
19    if idx == -1:
20        print("Error: Could not find @app.post('/execute') route in environment.py")
21        return False
22
23    end_marker = 'return {"output": output}'
24    end_idx = content.find(end_marker, idx)
25    if end_idx == -1:
26        print("Error: Could not find the end of execute function in environment.py")
27        return False
28
29    end_of_func = end_idx + len(end_marker)
30
31    new_endpoint = """
32
33@app.post("/execute_with_bookmark")
34async def execute_with_bookmark(task_id: str = Body(...), code: str = Body(...)) -> dict[str, Any]:
35    maybe_raise_exception(task_id)
36    n_requests_before = len(world.requester.request_tracker.requests)
37    output = world.execute(code)
38    new_requests = world.requester.request_tracker.requests[n_requests_before:]
39    is_bookmark = False
40    if "Execution failed." not in output:
41        is_bookmark = any(
42            req.get("method", "").upper() in ("POST", "PUT", "DELETE", "PATCH")
43            for req in new_requests
44        )
45    return {"output": {"output": output, "is_bookmark": is_bookmark}}"""
46
47    patched_content = content[:end_of_func] + new_endpoint + content[end_of_func:]
48
49    with open(filepath, "w") as f:
50        f.write(patched_content)
51
52    print("Successfully patched AppWorld server environment.py!")
53    return True
54
55
56if __name__ == "__main__":
57    patch_server()

# simulation

Live Simulation Output

Simulated console execution.

Outputs

Training Run — Iteration 194–200

$_▋

# repositories

Source Code

GitHub repositories for this project.

Visual Causal-chain Bookmarking Repository

Access the complete source code on GitHub.