Merge pull request #247 from CSSLab/codex/maia-winrate-play

Codex/maia winrate play

Author: ashtonanderson

README.md

@@ -187,6 +187,24 @@ The project uses automated code formatting and linting to maintain consistency:
 
 This ensures consistent interpretation where positive values always represent an advantage for the player whose turn it is. The conversion logic is documented in `src/hooks/useStockfishEngine/engine.ts:138-155`.
 
+#### Maia Value-Head Debug Logging
+
+When testing the value-head play mode (`/play/maia` with `maiaMoveSelectionMode=value_head`), you can enable an opt-in console table that prints each legal move with Maia's evaluated win probability.
+
+Enable in browser DevTools:
+
+```js
+localStorage.setItem('maia.valueHeadDebug', 'true')
+```
+
+Disable:
+
+```js
+localStorage.removeItem('maia.valueHeadDebug')
+```
+
+With the flag enabled, each Maia move prints a compact table (`san`, `move`, `maiaWinProb`) in the browser console. This is local-only and does not change behavior unless the flag is set in localStorage.
+
 #### State Management Architecture
 
 The platform uses a Context + Custom Hooks pattern:

src/components/Common/PlaySetupModal.tsx

@@ -7,6 +7,7 @@ import { useCallback, useContext, useState } from 'react'
 
 import {
   Color,
+  MaiaMoveSelectionMode,
   PlayType,
   TimeControl,
   TimeControlOptionNames,
@@ -80,6 +81,7 @@ interface Props {
   maiaVersion?: string
   isBrain?: boolean
   sampleMoves?: boolean
+  maiaMoveSelectionMode?: MaiaMoveSelectionMode
   simulateMaiaTime?: boolean
   startFen?: string
 }
@@ -105,6 +107,10 @@ export const PlaySetupModal: React.FC<Props> = (props: Props) => {
   const [sampleMoves, setSampleMoves] = useState<boolean>(
     props.sampleMoves || true,
   )
+  const [maiaMoveSelectionMode, setMaiaMoveSelectionMode] =
+    useState<MaiaMoveSelectionMode>(
+      props.maiaMoveSelectionMode || 'move_matching',
+    )
   const [simulateMaiaTime, setSimulateMaiaTime] = useState<boolean>(
     props.simulateMaiaTime !== undefined ? props.simulateMaiaTime : true,
   )
@@ -174,6 +180,7 @@ export const PlaySetupModal: React.FC<Props> = (props: Props) => {
             maiaVersion: maiaVersion,
             timeControl: timeControl,
             sampleMoves: sampleMoves,
+            maiaMoveSelectionMode: maiaMoveSelectionMode,
             simulateMaiaTime: simulateMaiaTime,
             startFen: fen,
           },
@@ -188,6 +195,7 @@ export const PlaySetupModal: React.FC<Props> = (props: Props) => {
             timeControl: timeControl,
             isBrain: isBrain,
             sampleMoves: sampleMoves,
+            maiaMoveSelectionMode: maiaMoveSelectionMode,
             simulateMaiaTime: simulateMaiaTime,
             startFen: fen,
           },
@@ -202,6 +210,7 @@ export const PlaySetupModal: React.FC<Props> = (props: Props) => {
       maiaVersion,
       timeControl,
       sampleMoves,
+      maiaMoveSelectionMode,
       simulateMaiaTime,
       fen,
       isBrain,
@@ -455,6 +464,35 @@ export const PlaySetupModal: React.FC<Props> = (props: Props) => {
                 </div>
               </div>
 
+              {props.playType == 'againstMaia' ? (
+                <div>
+                  <label
+                    htmlFor="maia-play-style-select"
+                    className="mb-1 block text-sm font-medium text-primary"
+                  >
+                    Maia play style:
+                  </label>
+                  <div id="maia-play-style-select">
+                    <OptionSelect
+                      options={['move_matching', 'value_head']}
+                      labels={['Human move-match', 'Best win rate']}
+                      selected={maiaMoveSelectionMode}
+                      onChange={(selected) =>
+                        setMaiaMoveSelectionMode(
+                          selected as MaiaMoveSelectionMode,
+                        )
+                      }
+                      selectedClassName="border-human-4 bg-human-4 text-white hover:bg-human-4/90"
+                    />
+                  </div>
+                  <p className="mt-2 text-xs text-secondary">
+                    {maiaMoveSelectionMode === 'move_matching'
+                      ? 'Classic Maia: chooses the move it expects a human to play.'
+                      : 'Local Maia 2 mode: evaluates every legal move and picks the one with the best perceived win rate.'}
+                  </p>
+                </div>
+              ) : null}
+
               <div className="flex items-center gap-2">
                 <input
                   type="checkbox"

src/hooks/usePlayController/useVsMaiaController.ts

@@ -1,4 +1,4 @@
-import { useEffect } from 'react'
+import { useCallback, useContext, useEffect } from 'react'
 import { Chess } from 'chess.ts'
 import { PlayGameConfig } from 'src/types'
 import { backOff } from 'exponential-backoff'
@@ -7,6 +7,26 @@ import { usePlayController } from 'src/hooks/usePlayController'
 import { fetchGameMove, logGameMove, fetchPlayPlayerStats } from 'src/api'
 import { useSound } from 'src/hooks/useSound'
 import { safeUpdateRating } from 'src/lib/ratingUtils'
+import { MaiaEngineContext } from 'src/contexts'
+
+const MAIA_VALUE_HEAD_DEBUG_KEY = 'maia.valueHeadDebug'
+
+const isTruthy = (value: string | null | undefined): boolean => {
+  if (!value) return false
+  return ['1', 'true', 'yes', 'on'].includes(value.toLowerCase())
+}
+
+const isMaiaValueHeadDebugEnabled = (): boolean => {
+  if (typeof window === 'undefined') {
+    return false
+  }
+
+  try {
+    return isTruthy(window.localStorage.getItem(MAIA_VALUE_HEAD_DEBUG_KEY))
+  } catch {
+    return false
+  }
+}
 
 const playStatsLoader = async () => {
   const stats = await fetchPlayPlayerStats()
@@ -25,6 +45,126 @@ export const useVsMaiaPlayController = (
   const controller = usePlayController(id, playGameConfig)
   const [stats, incrementStats, updateRating] = useStats(playStatsLoader)
   const { playMoveSound } = useSound()
+  const maiaEngine = useContext(MaiaEngineContext)
+  const valueHeadModel =
+    playGameConfig.maiaMoveSelectionMode === 'value_head'
+      ? maiaEngine.maia
+      : null
+  const valueHeadStatus =
+    playGameConfig.maiaMoveSelectionMode === 'value_head'
+      ? maiaEngine.status
+      : null
+  const playerRatingForValueHead = stats.rating ?? 1500
+  const valueHeadDebugEnabled = isMaiaValueHeadDebugEnabled()
+
+  const selectValueHeadMove = useCallback(async () => {
+    if (
+      !controller.currentNode ||
+      valueHeadStatus !== 'ready' ||
+      !valueHeadModel
+    ) {
+      if (valueHeadDebugEnabled) {
+        console.log('[Maia value-head debug] selector unavailable', {
+          hasCurrentNode: !!controller.currentNode,
+          valueHeadStatus,
+          hasModel: !!valueHeadModel,
+        })
+      }
+      return null
+    }
+
+    const currentFen = controller.currentNode.fen
+    const chess = new Chess(currentFen)
+    const legalMoves = chess.moves({ verbose: true }) as Array<{
+      from: string
+      to: string
+      promotion?: string
+    }>
+
+    if (legalMoves.length === 0) {
+      return null
+    }
+
+    const candidateMoves = legalMoves.map(
+      (move) => `${move.from}${move.to}${move.promotion ?? ''}`,
+    )
+    const candidateBoards = candidateMoves.map((moveUci) => {
+      const board = new Chess(currentFen)
+      board.move(moveUci, { sloppy: true })
+      return board.fen()
+    })
+
+    const maiaRating = parseInt(
+      playGameConfig.maiaVersion.replace('maia_kdd_', ''),
+      10,
+    )
+    const modelElo = Number.isNaN(maiaRating) ? 1500 : maiaRating
+
+    // After Maia makes a candidate move, it is the human's turn.
+    // Maia's value head conditions on the side to move as elo_self,
+    // so the resulting boards must use the human as elo_self and Maia as elo_oppo.
+    const { result } = await valueHeadModel.batchEvaluate(
+      candidateBoards,
+      Array(candidateBoards.length).fill(playerRatingForValueHead),
+      Array(candidateBoards.length).fill(modelElo),
+    )
+
+    const maiaIsWhite = controller.player === 'black'
+    let bestMove = candidateMoves[0]
+    let bestScore = maiaIsWhite ? result[0].value : 1 - result[0].value
+
+    for (let index = 1; index < candidateMoves.length; index++) {
+      const whiteWinProb = result[index].value
+      const maiaWinProb = maiaIsWhite ? whiteWinProb : 1 - whiteWinProb
+
+      if (maiaWinProb > bestScore) {
+        bestMove = candidateMoves[index]
+        bestScore = maiaWinProb
+      }
+    }
+
+    if (valueHeadDebugEnabled) {
+      const candidateSummaries = candidateMoves
+        .map((moveUci, index) => {
+          const moveResult = result[index]
+          const whiteWinProb = moveResult.value
+          const maiaWinProb = maiaIsWhite ? whiteWinProb : 1 - whiteWinProb
+          const board = new Chess(currentFen)
+          const moveObj = board.move(moveUci, { sloppy: true })
+
+          return {
+            move: moveUci,
+            san: moveObj?.san ?? moveUci,
+            maiaWinProb: Number(maiaWinProb.toFixed(4)),
+          }
+        })
+        .sort((a, b) => b.maiaWinProb - a.maiaWinProb)
+
+      console.groupCollapsed(
+        `[Maia value-head debug] ${playGameConfig.maiaVersion} selected ${bestMove}`,
+      )
+      console.table(candidateSummaries)
+      console.groupEnd()
+    }
+
+    const estimatedDelaySeconds = Math.min(
+      3,
+      0.35 + legalMoves.length * 0.04 + Math.random() * 0.25,
+    )
+
+    return {
+      top_move: bestMove,
+      move_delay: estimatedDelaySeconds,
+    }
+  }, [
+    controller.currentNode,
+    controller.player,
+    playGameConfig.maiaVersion,
+    playerRatingForValueHead,
+    valueHeadModel,
+    valueHeadStatus,
+    valueHeadDebugEnabled,
+  ])
 
   const makePlayerMove = async (moveUci: string) => {
     const moveTime = controller.updateClock()
@@ -48,20 +188,28 @@ export const useVsMaiaPlayController = (
           ? parseInt(controller.timeControl.split('+')[0]) * 60
           : 0
 
-        const maiaMoves = await backOff(
-          () =>
-            fetchGameMove(
-              controller.moveList,
-              playGameConfig.maiaVersion,
-              playGameConfig.startFen,
-              null,
-              simulateMaiaTime ? initialClock : 0,
-              simulateMaiaTime ? maiaClock : 0,
-            ),
-          {
-            jitter: 'full',
-          },
-        )
+        const maiaMoves =
+          playGameConfig.maiaMoveSelectionMode === 'value_head'
+            ? await selectValueHeadMove()
+            : await backOff(
+                () =>
+                  fetchGameMove(
+                    controller.moveList,
+                    playGameConfig.maiaVersion,
+                    playGameConfig.startFen,
+                    null,
+                    simulateMaiaTime ? initialClock : 0,
+                    simulateMaiaTime ? maiaClock : 0,
+                  ),
+                {
+                  jitter: 'full',
+                },
+              )
+
+        if (!maiaMoves?.top_move) {
+          return
+        }
+
         const nextMove = maiaMoves['top_move']
         const moveDelay = maiaMoves['move_delay']
 
@@ -107,8 +255,11 @@ export const useVsMaiaPlayController = (
     controller.game.termination,
     controller.moveList.length,
     playGameConfig.maiaVersion,
+    playGameConfig.maiaMoveSelectionMode,
     playGameConfig.startFen,
     simulateMaiaTime,
+    selectValueHeadMove,
+    valueHeadStatus,
   ])
 
   useEffect(() => {