Skip to content

feat(agentic):Add Battleship agentic RL demo with web UI and autoplay comparison - #431

Open
tal-ent wants to merge 14 commits into
inclusionAI:mainfrom
tal-ent:main
Open

feat(agentic):Add Battleship agentic RL demo with web UI and autoplay comparison#431
tal-ent wants to merge 14 commits into
inclusionAI:mainfrom
tal-ent:main

Conversation

@tal-ent

@tal-ent tal-ent commented Jul 31, 2026

Copy link
Copy Markdown

Type of Change

  • ✨ New Feature

Related Issue

#191

Overview

This experiment implements a multi-turn agentic RL example for the AReno framework — Battleship. The rules: on an 8×8 grid, the model makes sequential decisions via the fire tool call, targeting and sinking a hidden fleet [4, 3, 2, 2]; sinking all ships within 64 turns constitutes a win. Per the Proposed feature requirements, the framework remains unmodified; integration with areno occurs at three接缝 points only: data loading / trajectory collection / reward. Game rules are centralized in game.py as the single source of truth; a single game is modeled as a multi-turn tool conversation. Verification is two-tiered: pure CPU evaluation and baseline evaluation / web UI.

New Files

examples/agentic/battleship/
├── game.py              # Core game logic
├── reward.py            # RL reward function
├── run_agent.py         # areno agent entry point
├── dataset_generator.py # Training/eval JSONL generation
├── dataset_loader.py    # JSONL → areno prompt records
├── evaluate.py          # Offline baseline players (random/fake/heuristic)
├── play_llm.py          # LLM evaluation against any OpenAI-compatible endpoint
├── compare_modes.py     # Head-to-head comparison of multiple strategies on the same fleet sequence
├── web_ui.py            # Browser-based demo
└── README.md            # Rules, quick start, reward shaping, limitations

tests/
├── test_agentic_battleship_example_cpu.py
└── test_compare_modes_cpu.py

Issues Identified During Review

  1. _tool_messages() raises TypeError: GameState not JSON serializable — The 4 error branches in _run_tool stuff the internally bookkept GameState into the "state" key, which is then fed directly to json.dumps.
  2. LLM repeatedly fires at the same cell — The prompt includes a "do not repeat" rule, but the model cannot see which coordinates it has already fired.

Solutions

  1. Strip the state key before serialization in _tool_messages() — Does not adopt the default= serializer approach. Reason: GameState.ships contains all ship positions; serializing it would leak the answer to the model.
  2. Add explicit instructions to the system prompt in run_agent.py — Teach the LLM how to play Battleship, including the rules and strategy.

✅ Self-Check Checklist

  • Follows project CONTRIBUTING.md / AGENTS.md conventions
  • All local pytest tests/ -k cpu tests pass
  • No existing public API broken
  • Code formatted

why and others added 14 commits July 28, 2026 15:57
Adds a focused agentic RL example where a model learns to play Battleship
by calling a single fire(coordinate) tool. Reuses existing AReno agentic
contracts (RolloutSession, AgentTrajectory, RewardRecord, gspo) — no
framework or CLI changes; the demo runs through `areno train` with
--agent-fn / --reward-fn-path / --dataset-loader-fn.

- 8x8 board, fleet [4,3,2,2]=11 cells, 64-shot cap; fire returns miss/hit/sunk
  without leaking hidden cells; rejects repeated/out-of-range shots
- Seeded, reproducible fleet placement via dataset_generator/dataset_loader
- Shaped reward (win bonus + per-hit/per-sunk shaping - invalid-shot and
  efficiency penalties) for a learnable gradient before consistent wins
- run_agent loops the fire tool until win or turn cap (shopping-example shape)
- evaluate.py baseline harness with random/fake players, reporting
  completion rate and shots-to-win against seeded fleets
- CPU test suite covers generation, fire semantics, sink detection,
  deterministic replay, invalid-input rejection, terminal states, reward
  boundary paths, loader import behavior, and eval orchestration

