Repository: sotetsuk/pgx

Issue:
[Shogi] Implement `observe`
Follow dlshogi's feature:

https://github.com/TadaoYamaoka/DeepLearningShogi/blob/1053b87e73c5eb69b66c9937c984d02afe8d1a3b/cppshogi/cppshogi.cpp#L39-L123


Discussion:

Resolving patch:
diff --git a/pgx/shogi.py b/pgx/shogi.py
--- a/pgx/shogi.py
+++ b/pgx/shogi.py
@@ -312,8 +312,73 @@ def step(state: State, action: jnp.ndarray) -> State:
 
 
 def observe(state: State, player_id: jnp.ndarray) -> jnp.ndarray:
-    # TODO: write me
-    return jnp.zeros(100)
+    state = jax.lax.cond(
+        state.curr_player != player_id, lambda: _flip(state), lambda: state
+    )
+
+    def piece_and_effect(state):
+        # 駒の場所
+        my_pieces = jnp.arange(OPP_PAWN)
+        my_piece_feat = jax.vmap(lambda p: state.piece_board == p)(my_pieces)
+        # 自分の利き
+        my_effect = state.effects[0]
+
+        def e(p):
+            mask = state.piece_board == p
+            return jnp.where(mask.reshape(81, 1), my_effect, FALSE).any(axis=0)
+
+        my_effect_feat = jax.vmap(e)(my_pieces)
+        # 利きの枚数
+        my_effect_sum = my_effect.sum(axis=0)
+
+        def effect_sum(n) -> jnp.ndarray:
+            return my_effect_sum >= n  # type: ignore
+
+        effect_sum_feat = jax.vmap(effect_sum)(jnp.arange(1, 4))
+        return my_piece_feat, my_effect_feat, effect_sum_feat
+
+    my_piece_feat, my_effect_feat, my_effect_sum_feat = piece_and_effect(state)
+    opp_piece_feat, opp_effect_feat, opp_effect_sum_feat = piece_and_effect(
+        _flip(state)
+    )
+    opp_piece_feat = opp_piece_feat[:, ::-1]
+    opp_effect_feat = opp_effect_feat[:, ::-1]
+    opp_effect_sum_feat = opp_effect_sum_feat[:, ::-1]
+
+    def num_hand(n, hand, p):
+        return jnp.tile(hand[p] >= n, reps=(9, 9))
+
+    def hand_feat(hand):
+        # fmt: off
+        pawn_feat = jax.vmap(partial(num_hand, hand=hand, p=PAWN))(jnp.arange(1, 9))
+        lance_feat = jax.vmap(partial(num_hand, hand=hand, p=LANCE))(jnp.arange(1, 5))
+        knight_feat = jax.vmap(partial(num_hand, hand=hand, p=KNIGHT))(jnp.arange(1, 5))
+        silver_feat = jax.vmap(partial(num_hand, hand=hand, p=SILVER))(jnp.arange(1, 5))
+        gold_feat = jax.vmap(partial(num_hand, hand=hand, p=GOLD))(jnp.arange(1, 5))
+        bishop_feat = jax.vmap(partial(num_hand, hand=hand, p=BISHOP))(jnp.arange(1, 3))
+        rook_feat = jax.vmap(partial(num_hand, hand=hand, p=ROOK))(jnp.arange(1, 3))
+        return [pawn_feat, lance_feat, knight_feat, silver_feat, gold_feat, bishop_feat, rook_feat]
+        # fmt: on
+
+    my_hand_feat = hand_feat(state.hand[0])
+    opp_hand_feat = hand_feat(state.hand[1])
+
+    checking_point_board, _ = _check_info(
+        state, _flip(state), state.effects[1]
+    )
+    checked = jnp.tile(checking_point_board.any(), reps=(1, 9, 9))
+
+    feat1 = [
+        my_piece_feat.reshape(14, 9, 9),
+        my_effect_feat.reshape(14, 9, 9),
+        my_effect_sum_feat.reshape(3, 9, 9),
+        opp_piece_feat.reshape(14, 9, 9),
+        opp_effect_feat.reshape(14, 9, 9),
+        opp_effect_sum_feat.reshape(3, 9, 9),
+    ]
+    feat2 = my_hand_feat + opp_hand_feat + [checked]
+    feat = jnp.vstack(feat1 + feat2)
+    return feat
 
 
 def _step(state: State, action: Action) -> State:


