Update app.py

app.py

@@ -86,17 +86,21 @@ def step_episode(inputs):
         gr.Info(f"Environment ended game: ILLEGAL_MOVE", duration=None)
         svg_string = create_chess_board_svg()
         return svg_string, START_POSITION, "", [START_POSITION, []]
-    
-    svg_string = create_chess_board_svg(new_state)
+
+    try:
+        mv = chess.Move.from_uci(action)
+        svg_string = create_chess_board_svg(new_state, lastmove=mv)
+    except:
+        svg_string = create_chess_board_svg(new_state)
     if not svg_string:
         svg_string = create_chess_board_svg()
         return svg_string, [START_POSITION, []], START_POSITION, ""
     return svg_string, new_state, ", ".join(history), [new_state, history]
 
-def create_chess_board_svg(fen=None):
+def create_chess_board_svg(fen=None, lastmove=None):
     try:
         board = chess.Board(fen) if fen else chess.Board()
-        return chess.svg.board(board, size=400)
+        return chess.svg.board(board, lastmove=lastmove, size=400)
     except:
         gr.Info(f"Python-Chess board visualization cannot be rendered from FEN: {fen}", duration=None)
         return ""