Fixes a sunk-ship undercount in score_episode: keying a set on ship length
deduped the two length-2 ships. Now counts is_sunk ships directly.

Co-Authored-By: Claude <noreply@anthropic.com>
Interactive 8x8 browser board that reuses game.py and mirrors the
run_agent.py FIRE_TOOL schema/SYSTEM_PROMPT, so the on-screen agent is
the same as the training agent. Supports click-to-fire, single agent
shot, auto-play, seed replay, and LLM (via --base-url) / heuristic
(hunt-target, no server) modes. Payload exposes only hit/miss/unknown,
never unrevealed ship cells.

Adds 4 CPU tests (heuristic wins, payload hides ships, LLM mode refuses
without --base-url, human fire + invalid rejection) and a README section.

Co-Authored-By: Claude <noreply@anthropic.com>
…ON encode

_run_tool's error paths embed the live GameState under a "state" key, which
_tool_messages then passed to json.dumps — crashing with TypeError since
GameState/Ship are not JSON-serializable. Strip the internal "state" key before
encoding so the tool message only carries serializable status metadata. This
also prevents leaking hidden ship positions to the model (board_text already
hides unshot cells). Adds a regression test covering the error path.

Co-Authored-By: Claude <noreply@anthropic.com>
…ship agent

When the model emits no legal fire tool call (e.g. untrained Qwen3-0.6B),
_run_tool's four error paths returned early without calling game.fire(),
so state.shots_used never incremented and run_one's while loop spun forever.
Mark such turns as wasted: bump shots_used and return an invalid result
without touching the board, so the loop terminates within MAX_TURNS.

Adds a regression test and the accumulated Battleship example/test changes
(CLADE.md intentionally excluded).

Co-Authored-By: Claude <noreply@anthropic.com>
…eState attr in test

reward_fn passed each {"name":...,"arguments":...} dict straight to
game.fire as the coordinate, which marked every shot invalid without
ever touching the board. The reward still came out right only because
score_episode rebuilds a fresh state from the call list, so the loop was
dead code, wasted work, and a trap for future refactors. Remove it and
let score_episode own the replay. Also fix the regression test that
referenced GameState.shots (no such field) -> len(shots_history), which
only passed because the run_agent tests silently return on missing torch.

Co-Authored-By: Claude <noreply@anthropic.com>
Head-to-head comparison of heuristic/random/(optional) llm modes over the
same seeded fleets so the only variable is the decision policy. Reuses
evaluate, game, and play_llm._play_game; prints a side-by-side aggregate
table and optionally writes JSON. LLM mode is opt-in via --base-url.

CPU test covers offline schema/seed-fairness, llm-drop-without-base-url,
sample-std parity, llm-only summary extras, and a stubbed-endpoint LLM
game attaching latency/token totals.

Co-Authored-By: Claude <noreply@anthropic.com>
Append an explicit "Already fired: A1, B3, ..." line to the per-turn
board observation so untrained small models don't have to parse o/X
symbols to avoid repeats. Applied to all three LLM prompt sites
(run_agent training rollout, play_llm eval, web_ui LLM mode) plus a
one-line SYSTEM_PROMPT hint in each, keeping train/eval observation
format consistent to avoid distribution skew. No game-logic or game.py
change; the reward already penalizes repeats, so this is a
sample-efficiency scaffold, not a new learning signal.

Adds a CPU test (test_battleship_play_llm_observation_includes_fired_list)
locking the fired-list behavior via the injectable-step harness.

Co-Authored-By: Claude <noreply@anthropic.com>
Co-Authored-By: Claude <noreply@anthropic.com>
…ators

play_llm reported population std (statistics.pstdev, n) while evaluate
and compare_modes reported sample std (n-1), so completion_std was not
comparable across the three evaluators. Switch to statistics.stdev and
add a regression test asserting all three now agree.