Test patch:
diff --git a/tests/assets/shogi/observe_001.svg b/tests/assets/shogi/observe_001.svg
new file mode 100644
--- /dev/null
+++ b/tests/assets/shogi/observe_001.svg
@@ -0,0 +1,2 @@
+<?xml version="1.0" encoding="utf-8" ?>
+<svg baseProfile="full" height="500.0" version="1.1" width="550.0" xmlns="http://www.w3.org/2000/svg" xmlns:ev="http://www.w3.org/2001/xml-events" xmlns:xlink="http://www.w3.org/1999/xlink"><defs /><rect fill="#202020" height="500" width="575.0" x="0" y="0" /><g transform="scale(1.0)"><rect fill="#202020" height="500" width="550" x="0" y="0" /><g transform="translate(15,0) translate(25.0,25.0)"><g id="hlines" stroke="gray"><line stroke-width="2px" x1="0" x2="450" y1="50" y2="50" /><line stroke-width="2px" x1="0" x2="450" y1="100" y2="100" /><line stroke-width="2px" x1="0" x2="450" y1="150" y2="150" /><line stroke-width="2px" x1="0" x2="450" y1="200" y2="200" /><line stroke-width="2px" x1="0" x2="450" y1="250" y2="250" /><line stroke-width="2px" x1="0" x2="450" y1="300" y2="300" /><line stroke-width="2px" x1="0" x2="450" y1="350" y2="350" /><line stroke-width="2px" x1="0" x2="450" y1="400" y2="400" /></g><g id="vline" stroke="gray"><line stroke-width="2px" x1="50" x2="50" y1="0" y2="450" /><line stroke-width="2px" x1="100" x2="100" y1="0" y2="450" /><line stroke-width="2px" x1="150" x2="150" y1="0" y2="450" /><line stroke-width="2px" x1="200" x2="200" y1="0" y2="450" /><line stroke-width="2px" x1="250" x2="250" y1="0" y2="450" /><line stroke-width="2px" x1="300" x2="300" y1="0" y2="450" /><line stroke-width="2px" x1="350" x2="350" y1="0" y2="450" /><line stroke-width="2px" x1="400" x2="400" y1="0" y2="450" /></g><rect fill="none" height="450" stroke="gray" stroke-width="4px" width="450" x="0" y="0" /><g fill="gray" id="cord"><text font-family="Serif" font-size="20px" x="455.0" y="30.0">一</text><text font-family="Serif" font-size="20px" x="420.0" y="-5.0">1</text><text font-family="Serif" font-size="20px" x="455.0" y="80.0">二</text><text font-family="Serif" font-size="20px" x="370.0" y="-5.0">2</text><text font-family="Serif" font-size="20px" x="455.0" y="130.0">三</text><text font-family="Serif" font-size="20px" x="320.0" y="-5.0">3</text><text font-family="Serif" font-size="20px" x="455.0" y="180.0">四</text><text font-family="Serif" font-size="20px" x="270.0" y="-5.0">4</text><text font-family="Serif" font-size="20px" x="455.0" y="229.99999999999997">五</text><text font-family="Serif" font-size="20px" x="220.00000000000003" y="-5.0">5</text><text font-family="Serif" font-size="20px" x="455.0" y="280.0">六</text><text font-family="Serif" font-size="20px" x="170.0" y="-5.0">6</text><text font-family="Serif" font-size="20px" x="455.0" y="330.0">七</text><text font-family="Serif" font-size="20px" x="120.0" y="-5.0">7</text><text font-family="Serif" font-size="20px" x="455.0" y="380.0">八</text><text font-family="Serif" font-size="20px" x="70.0" y="-5.0">8</text><text font-family="Serif" font-size="20px" x="455.0" y="430.0">九</text><text font-family="Serif" font-size="20px" x="20.0" y="-5.0">9</text></g><g><text fill="gray" font-family="Serif" font-size="45px" x="402.50000000000006" y="342.5">歩</text><text fill="gray" font-family="Serif" font-size="45px" x="302.5" y="292.5">歩</text><text fill="gray" font-family="Serif" font-size="45px" x="252.5" y="292.5">歩</text><text fill="gray" font-family="Serif" font-size="45px" x="202.5" y="342.5">歩</text><text fill="gray" font-family="Serif" font-size="45px" x="152.5" y="342.5">歩</text><text fill="gray" font-family="Serif" font-size="45px" x="102.49999999999999" y="292.5">歩</text><text fill="gray" font-family="Serif" font-size="45px" x="52.5" y="342.5">歩</text><text fill="gray" font-family="Serif" font-size="45px" x="402.50000000000006" y="442.5">香</text><text fill="gray" font-family="Serif" font-size="45px" x="102.49999999999999" y="342.5">桂</text><text fill="gray" font-family="Serif" font-size="45px" x="252.5" y="392.5">銀</text><text fill="gray" font-family="Serif" font-size="45px" x="102.49999999999999" y="442.5">銀</text><text fill="gray" font-family="Serif" font-size="45px" x="252.5" y="442.5">金</text><text fill="gray" font-family="Serif" font-size="45px" x="102.49999999999999" y="392.5">金</text><text fill="gray" font-family="Serif" font-size="45px" x="202.5" y="442.5">玉</text><text fill="gray" font-family="Serif" font-size="34px" x="475.0" y="160.0">☗</text><text fill="gray" font-family="Serif" font-size="34px" x="475.0" y="195.00000000000003">先</text><text fill="gray" font-family="Serif" font-size="34px" x="475.0" y="230.00000000000003">手</text><text fill="gray" font-family="Serif" font-size="34px" x="475.0" y="265.00000000000006" /><text fill="gray" font-family="Serif" font-size="34px" x="475.0" y="300.0">角</text><text fill="gray" font-family="Serif" font-size="34px" x="475.0" y="335.0">金</text><text fill="gray" font-family="Serif" font-size="34px" x="475.0" y="370.0">銀</text><text fill="gray" font-family="Serif" font-size="34px" x="475.0" y="405.0">歩</text><text fill="gray" font-family="Serif" font-size="34px" x="475.0" y="440.00000000000006">三</text></g><g transform="rotate(180,225.0,225.0)"><text fill="gray" font-family="Serif" font-size="45px" x="2.5" y="292.5">歩</text><text fill="gray" font-family="Serif" font-size="45px" x="152.5" y="342.5">歩</text><text fill="gray" font-family="Serif" font-size="45px" x="202.5" y="292.5">歩</text><text fill="gray" font-family="Serif" font-size="45px" x="252.5" y="292.5">歩</text><text fill="gray" font-family="Serif" font-size="45px" x="302.5" y="342.5">歩</text><text fill="gray" font-family="Serif" font-size="45px" x="352.5" y="292.5">歩</text><text fill="gray" font-family="Serif" font-size="45px" x="2.5" y="442.5">香</text><text fill="gray" font-family="Serif" font-size="45px" x="402.50000000000006" y="242.49999999999997">香</text><text fill="gray" font-family="Serif" font-size="45px" x="52.5" y="442.5">桂</text><text fill="gray" font-family="Serif" font-size="45px" x="352.5" y="442.5">桂</text><text fill="gray" font-family="Serif" font-size="45px" x="302.5" y="442.5">銀</text><text fill="gray" font-family="Serif" font-size="45px" x="352.5" y="392.5">飛</text><text fill="gray" font-family="Serif" font-size="45px" x="202.5" y="342.5">金</text><text fill="gray" font-family="Serif" font-size="45px" x="102.49999999999999" y="392.5">玉</text><text fill="gray" font-family="Serif" font-size="34px" x="455.0" y="125.0">☖</text><text fill="gray" font-family="Serif" font-size="34px" x="455.0" y="160.0">後</text><text fill="gray" font-family="Serif" font-size="34px" x="455.0" y="195.00000000000003">手</text><text fill="gray" font-family="Serif" font-size="34px" x="455.0" y="230.00000000000003" /><text fill="gray" font-family="Serif" font-size="34px" x="455.0" y="265.00000000000006">飛</text><text fill="gray" font-family="Serif" font-size="34px" x="455.0" y="300.0">角</text><text fill="gray" font-family="Serif" font-size="34px" x="455.0" y="335.0">桂</text><text fill="gray" font-family="Serif" font-size="34px" x="455.0" y="370.0">香</text><text fill="gray" font-family="Serif" font-size="34px" x="455.0" y="405.0">歩</text><text fill="gray" font-family="Serif" font-size="34px" x="455.0" y="440.00000000000006">二</text></g></g></g></svg>
\ No newline at end of file
diff --git a/tests/assets/shogi/observe_002.svg b/tests/assets/shogi/observe_002.svg
new file mode 100644
--- /dev/null
+++ b/tests/assets/shogi/observe_002.svg
@@ -0,0 +1,2 @@
+<?xml version="1.0" encoding="utf-8" ?>
+<svg baseProfile="full" height="500.0" version="1.1" width="550.0" xmlns="http://www.w3.org/2000/svg" xmlns:ev="http://www.w3.org/2001/xml-events" xmlns:xlink="http://www.w3.org/1999/xlink"><defs /><rect fill="#202020" height="500" width="575.0" x="0" y="0" /><g transform="scale(1.0)"><rect fill="#202020" height="500" width="550" x="0" y="0" /><g transform="translate(15,0) translate(25.0,25.0)"><g id="hlines" stroke="gray"><line stroke-width="2px" x1="0" x2="450" y1="50" y2="50" /><line stroke-width="2px" x1="0" x2="450" y1="100" y2="100" /><line stroke-width="2px" x1="0" x2="450" y1="150" y2="150" /><line stroke-width="2px" x1="0" x2="450" y1="200" y2="200" /><line stroke-width="2px" x1="0" x2="450" y1="250" y2="250" /><line stroke-width="2px" x1="0" x2="450" y1="300" y2="300" /><line stroke-width="2px" x1="0" x2="450" y1="350" y2="350" /><line stroke-width="2px" x1="0" x2="450" y1="400" y2="400" /></g><g id="vline" stroke="gray"><line stroke-width="2px" x1="50" x2="50" y1="0" y2="450" /><line stroke-width="2px" x1="100" x2="100" y1="0" y2="450" /><line stroke-width="2px" x1="150" x2="150" y1="0" y2="450" /><line stroke-width="2px" x1="200" x2="200" y1="0" y2="450" /><line stroke-width="2px" x1="250" x2="250" y1="0" y2="450" /><line stroke-width="2px" x1="300" x2="300" y1="0" y2="450" /><line stroke-width="2px" x1="350" x2="350" y1="0" y2="450" /><line stroke-width="2px" x1="400" x2="400" y1="0" y2="450" /></g><rect fill="none" height="450" stroke="gray" stroke-width="4px" width="450" x="0" y="0" /><g fill="gray" id="cord"><text font-family="Serif" font-size="20px" x="455.0" y="30.0">一</text><text font-family="Serif" font-size="20px" x="420.0" y="-5.0">1</text><text font-family="Serif" font-size="20px" x="455.0" y="80.0">二</text><text font-family="Serif" font-size="20px" x="370.0" y="-5.0">2</text><text font-family="Serif" font-size="20px" x="455.0" y="130.0">三</text><text font-family="Serif" font-size="20px" x="320.0" y="-5.0">3</text><text font-family="Serif" font-size="20px" x="455.0" y="180.0">四</text><text font-family="Serif" font-size="20px" x="270.0" y="-5.0">4</text><text font-family="Serif" font-size="20px" x="455.0" y="229.99999999999997">五</text><text font-family="Serif" font-size="20px" x="220.00000000000003" y="-5.0">5</text><text font-family="Serif" font-size="20px" x="455.0" y="280.0">六</text><text font-family="Serif" font-size="20px" x="170.0" y="-5.0">6</text><text font-family="Serif" font-size="20px" x="455.0" y="330.0">七</text><text font-family="Serif" font-size="20px" x="120.0" y="-5.0">7</text><text font-family="Serif" font-size="20px" x="455.0" y="380.0">八</text><text font-family="Serif" font-size="20px" x="70.0" y="-5.0">8</text><text font-family="Serif" font-size="20px" x="455.0" y="430.0">九</text><text font-family="Serif" font-size="20px" x="20.0" y="-5.0">9</text></g><g><text fill="gray" font-family="Serif" font-size="45px" x="402.50000000000006" y="342.5">歩</text><text fill="gray" font-family="Serif" font-size="45px" x="352.5" y="342.5">歩</text><text fill="gray" font-family="Serif" font-size="45px" x="302.5" y="342.5">歩</text><text fill="gray" font-family="Serif" font-size="45px" x="252.5" y="342.5">歩</text><text fill="gray" font-family="Serif" font-size="45px" x="202.5" y="342.5">歩</text><text fill="gray" font-family="Serif" font-size="45px" x="152.5" y="342.5">歩</text><text fill="gray" font-family="Serif" font-size="45px" x="102.49999999999999" y="292.5">歩</text><text fill="gray" font-family="Serif" font-size="45px" x="52.5" y="342.5">歩</text><text fill="gray" font-family="Serif" font-size="45px" x="2.5" y="342.5">歩</text><text fill="gray" font-family="Serif" font-size="45px" x="402.50000000000006" y="442.5">香</text><text fill="gray" font-family="Serif" font-size="45px" x="2.5" y="442.5">香</text><text fill="gray" font-family="Serif" font-size="45px" x="352.5" y="442.5">桂</text><text fill="gray" font-family="Serif" font-size="45px" x="52.5" y="442.5">桂</text><text fill="gray" font-family="Serif" font-size="45px" x="302.5" y="442.5">銀</text><text fill="gray" font-family="Serif" font-size="45px" x="102.49999999999999" y="442.5">銀</text><text fill="gray" font-family="Serif" font-size="45px" x="402.50000000000006" y="242.49999999999997">角</text><text fill="gray" font-family="Serif" font-size="45px" x="352.5" y="392.5">飛</text><text fill="gray" font-family="Serif" font-size="45px" x="252.5" y="442.5">金</text><text fill="gray" font-family="Serif" font-size="45px" x="152.5" y="442.5">金</text><text fill="gray" font-family="Serif" font-size="45px" x="202.5" y="442.5">玉</text><text fill="gray" font-family="Serif" font-size="34px" x="475.0" y="335.0">☗</text><text fill="gray" font-family="Serif" font-size="34px" x="475.0" y="370.0">先</text><text fill="gray" font-family="Serif" font-size="34px" x="475.0" y="405.0">手</text><text fill="gray" font-family="Serif" font-size="34px" x="475.0" y="440.00000000000006" /></g><g transform="rotate(180,225.0,225.0)"><text fill="gray" font-family="Serif" font-size="45px" x="2.5" y="342.5">歩</text><text fill="gray" font-family="Serif" font-size="45px" x="52.5" y="342.5">歩</text><text fill="gray" font-family="Serif" font-size="45px" x="102.49999999999999" y="292.5">歩</text><text fill="gray" font-family="Serif" font-size="45px" x="152.5" y="342.5">歩</text><text fill="gray" font-family="Serif" font-size="45px" x="202.5" y="342.5">歩</text><text fill="gray" font-family="Serif" font-size="45px" x="252.5" y="342.5">歩</text><text fill="gray" font-family="Serif" font-size="45px" x="302.5" y="342.5">歩</text><text fill="gray" font-family="Serif" font-size="45px" x="352.5" y="342.5">歩</text><text fill="gray" font-family="Serif" font-size="45px" x="402.50000000000006" y="342.5">歩</text><text fill="gray" font-family="Serif" font-size="45px" x="2.5" y="442.5">香</text><text fill="gray" font-family="Serif" font-size="45px" x="402.50000000000006" y="442.5">香</text><text fill="gray" font-family="Serif" font-size="45px" x="52.5" y="442.5">桂</text><text fill="gray" font-family="Serif" font-size="45px" x="352.5" y="442.5">桂</text><text fill="gray" font-family="Serif" font-size="45px" x="52.5" y="392.5">銀</text><text fill="gray" font-family="Serif" font-size="45px" x="302.5" y="442.5">銀</text><text fill="gray" font-family="Serif" font-size="45px" x="352.5" y="392.5">飛</text><text fill="gray" font-family="Serif" font-size="45px" x="152.5" y="442.5">金</text><text fill="gray" font-family="Serif" font-size="45px" x="252.5" y="442.5">金</text><text fill="gray" font-family="Serif" font-size="45px" x="202.5" y="442.5">玉</text><text fill="gray" font-family="Serif" font-size="34px" x="455.0" y="300.0">☖</text><text fill="gray" font-family="Serif" font-size="34px" x="455.0" y="335.0">後</text><text fill="gray" font-family="Serif" font-size="34px" x="455.0" y="370.0">手</text><text fill="gray" font-family="Serif" font-size="34px" x="455.0" y="405.0" /><text fill="gray" font-family="Serif" font-size="34px" x="455.0" y="440.00000000000006">角</text></g></g></g></svg>
\ No newline at end of file
diff --git a/tests/assets/shogi/observe_003.svg b/tests/assets/shogi/observe_003.svg
new file mode 100644
--- /dev/null
+++ b/tests/assets/shogi/observe_003.svg
@@ -0,0 +1,2 @@
+<?xml version="1.0" encoding="utf-8" ?>
+<svg baseProfile="full" height="500.0" version="1.1" width="550.0" xmlns="http://www.w3.org/2000/svg" xmlns:ev="http://www.w3.org/2001/xml-events" xmlns:xlink="http://www.w3.org/1999/xlink"><defs /><rect fill="#202020" height="500" width="575.0" x="0" y="0" /><g transform="scale(1.0)"><rect fill="#202020" height="500" width="550" x="0" y="0" /><g transform="translate(15,0) translate(25.0,25.0)"><g id="hlines" stroke="gray"><line stroke-width="2px" x1="0" x2="450" y1="50" y2="50" /><line stroke-width="2px" x1="0" x2="450" y1="100" y2="100" /><line stroke-width="2px" x1="0" x2="450" y1="150" y2="150" /><line stroke-width="2px" x1="0" x2="450" y1="200" y2="200" /><line stroke-width="2px" x1="0" x2="450" y1="250" y2="250" /><line stroke-width="2px" x1="0" x2="450" y1="300" y2="300" /><line stroke-width="2px" x1="0" x2="450" y1="350" y2="350" /><line stroke-width="2px" x1="0" x2="450" y1="400" y2="400" /></g><g id="vline" stroke="gray"><line stroke-width="2px" x1="50" x2="50" y1="0" y2="450" /><line stroke-width="2px" x1="100" x2="100" y1="0" y2="450" /><line stroke-width="2px" x1="150" x2="150" y1="0" y2="450" /><line stroke-width="2px" x1="200" x2="200" y1="0" y2="450" /><line stroke-width="2px" x1="250" x2="250" y1="0" y2="450" /><line stroke-width="2px" x1="300" x2="300" y1="0" y2="450" /><line stroke-width="2px" x1="350" x2="350" y1="0" y2="450" /><line stroke-width="2px" x1="400" x2="400" y1="0" y2="450" /></g><rect fill="none" height="450" stroke="gray" stroke-width="4px" width="450" x="0" y="0" /><g fill="gray" id="cord"><text font-family="Serif" font-size="20px" x="455.0" y="30.0">一</text><text font-family="Serif" font-size="20px" x="420.0" y="-5.0">1</text><text font-family="Serif" font-size="20px" x="455.0" y="80.0">二</text><text font-family="Serif" font-size="20px" x="370.0" y="-5.0">2</text><text font-family="Serif" font-size="20px" x="455.0" y="130.0">三</text><text font-family="Serif" font-size="20px" x="320.0" y="-5.0">3</text><text font-family="Serif" font-size="20px" x="455.0" y="180.0">四</text><text font-family="Serif" font-size="20px" x="270.0" y="-5.0">4</text><text font-family="Serif" font-size="20px" x="455.0" y="229.99999999999997">五</text><text font-family="Serif" font-size="20px" x="220.00000000000003" y="-5.0">5</text><text font-family="Serif" font-size="20px" x="455.0" y="280.0">六</text><text font-family="Serif" font-size="20px" x="170.0" y="-5.0">6</text><text font-family="Serif" font-size="20px" x="455.0" y="330.0">七</text><text font-family="Serif" font-size="20px" x="120.0" y="-5.0">7</text><text font-family="Serif" font-size="20px" x="455.0" y="380.0">八</text><text font-family="Serif" font-size="20px" x="70.0" y="-5.0">8</text><text font-family="Serif" font-size="20px" x="455.0" y="430.0">九</text><text font-family="Serif" font-size="20px" x="20.0" y="-5.0">9</text></g><g><text fill="gray" font-family="Serif" font-size="45px" x="402.50000000000006" y="342.5">歩</text><text fill="gray" font-family="Serif" font-size="45px" x="352.5" y="342.5">歩</text><text fill="gray" font-family="Serif" font-size="45px" x="302.5" y="342.5">歩</text><text fill="gray" font-family="Serif" font-size="45px" x="252.5" y="342.5">歩</text><text fill="gray" font-family="Serif" font-size="45px" x="202.5" y="342.5">歩</text><text fill="gray" font-family="Serif" font-size="45px" x="152.5" y="342.5">歩</text><text fill="gray" font-family="Serif" font-size="45px" x="102.49999999999999" y="292.5">歩</text><text fill="gray" font-family="Serif" font-size="45px" x="52.5" y="342.5">歩</text><text fill="gray" font-family="Serif" font-size="45px" x="2.5" y="342.5">歩</text><text fill="gray" font-family="Serif" font-size="45px" x="402.50000000000006" y="442.5">香</text><text fill="gray" font-family="Serif" font-size="45px" x="2.5" y="442.5">香</text><text fill="gray" font-family="Serif" font-size="45px" x="352.5" y="442.5">桂</text><text fill="gray" font-family="Serif" font-size="45px" x="52.5" y="442.5">桂</text><text fill="gray" font-family="Serif" font-size="45px" x="302.5" y="442.5">銀</text><text fill="gray" font-family="Serif" font-size="45px" x="102.49999999999999" y="442.5">銀</text><text fill="gray" font-family="Serif" font-size="45px" x="402.50000000000006" y="242.49999999999997">角</text><text fill="gray" font-family="Serif" font-size="45px" x="352.5" y="392.5">飛</text><text fill="gray" font-family="Serif" font-size="45px" x="252.5" y="442.5">金</text><text fill="gray" font-family="Serif" font-size="45px" x="152.5" y="442.5">金</text><text fill="gray" font-family="Serif" font-size="45px" x="202.5" y="442.5">玉</text><text fill="gray" font-family="Serif" font-size="34px" x="475.0" y="335.0">☗</text><text fill="gray" font-family="Serif" font-size="34px" x="475.0" y="370.0">先</text><text fill="gray" font-family="Serif" font-size="34px" x="475.0" y="405.0">手</text><text fill="gray" font-family="Serif" font-size="34px" x="475.0" y="440.00000000000006" /></g><g transform="rotate(180,225.0,225.0)"><text fill="gray" font-family="Serif" font-size="45px" x="2.5" y="342.5">歩</text><text fill="gray" font-family="Serif" font-size="45px" x="52.5" y="342.5">歩</text><text fill="gray" font-family="Serif" font-size="45px" x="102.49999999999999" y="292.5">歩</text><text fill="gray" font-family="Serif" font-size="45px" x="152.5" y="342.5">歩</text><text fill="gray" font-family="Serif" font-size="45px" x="202.5" y="342.5">歩</text><text fill="gray" font-family="Serif" font-size="45px" x="252.5" y="342.5">歩</text><text fill="gray" font-family="Serif" font-size="45px" x="302.5" y="342.5">歩</text><text fill="gray" font-family="Serif" font-size="45px" x="352.5" y="342.5">歩</text><text fill="gray" font-family="Serif" font-size="45px" x="402.50000000000006" y="342.5">歩</text><text fill="gray" font-family="Serif" font-size="45px" x="2.5" y="442.5">香</text><text fill="gray" font-family="Serif" font-size="45px" x="402.50000000000006" y="442.5">香</text><text fill="gray" font-family="Serif" font-size="45px" x="52.5" y="442.5">桂</text><text fill="gray" font-family="Serif" font-size="45px" x="352.5" y="442.5">桂</text><text fill="gray" font-family="Serif" font-size="45px" x="52.5" y="392.5">銀</text><text fill="gray" font-family="Serif" font-size="45px" x="302.5" y="442.5">銀</text><text fill="gray" font-family="Serif" font-size="45px" x="352.5" y="392.5">飛</text><text fill="gray" font-family="Serif" font-size="45px" x="152.5" y="442.5">金</text><text fill="gray" font-family="Serif" font-size="45px" x="252.5" y="442.5">金</text><text fill="gray" font-family="Serif" font-size="45px" x="202.5" y="442.5">玉</text><text fill="gray" font-family="Serif" font-size="34px" x="455.0" y="300.0">☖</text><text fill="gray" font-family="Serif" font-size="34px" x="455.0" y="335.0">後</text><text fill="gray" font-family="Serif" font-size="34px" x="455.0" y="370.0">手</text><text fill="gray" font-family="Serif" font-size="34px" x="455.0" y="405.0" /><text fill="gray" font-family="Serif" font-size="34px" x="455.0" y="440.00000000000006">角</text></g></g></g></svg>
\ No newline at end of file
diff --git a/tests/test_shogi.py b/tests/test_shogi.py
--- a/tests/test_shogi.py
+++ b/tests/test_shogi.py
@@ -506,6 +506,129 @@ def test_buggy_samples():
     # assert ~state.legal_action_mask[13 * 81 + 72]  # = 1125, promote + left (91角成）
 
 
