従来の落ち物パズルは、ランダムに供給されるブロックをいかに効率的に平らに積み上げて消すかという「個人の操作精度・平積み技術」が重視されるゲームでした。

しかし、もし「相手に落とすブロックをこちらが指名できる」としたらどうなるでしょうか？

この記事では、戦略的ブロック選択システムを取り入れたWeb対戦パズルゲーム「Choices Battle Blocks」の開発において、JAXを活用した超高速シミュレーション環境とCEM（Cross-Entropy Method）による強力な思考AIの学習、そしてブラウザ上でのONNX/P2P対戦の実装秘話をまとめます。

1. 「ブロック選択」がもたらすゲーム性の劇的変化

本作の基本ルールは、提示された4つの選択肢の中から「自分のブロック」と「相手のブロック」を互いに指定し合う点にあります。

このルール下では、従来のテトリスAIのような「定型的な置き方」だけでは勝てず、相手の隙を突く高度な先読みAIが必要となります。

2. JAXによる超高速並列シミュレーション環境

強化学習や進化的最適化でAIを訓練する場合、最大のボトルネックは「ゲームのシミュレーション速度」です。Pythonの通常のループでは1秒間に数百手程度しか回せません。

そこでゲーム盤面（20×10のビットマップ）、ブロックの落下・回転・ライン消去判定をすべてJAXの純粋関数（Pure Function）として記述し、jax.vmap と jax.jit でXLAコンパイルしました。

import jax
import jax.numpy as jnp

@jax.jit
def step_game_state(state, action):
    # state: (batch_size, 20, 10) の盤面テンソル
    # 1万ゲームをGPUメモリ上で完全に並列実行
    new_board, reward, done = simulate_placement(state.board, action)
    return state.replace(board=new_board, score=state.score + reward)

これにより、単一のGPU上で1秒間に100万手以上の盤面シミュレーションが可能になり、数時間で数千世代の進化学習を回せるようになりました。

3. CEM（Cross-Entropy Method）による評価関数の最適化

AIの思考エンジンには、各配置後の盤面状態をスコアリングする評価関数（特徴量：最大高さ、穴の数、窪みの深さ、平坦度、相手への妨害度など）を採用しました。

これらの重みパラメータベクトルをCEM（交叉エントロピー法）によって最適化：

1. 重みの多変量正規分布から候補エージェント集団（100個体）をサンプリング。
2. JAX並列シミュレータ上で総当たり戦やサバイバルテストを実施。
3. 成績上位の「エリート個体」の平均・分散を計算し、次のサンプリング分布を更新。

このループを繰り返すことで、人間でも容易には崩せない、粘り強く意地悪な「Lv11+（1手先読みCEM強化AI）」が誕生しました。

4. ブラウザでのONNX推論 & PeerJSによる完全P2P対戦

Web版ではサーバーのインフラ費用をかけずに誰もが気軽に遊べるよう、以下の技術を採用しました：

• ONNX Runtime Web: 最適化されたモデル重みと先読みロジックをJavaScript Worker内で軽量に推論実行（ラグなし）。
• PeerJS (WebRTC DataChannel): プレイヤー間の通信は完全なP2P。中央サーバーを介さず直接通信するため、超低遅延で対戦が成立します。
• PWA対応: 端末にインストールすれば、スマホのテザリング機能を使って完全オフラインのLAN対戦も可能です。

5. おわりに

JAXとWebAssembly/WebRTCの組み合わせは、個人開発のインディゲームでも「高度なAI」と「リアルタイム通信対戦」をゼロコスト・高パフォーマンスで提供できる非常に強力な武器です。

Choices Battle BlocksはWeb上で今すぐプレイできますので、ぜひAIや友達との熱い駆け引きをお楽しみください！