Co-Authored-By: Claude <noreply@anthropic.com>
…luate hint

Drop unused occupied/fleet_cells variables and unused imports from
game.py (forbidden already covers occupancy), hoist import json to
module top, and sync the place_fleet docstring. Modernize evaluate's
max_turns hint to int | None. No behavior change.

Co-Authored-By: Claude <noreply@anthropic.com>
…differ from rendered assistant

_append_sample_response merged multi-turn rows by splitting at the first
diverging token (_common_prefix_len). Each turn's prompt re-renders the prior
assistant message from tokenizer.decode(response_tokens), which is not
token-idempotent with the raw engine response_tokens for real BPE tokenizers,
so the divergence landed inside an early assistant span and the re-rendered
prior context was re-appended on every turn. The training row grew O(N**2),
blew past max_context_len, and _filter_overlong_agent_samples dropped every
trajectory -> "all agent trajectories exceeded the configured context length".

Split at len(existing.token_row) instead: the previous trajectory is the
intended prefix and the suffix is the bounded per-turn delta (new tool/board
context + new response). Row stays O(N); behavior is unchanged when raw
response_tokens equal the re-rendered assistant content. Remove the now-orphaned
_common_prefix_len in agentic.py (the dpo.py copy is separate).

Add a CPU regression test that constructs turns whose re-rendered assistant
tokens differ from the raw response tokens and asserts the merged row is the
linear trajectory, not the quadratic bloat.

Co-Authored-By: Claude <noreply@anthropic.com>
Each Battleship turn accumulates conversation context, so a 40-turn episode
produces a ~7.5K-token trajectory. On a single 14.5GB GPU (tp_size=1) the
train-step selected-logprob kernel materializes full-vocab float32 tensors per
chunk and OOMs on a trajectory that long, while lowering --max-context-len
below the natural length just re-triggers the trajectory-filter crash.

12 turns keeps a trajectory ~2.7K tokens (under a 4096 context cap and clear of
the logprob memory peak) and also shortens agentic rollout wall time.

Co-Authored-By: Claude <noreply@anthropic.com>
…cked sequences

The vocab-parallel selected-logprob autograd path materialized full-vocab
float tensors on the training hot path: forward did `logits = logits_shard.float()`
and `exp_logits = torch.exp(...)` (two `[positions, vocab]` float32 tensors) and
saved a third full `[positions, vocab]` probs tensor for backward. With a large
vocabulary (~231K) and a multi-turn packed training row, the per-chunk peak plus
the probs retained across every position chunk for backward exceeded small-GPU
memory and OOM'd at `logits = logits_shard.float()`.

Rewrite the autograd Function so forward reduces row maxima, vocab-chunked
exp-sums, and target logits (never materializing full float logits/probs), and
backward recomputes the per-vocab-chunk softmax to form the local shard gradient
`grad * (onehot - softmax)`. Peak memory is now `positions * vocab_chunk_size`
float instead of `positions * full_vocab`, on both forward and backward.

Forward equivalence and backward correctness are covered by existing CPU tests
plus a new chunked-backward test that forces multiple vocab chunks and checks
the gradient against full-vocab cross entropy.

Co-Authored-By: Claude <noreply@anthropic.com>
…under context limit

With forced tool_choice=fire, small untrained models still emit long assistant
content before the fire call, and sampling variance made per-turn response
length swing ~140..290 tokens run-to-run. An 11-turn trajectory could reach
6.1K tokens, exceed max_context_len=4096, and get filtered -- crashing rollout
with "all agent trajectories exceeded the configured context length" -- or
 OOM the train-step logprobs.

Pass max_tokens=128 to each chat completion (honoredby the proxy and mapped to
params.max_new_tokens). A fire tool call is ~30 tokens, so 128 leaves room for
brief content without letting a turn ramble; 11 turns now stay ~3.3K tokens
worst case, safely under 4096 and clear of the logprob memory peak.

Co-Authored-By: Claude <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant