WIP: add value-head Maia play mode

ashtonanderson · ashtonanderson · commit 0a05bd3d6099 · 2026-03-03T11:58:38.000-05:00
diff --git a/src/components/Common/PlaySetupModal.tsx b/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,
   )
@@ -173,6 +179,7 @@ export const PlaySetupModal: React.FC<Props> = (props: Props) => {
             maiaVersion: maiaVersion,
             timeControl: timeControl,
             sampleMoves: sampleMoves,
+            maiaMoveSelectionMode: maiaMoveSelectionMode,
             simulateMaiaTime: simulateMaiaTime,
             startFen: fen,
           },
@@ -187,6 +194,7 @@ export const PlaySetupModal: React.FC<Props> = (props: Props) => {
             timeControl: timeControl,
             isBrain: isBrain,
             sampleMoves: sampleMoves,
+            maiaMoveSelectionMode: maiaMoveSelectionMode,
             simulateMaiaTime: simulateMaiaTime,
             startFen: fen,
           },
@@ -201,6 +209,7 @@ export const PlaySetupModal: React.FC<Props> = (props: Props) => {
       maiaVersion,
       timeControl,
       sampleMoves,
+      maiaMoveSelectionMode,
       simulateMaiaTime,
       fen,
       isBrain,
@@ -410,6 +419,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"
diff --git a/src/hooks/usePlayController/useVsMaiaController.ts b/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,7 @@ 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 playStatsLoader = async () => {
   const stats = await fetchPlayPlayerStats()
@@ -25,6 +26,87 @@ 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 selectValueHeadMove = useCallback(async () => {
+    if (
+      !controller.currentNode ||
+      valueHeadStatus !== 'ready' ||
+      !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
+    const { result } = await valueHeadModel.batchEvaluate(
+      candidateBoards,
+      Array(candidateBoards.length).fill(modelElo),
+      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
+      }
+    }
+
+    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,
+    valueHeadModel,
+    valueHeadStatus,
+  ])
 
   const makePlayerMove = async (moveUci: string) => {
     const moveTime = controller.updateClock()
@@ -48,20 +130,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 +197,11 @@ export const useVsMaiaPlayController = (
     controller.game.termination,
     controller.moveList.length,
     playGameConfig.maiaVersion,
+    playGameConfig.maiaMoveSelectionMode,
     playGameConfig.startFen,
     simulateMaiaTime,
+    selectValueHeadMove,
+    valueHeadStatus,
   ])
 
   useEffect(() => {
diff --git a/src/pages/play/maia.tsx b/src/pages/play/maia.tsx
@@ -3,9 +3,20 @@ import { startGame } from 'src/api'
 import { NextPage } from 'next/types'
 import { useRouter } from 'next/router'
 import { tourConfigs } from 'src/constants/tours'
-import { ModalContext, useTour } from 'src/contexts'
-import { Loading, PlayControls } from 'src/components'
-import { Color, TimeControl, PlayGameConfig } from 'src/types'
+import {
+  MaiaEngineContext,
+  MaiaEngineContextProvider,
+  ModalContext,
+  useTour,
+} from 'src/contexts'
+import { DownloadModelModal, Loading, PlayControls } from 'src/components'
+import {
+  Color,
+  MaiaEngine,
+  MaiaMoveSelectionMode,
+  TimeControl,
+  PlayGameConfig,
+} from 'src/types'
 import { useContext, useEffect, useMemo, useState } from 'react'
 import { GameplayInterface } from 'src/components/Board/GameplayInterface'
 import { useVsMaiaPlayController } from 'src/hooks/usePlayController/useVsMaiaController'
@@ -74,7 +85,13 @@ const PlayMaia: React.FC<Props> = ({
   )
 }
 
-const PlayMaiaPage: NextPage = () => {
+interface PageContentProps {
+  maiaEngine?: MaiaEngine
+}
+
+const PlayMaiaPageContent: React.FC<PageContentProps> = ({
+  maiaEngine,
+}: PageContentProps) => {
   const { startTour, tourState } = useTour()
   const [initialTourCheck, setInitialTourCheck] = useState(false)
 
@@ -96,6 +113,7 @@ const PlayMaiaPage: NextPage = () => {
     timeControl,
     isBrain,
     sampleMoves,
+    maiaMoveSelectionMode,
     simulateMaiaTime: simulateMaiaTimeQuery,
     startFen,
   } = router.query
@@ -114,13 +132,16 @@ const PlayMaiaPage: NextPage = () => {
       timeControl: (timeControl || 'unlimited') as TimeControl,
       isBrain: isBrain == 'true',
       sampleMoves: sampleMoves == 'true',
+      maiaMoveSelectionMode: (maiaMoveSelectionMode ||
+        'move_matching') as MaiaMoveSelectionMode,
       simulateMaiaTime: simulateMaiaTime,
       startFen: typeof startFen == 'string' ? startFen : undefined,
     }),
     [
       startFen,
       isBrain,
       maiaVersion,
+      maiaMoveSelectionMode,
       player,
       sampleMoves,
       timeControl,
@@ -177,6 +198,10 @@ const PlayMaiaPage: NextPage = () => {
       }
     }
 
+    if (!router.isReady) {
+      return
+    }
+
     if (!id) {
       fetchGameId()
 
@@ -195,6 +220,15 @@ const PlayMaiaPage: NextPage = () => {
           content="Challenge the most human-like chess AI. Unlike traditional engines that play robotically, Maia naturally plays moves a person would make, trained on millions of human games with real chess intuition."
         />
       </Head>
+      {playGameConfig.maiaMoveSelectionMode === 'value_head' &&
+      maiaEngine &&
+      (maiaEngine.status === 'no-cache' ||
+        maiaEngine.status === 'downloading') ? (
+        <DownloadModelModal
+          progress={maiaEngine.progress}
+          download={maiaEngine.downloadModel}
+        />
+      ) : null}
       <Loading isLoading={!router.isReady || !id}>
         {router.isReady && id && (
           <PlayMaia
@@ -210,4 +244,25 @@ const PlayMaiaPage: NextPage = () => {
   )
 }
 
+const PlayMaiaPageWithLocalMaia: React.FC = () => {
+  const maiaEngine = useContext(MaiaEngineContext)
+
+  return <PlayMaiaPageContent maiaEngine={maiaEngine} />
+}
+
+const PlayMaiaPage: NextPage = () => {
+  const router = useRouter()
+  const requiresLocalMaia = router.query.maiaMoveSelectionMode === 'value_head'
+
+  if (!requiresLocalMaia) {
+    return <PlayMaiaPageContent />
+  }
+
+  return (
+    <MaiaEngineContextProvider>
+      <PlayMaiaPageWithLocalMaia />
+    </MaiaEngineContextProvider>
+  )
+}
+
 export default PlayMaiaPage
diff --git a/src/types/play.ts b/src/types/play.ts
@@ -10,9 +10,12 @@ export const TimeControlOptionNames = [
   'Unlimited',
 ]
 
+export const MaiaMoveSelectionModes = ['move_matching', 'value_head'] as const
+
 // To fix some weird inconsistency between vercel and local:
 // eslint-disable-next-line prettier/prettier
 export type TimeControl = (typeof TimeControlOptions)[number]
+export type MaiaMoveSelectionMode = (typeof MaiaMoveSelectionModes)[number]
 
 export type PlayType = 'againstMaia' | 'handAndBrain'
 
@@ -23,6 +26,7 @@ export interface PlayGameConfig {
   playType: PlayType
   isBrain: boolean
   sampleMoves: boolean
+  maiaMoveSelectionMode?: MaiaMoveSelectionMode
   simulateMaiaTime?: boolean
   startFen?: string
   maiaPartnerVersion?: string
