Files
SirRoboGarage/PPO_Bot/PPO_Bot.nim
T
SirStone 0b17430735 feat(PPO_Bot): persist Adam optimizer state and round counter across restarts (#35)
Save ACAdamStates (m/v tensors + t counters) as .npy files alongside
network weights in latest/ and checkpoint dirs; save round counter to
round_counter.txt. loadBestAvailable restores both on startup; fresh
start works unchanged when files are absent.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-08-18 15:40:54 +02:00

231 lines
8.3 KiB
Nim

## PPO_Bot — enemy tracker + state vector wired into the game loop.
## Training: trajectory collected per tick, PPO update in background thread.
import std/[os, strformat, strutils, math]
import arraymancer
import tankroyale_botapi
import network
import actions
import training
import weights
import ./enemy_tracker
import ./state_vector
const botJsonPath = currentSourcePath().parentDir / "PPO_Bot.json"
const weightsRoot = currentSourcePath().parentDir / "weights"
type PPOBot = ref object of Bot
tracker: EnemyTracker
buffer: TrajectoryBuffer
prevEnergy: float32 # own energy last tick
prevEnemyE: float32 # enemy energy last tick (from tracker)
lastState: Tensor[float32]
lastAction: Tensor[float32]
lastLogP: float32
lastValue: float32
hasLastTrans: bool
lastActions: BotActions # previous tick's decoded actions (for state vector)
roundRewardSum: float32 # cumulative reward this round (for live display)
roundTicks: int # ticks this round
var ac = initActorCritic()
var gAdamStates: ACAdamStates # persists across rounds
# ── Background training state ─────────────────────────────────────────────────
type
TrainingResult = object
ac: ActorCritic
adamStates: ACAdamStates
metrics: PPOMetrics
TrainingArgs = object
ac: ActorCritic
adamStates: ACAdamStates
buffer: TrajectoryBuffer
lastValue: float32
roundNum: int
weightsRoot: string
var
trainingThread: Thread[TrainingArgs]
resultChan: Channel[TrainingResult]
threadLaunched: bool = false # true while training thread is running
roundCounter: int = 0
proc trainingThreadProc(args: TrainingArgs) {.thread.} =
var localAc = args.ac
var localAdam = args.adamStates
let m = ppoUpdate(localAc, args.buffer, lastValue = args.lastValue, adamStates = localAdam)
saveCheckpoint(localAc, localAdam, args.weightsRoot, args.roundNum)
resultChan.send(TrainingResult(ac: localAc, adamStates: localAdam, metrics: m))
# ── Bot methods ───────────────────────────────────────────────────────────────
method onScannedBot*(bot: PPOBot, e: ScannedBotEvent) =
bot.tracker.update(e.x, e.y, e.direction, e.speed, e.energy)
method onRoundStarted*(bot: PPOBot, e: RoundStartedEvent) =
setAdjustGunForBodyTurn(true)
setAdjustRadarForBodyTurn(true)
setAdjustRadarForGunTurn(true)
bot.tracker = initEnemyTracker()
bot.buffer = initTrajectoryBuffer()
bot.prevEnergy = 0.0'f32
bot.prevEnemyE = 0.0'f32
bot.hasLastTrans = false
bot.roundRewardSum = 0.0'f32
bot.roundTicks = 0
method onRoundEnded*(bot: PPOBot, e: RoundEndedEventForBot) =
inc roundCounter
# Add round-end score bonus to last transition (if any)
let roundReward = computeRoundReward(e.results.totalScore.float32)
if bot.hasLastTrans and bot.buffer.len > 0:
bot.buffer.transitions[^1].reward += roundReward
# Training progress display — one line per round in the UI console
let ticks = bot.buffer.len
if ticks > 0:
var rewardSum = 0.0'f32
for tr in bot.buffer.transitions: rewardSum += tr.reward
let avgR = rewardSum / ticks.float32
let avgRStr = formatFloat(avgR.float, ffDecimal, 3)
printToStdOut(&"R:{roundCounter} ticks:{ticks} avgR:{avgRStr} score:{e.results.totalScore}\n")
echo &"R:{roundCounter} ticks:{ticks} avgR:{avgRStr} score:{e.results.totalScore}"
# Pick up result from previous training thread if available; channel IS the sync
if threadLaunched:
let (avail, trained) = resultChan.tryRecv()
if avail:
ac = trained.ac
gAdamStates = trained.adamStates
threadLaunched = false
let m = trained.metrics
printToStdOut(&" trained R:{roundCounter-1} aLoss:{formatFloat(m.actorLoss.float, ffDecimal, 4)} vLoss:{formatFloat(m.valueLoss.float, ffDecimal, 4)} gNorm:{formatFloat(m.gradNorm.float, ffDecimal, 3)}\n")
echo &" trained R:{roundCounter-1} aLoss:{formatFloat(m.actorLoss.float, ffDecimal, 4)} vLoss:{formatFloat(m.valueLoss.float, ffDecimal, 4)} gNorm:{formatFloat(m.gradNorm.float, ffDecimal, 3)}"
if bot.buffer.len == 0:
bot.hasLastTrans = false
return
# If last thread still running, drop this pass — start fresh with newer data
# ponytail: simple drop; queue if every round must train
if threadLaunched:
bot.buffer.clear()
bot.hasLastTrans = false
return
let args = TrainingArgs(
ac: ac,
adamStates: gAdamStates,
buffer: bot.buffer,
lastValue: 0.0'f32,
roundNum: roundCounter,
weightsRoot: weightsRoot,
)
bot.buffer.clear()
bot.hasLastTrans = false
printToStdOut(&" train→ R:{roundCounter} ticks:{ticks}\n")
echo &" train→ R:{roundCounter} ticks:{ticks}"
createThread(trainingThread, trainingThreadProc, args)
threadLaunched = true
method run(bot: PPOBot) =
# Seed energy on first tick
bot.prevEnergy = getEnergy().float32
bot.prevEnemyE = if bot.tracker.hasContact: bot.tracker.current.energy.float32 else: 0.0'f32
while isRunning():
bot.tracker.deadReckon()
setRadarTurnRate(bot.tracker.getRadarTurnRate(getX(), getY(), getDirection(), getRadarDirection()))
let botData = BotStateData(
x: getX(),
y: getY(),
direction: getDirection(),
speed: getSpeed(),
energy: getEnergy(),
gunDirection: getGunDirection(),
gunHeat: getGunHeat(),
arenaWidth: float64(getArenaWidth()),
arenaHeight: float64(getArenaHeight()),
)
let remainingGotoDistance = hypot(bot.lastActions.gotoX - botData.x,
bot.lastActions.gotoY - botData.y)
let remainingGunAngle = abs(normalizeRelativeAngle(
directionTo(botData.x, botData.y, bot.lastActions.aimToX, bot.lastActions.aimToY) -
botData.gunDirection))
let state = buildStateVector(botData, bot.tracker, remainingGotoDistance, remainingGunAngle)
let (rawActs, logP) = ac.actorForward(state)
let value = ac.criticForward(state)
let acts = mapActions(rawActs,
getGunHeat().float,
botData.arenaWidth, botData.arenaHeight,
botData.x, botData.y,
botData.direction, botData.speed, botData.gunDirection)
# Compute tick reward from energy deltas
let curEnergy = getEnergy().float32
let curEnemyE = if bot.tracker.hasContact: bot.tracker.current.energy.float32 else: bot.prevEnemyE
let myDelta = curEnergy - bot.prevEnergy
let enemyDelta = curEnemyE - bot.prevEnemyE
let tickReward = computeTickReward(myDelta, enemyDelta)
# Track running reward for in-game display
bot.roundRewardSum += tickReward
inc bot.roundTicks
# Finalise previous transition with the reward from this tick's state change
if bot.hasLastTrans:
let tr = Transition(
state: bot.lastState,
action: bot.lastAction,
logProb: bot.lastLogP,
reward: tickReward,
value: bot.lastValue,
)
bot.buffer.add(tr)
# Store current for next tick
bot.lastState = state
bot.lastAction = rawActs
bot.lastLogP = logP
bot.lastValue = value
bot.prevEnergy = curEnergy
bot.prevEnemyE = curEnemyE
bot.hasLastTrans = true
bot.lastActions = acts
setTargetSpeed(acts.targetSpeed.float)
setTurnRate(acts.turnRate.float)
setGunTurnRate(acts.gunTurnRate.float)
if acts.shouldFire:
discard setFire(acts.firePower.float)
# In-game training progress overlay
let avgR = if bot.roundTicks > 0: bot.roundRewardSum / bot.roundTicks.float32
else: 0.0'f32
let avgRStr = formatFloat(avgR.float, ffDecimal, 2)
drawText(&"R:{roundCounter} avg:{avgRStr}", getX(), getY() - 40.0)
go()
when isMainModule:
resultChan.open()
createDir(weightsRoot)
cleanStaleTempDirs(weightsRoot)
let loadResult = loadBestAvailable(ac, gAdamStates, weightsRoot)
if loadResult.loaded:
roundCounter = loadResult.roundNum
var bot = PPOBot(
tracker: initEnemyTracker(),
buffer: initTrajectoryBuffer(),
)
start(bot, botJsonPath)