+def test_observe():
+    s: State = _init()
+    obs = observe(s, s.curr_player)
+
+    assert obs.shape == (119, 9, 9)
+
+    expected = jnp.bool_([[0.,0.,0.,0.,0.,0.,1.,0.,0.],
+                          [0.,0.,0.,0.,0.,0.,1.,0.,0.],
+                          [0.,0.,0.,0.,0.,0.,1.,0.,0.],
+                          [0.,0.,0.,0.,0.,0.,1.,0.,0.],
+                          [0.,0.,0.,0.,0.,0.,1.,0.,0.],
+                          [0.,0.,0.,0.,0.,0.,1.,0.,0.],
+                          [0.,0.,0.,0.,0.,0.,1.,0.,0.],
+                          [0.,0.,0.,0.,0.,0.,1.,0.,0.],
+                          [0.,0.,0.,0.,0.,0.,1.,0.,0.]])
+    assert (obs[PAWN] == expected).all()  # 0
+
+    expected = jnp.bool_([[0.,0.,0.,0.,0.,1.,0.,0.,0.],
+                          [0.,0.,0.,0.,0.,1.,0.,0.,0.],
+                          [0.,0.,0.,0.,0.,1.,0.,0.,0.],
+                          [0.,0.,0.,0.,0.,1.,0.,0.,0.],
+                          [0.,0.,0.,0.,0.,1.,0.,0.,0.],
+                          [0.,0.,0.,0.,0.,1.,0.,0.,0.],
+                          [0.,0.,0.,0.,0.,1.,0.,0.,0.],
+                          [0.,0.,0.,0.,0.,1.,0.,0.,0.],
+                          [0.,0.,0.,0.,0.,1.,0.,0.,0.]])
+    assert (obs[DRAGON + 1] == expected).all()  # 14
+
+    expected = jnp.bool_([[0.,0.,0.,0.,0.,1.,1.,1.,0.],
+                          [0.,0.,0.,0.,0.,1.,1.,1.,1.],
+                          [0.,0.,0.,0.,0.,1.,1.,1.,1.],
+                          [0.,0.,0.,0.,0.,1.,0.,1.,1.],
+                          [0.,0.,0.,0.,0.,1.,0.,1.,1.],
+                          [0.,0.,0.,0.,0.,1.,0.,1.,1.],
+                          [0.,0.,0.,0.,0.,1.,1.,1.,1.],
+                          [0.,0.,0.,0.,0.,1.,0.,1.,0.],
+                          [0.,0.,0.,0.,0.,1.,1.,1.,1.]])
+    assert (obs[28] == expected).all()  # 利きの数1以上
+
+    expected = jnp.bool_([[0.,0.,0.,0.,0.,0.,1.,1.,0.],
+                          [0.,0.,0.,0.,0.,0.,0.,0.,0.],
+                          [0.,0.,0.,0.,0.,0.,0.,1.,0.],
+                          [0.,0.,0.,0.,0.,0.,0.,1.,0.],
+                          [0.,0.,0.,0.,0.,0.,0.,1.,1.],
+                          [0.,0.,0.,0.,0.,0.,0.,1.,0.],
+                          [0.,0.,0.,0.,0.,0.,1.,1.,1.],
+                          [0.,0.,0.,0.,0.,0.,0.,1.,0.],
+                          [0.,0.,0.,0.,0.,0.,1.,0.,0.]])
+    assert (obs[29] == expected).all()  # 利きの数2以上
+
+    expected = jnp.bool_([[0.,0.,0.,0.,0.,0.,0.,0.,0.],
+                          [0.,0.,0.,0.,0.,0.,0.,0.,0.],
+                          [0.,0.,0.,0.,0.,0.,0.,1.,0.],
+                          [0.,0.,0.,0.,0.,0.,0.,1.,0.],
+                          [0.,0.,0.,0.,0.,0.,0.,1.,0.],
+                          [0.,0.,0.,0.,0.,0.,0.,1.,0.],
+                          [0.,0.,0.,0.,0.,0.,0.,1.,0.],
+                          [0.,0.,0.,0.,0.,0.,0.,0.,0.],
+                          [0.,0.,0.,0.,0.,0.,1.,0.,0.]])
+    assert (obs[30] == expected).all()  # 利きの数3以上
+
+    expected = jnp.bool_([[0.,0.,1.,0.,0.,0.,0.,0.,0.],
+                          [0.,0.,1.,0.,0.,0.,0.,0.,0.],
+                          [0.,0.,1.,0.,0.,0.,0.,0.,0.],
+                          [0.,0.,1.,0.,0.,0.,0.,0.,0.],
+                          [0.,0.,1.,0.,0.,0.,0.,0.,0.],
+                          [0.,0.,1.,0.,0.,0.,0.,0.,0.],
+                          [0.,0.,1.,0.,0.,0.,0.,0.,0.],
+                          [0.,0.,1.,0.,0.,0.,0.,0.,0.],
+                          [0.,0.,1.,0.,0.,0.,0.,0.,0.]])
+    assert (obs[31] == expected).all()
+
+    expected = jnp.bool_([[0.,0.,0.,1.,0.,0.,0.,0.,0.],
+                          [0.,0.,0.,1.,0.,0.,0.,0.,0.],
+                          [0.,0.,0.,1.,0.,0.,0.,0.,0.],
+                          [0.,0.,0.,1.,0.,0.,0.,0.,0.],
+                          [0.,0.,0.,1.,0.,0.,0.,0.,0.],
+                          [0.,0.,0.,1.,0.,0.,0.,0.,0.],
+                          [0.,0.,0.,1.,0.,0.,0.,0.,0.],
+                          [0.,0.,0.,1.,0.,0.,0.,0.,0.],
+                          [0.,0.,0.,1.,0.,0.,0.,0.,0.]])
+    assert (obs[45] == expected).all()
+
+    expected = jnp.bool_([[1.,1.,1.,1.,0.,0.,0.,0.,0.],
+                          [0.,1.,0.,1.,0.,0.,0.,0.,0.],
+                          [1.,1.,1.,1.,0.,0.,0.,0.,0.],
+                          [1.,1.,0.,1.,0.,0.,0.,0.,0.],
+                          [1.,1.,0.,1.,0.,0.,0.,0.,0.],
+                          [1.,1.,0.,1.,0.,0.,0.,0.,0.],
+                          [1.,1.,1.,1.,0.,0.,0.,0.,0.],
+                          [1.,1.,1.,1.,0.,0.,0.,0.,0.],
+                          [0.,1.,1.,1.,0.,0.,0.,0.,0.]])
+    assert (obs[59] == expected).all()
+
+    # 駒打ち
+    sfen = "1ns4nl/1r4k2/2p1gp3/1p1pp3p/l8/2P2PP2/1PNPP3P/2G2S3/2S1KG2L b BGS3Prbnl2p 1"
+    s = _from_sfen(sfen)
+    visualize(s, "tests/assets/shogi/observe_001.svg")
+    obs = observe(s, s.curr_player)
+
+    filled = [0, 1, 2, 16, 20, 24, 28, 29, 36, 40, 52, 54]
+    for i in range(56):
+        if i in filled:
+            assert obs[62 + i].all()
+        else:
+            assert (~obs[62 + i]).all()
+
+    # 王手
+    sfen = "lnsgkg1nl/1r5s1/pppppp1pp/6p2/8B/2P6/PP1PPPPPP/7R1/LNSGKGSNL b b 1"  # 先手番
+    s = _from_sfen(sfen)
+    visualize(s, "tests/assets/shogi/observe_002.svg")
+    obs = observe(s, s.curr_player)
+    assert (~obs[-1]).all()
+
+    sfen = "lnsgkg1nl/1r5s1/pppppp1pp/6p2/8B/2P6/PP1PPPPPP/7R1/LNSGKGSNL w b 1"  # 後手番
+    s = _from_sfen(sfen)
+    visualize(s, "tests/assets/shogi/observe_003.svg")
+    obs = observe(s, s.curr_player)
+    assert obs[-1].all()
+
+    # TODO: player_id != curr_player
+
+
 def test_sfen():
     sfen = "lnsgkg1nl/1r5s1/pppppp1pp/6p2/8B/2P6/PP1PPPPPP/7R1/LNSGKGSNL b b 1"
     s = _from_sfen(sfen)