bug fixes

TODO

@@ -8,16 +8,17 @@
 [] engine status
 [] README
 [] Release notes
-[] List edit settings/object edit settings?
 [] more AI modes?
+[] P+noise pass not noisy
 
 [] show PV on hint hover?
 
+
 lo prio
 [] more teaching / groups in danger? hard
 [] limit below board buttons somehow to not be too big
 [] box to label ? split in status and comment?
-
+[] List edit settings/object edit settings?
 
 wont do
 [] Self-play tournaments?

config.json

@@ -1,45 +1,123 @@
 {
-  "engine": {
-    "katago": "KataGo/katago-bs",
-    "model": " models/b15-1.3.2.txt.gz",
-    "config": "KataGo/analysis_config.cfg",
-    "threads": 8,
-    "visits": 500,
-    "enable_ownership": true
-  },
-  "files": {
-    "sgf_load": "~/Downloads",
-    "sgf_save": "./sgfout"
-  },
-  "game": {
-    "init_size": 19,
-    "init_komi": {"19": 6.5},
-    "sweep_visits_frac": 0.05
-  },
-  "trainer": {
-    "num_undo_prompts": [1,1,1,0.5,0],
-    "eval_thresholds": [12,6,3,1.5,0.5],
-    "balance_target_score": 2,
-    "balance_randomize_score": 1,
-    "balance_max_lost": 5,
-    "balance_min_visits": 20,
-    "noise_strength": 0.05,
-    "eval_off_show_last": 3
-  },
-  "board_ui": {
-    "starpoint_size": 0.1,
-    "stone_size": 0.475,
-    "eval_dot_max_size": 0.5,
-    "stones": {"B":  [0.05, 0.05, 0.05], "W": [0.95, 0.95, 0.95] },
-    "outline": {"B":  [0.3,0.3,0.3,0.5], "W": [0.7, 0.7, 0.7,0.5] },
-    "ghost_alpha": 0.5,
-    "top_move_x_alpha": 0.3,
-    "_child_scale": 0.95,
-    "eval_colors": [[0.447, 0.129, 0.42], [0.8, 0, 0], [0.9,0.4,0.1], [0.85, 0.89, 0.30], [0.67, 0.9, 0.18], [0.117, 0.588, 0]],
-    "line_color": [0,0,0],
-    "policy_color": [0,0.2,0.8]
-  },
-  "debug": {
-    "level": 2
-  }
-}
+    "engine": {
+        "katago": "../KataGo/cpp/katago",
+        "model": " models/b15-1.3.2.txt.gz",
+        "config": "KataGo/analysis_config.cfg",
+        "threads": 8,
+        "max_visits": 500,
+        "max_time": 3.0,
+        "enable_ownership": true
+    },
+    "files": {
+        "sgf_load": "~/Downloads",
+        "sgf_save": "./sgfout"
+    },
+    "game": {
+        "init_size": 19,
+        "init_komi": {
+            "19": 6.5
+        },
+        "sweep_visits_frac": 0.05
+    },
+    "trainer": {
+        "num_undo_prompts": [
+            1,
+            1,
+            1,
+            0.5,
+            0
+        ],
+        "eval_thresholds": [
+            12,
+            6,
+            3,
+            1.5,
+            0.5
+        ],
+        "balance_target_score": 2,
+        "balance_random_loss": 1,
+        "balance_max_loss": 5,
+        "balance_min_visits": 20,
+        "noise_strength": 0.1,
+        "eval_off_show_last": 3
+    },
+    "board_ui": {
+        "starpoint_size": 0.1,
+        "stone_size": 0.475,
+        "eval_dot_max_size": 0.5,
+        "stones": {
+            "B": [
+                0.05,
+                0.05,
+                0.05
+            ],
+            "W": [
+                0.95,
+                0.95,
+                0.95
+            ]
+        },
+        "outline": {
+            "B": [
+                0.3,
+                0.3,
+                0.3,
+                0.5
+            ],
+            "W": [
+                0.7,
+                0.7,
+                0.7,
+                0.5
+            ]
+        },
+        "ghost_alpha": 0.5,
+        "top_move_x_alpha": 0.3,
+        "_child_scale": 0.95,
+        "eval_colors": [
+            [
+                0.447,
+                0.129,
+                0.42
+            ],
+            [
+                0.8,
+                0,
+                0
+            ],
+            [
+                0.9,
+                0.4,
+                0.1
+            ],
+            [
+                0.85,
+                0.89,
+                0.3
+            ],
+            [
+                0.67,
+                0.9,
+                0.18
+            ],
+            [
+                0.117,
+                0.588,
+                0
+            ]
+        ],
+        "line_color": [
+            0,
+            0,
+            0
+        ],
+        "policy_color": [
+            0,
+            0.2,
+            0.8
+        ]
+    },
+    "debug": {
+        "level": 2
+    }
+}

engine.py

@@ -91,7 +91,7 @@ class KataGoEngine:
             self.katago_process.stdin.flush()
 
     def request_analysis(
-        self, analysis_node: GameNode, callback: Callable, visits: int = None, priority: int = 0, ownership: Optional[bool] = None, next_move=None,
+        self, analysis_node: GameNode, callback: Callable, visits: int = None, time_limit=True, priority: int = 0, ownership: Optional[bool] = None, next_move=None,
     ):
         moves = [m for node in analysis_node.nodes_from_root for m in node.move_with_placements]
         if next_move:
@@ -103,13 +103,14 @@ class KataGoEngine:
             "rules": self.get_rules(analysis_node),
             "priority": self.base_priority + priority,
             "analyzeTurns": [len(moves)],
-            "maxVisits": visits or self.config["visits"],
+            "maxVisits": visits or self.config["max_visits"],
             "komi": analysis_node.komi,
             "boardXSize": size_x,
             "boardYSize": size_y,
             "includeOwnership": ownership,
             "includePolicy": not next_move,
             "moves": [[m.player, m.gtp()] for m in moves],
+            "overrideSettings": {"maxTime": self.config["max_time"] if time_limit else 1000.0}
             # "overrideSettings": {"playoutDoublingAdvantage": 3.0, "playoutDoublingAdvantagePla":  'BLACK' if not moves or moves[-1].player == 'W' else "WHITE"}
         }
         self.send_query(query, callback)

game.py

@@ -45,12 +45,12 @@ class Game:
         self.current_node = self.root
         self._init_chains()
 
-        def analyze_game(_dt):
-            self.engine.on_new_game()
-            for node in self.root.nodes_in_tree:
-                node.analyze(self.engine, priority=-1_000_000)
+        Clock.schedule_once(lambda _dt: self.analyze_all_nodes(-1_000_000), -1)  # return faster
 
-        Clock.schedule_once(analyze_game, -1)  # return faster
+    def analyze_all_nodes(self, priority=0):
+        self.engine.on_new_game()
+        for node in self.root.nodes_in_tree:
+            node.analyze(self.engine, priority=priority)
 
     # -- move tree functions --
     def _init_chains(self):
@@ -215,14 +215,13 @@ class Game:
         candidate_ai_moves = cn.candidate_moves
         ai_mode = self.katrain.controls.ai_mode(cn.next_player)
 
-        if "policy" in ai_mode and cn.policy:
+        if ("policy" in ai_mode or "noise" in ai_mode) and cn.policy:
             policy_moves = cn.policy_ranking
-            self.katrain.log(f"Top 5 policy moves are: {policy_moves[:5]}", OUTPUT_DEBUG)
             aimove = policy_moves[0][0]
-        elif "noise" in ai_mode and cn.policy:
-            noise = train_settings["noise_strength"]
-            policy_moves = [(mv, pol + random.gauss(0, noise)) for mv, pol in cn.policy_ranking]
-            aimove = max(policy_moves, key=lambda mp: mp[1])[0]
+            if not aimove.is_pass and "noise" in ai_mode:
+                noise = train_settings["noise_strength"]
+                policy_moves = [(mv, pol + random.gauss(0, noise)) for mv, pol in policy_moves if not mv.is_pass]
+                aimove = max(policy_moves, key=lambda mp: mp[1])[0]
         elif "balance" in ai_mode and candidate_ai_moves[0]["move"] != "pass":  # don't play suicidal to balance score - pass when it's best
             sign = cn.player_sign(cn.next_player)  # TODO check
             sel_moves = [  # top move, or anything not too bad, or anything that makes you still ahead
@@ -231,8 +230,8 @@ class Game:
                 if i == 0
                 or move["visits"] >= train_settings["balance_min_visits"]
                 and (
-                    move["pointsLost"] < train_settings["balance_randomize_score"]
-                    or move["pointsLost"] < train_settings["balance_max_lost"]
+                    move["pointsLost"] < train_settings["balance_random_loss"]
+                    or move["pointsLost"] < train_settings["balance_max_loss"]
                     and sign * move["scoreLead"] > train_settings["balance_target_score"]
                 )
             ]
@@ -240,7 +239,6 @@ class Game:
         elif "jigo" in ai_mode and candidate_ai_moves[0]["move"] != "pass":
             sign = cn.player_sign(cn.next_player)  # TODO check
             jigo_move = min(candidate_ai_moves, key=lambda move: abs(sign * move["scoreLead"] - 0.5))
-            print("NP sign", sign, "JIGO MOVE", jigo_move)
             aimove = Move.from_gtp(jigo_move["move"], player=cn.next_player)
         else:
             if "default" not in ai_mode:
@@ -249,9 +247,9 @@ class Game:
         self.play(aimove)
 
     def analyze_undo(self, node, train_config):
-        if node != self.current_node or node.auto_undo is not None or not node.analysis_ready or not node.single_move:
-            return
         move = node.single_move
+        if node != self.current_node or node.auto_undo is not None or not node.analysis_ready or not move:
+            return
         points_lost = node.points_lost
         thresholds = train_config["eval_thresholds"]
         num_undo_prompts = train_config["num_undo_prompts"]
@@ -264,7 +262,7 @@ class Game:
             undo = False
         elif num_undos < 1:  # probability
             undo = int(node.undo_threshold < num_undos) and len(node.parent.children) == 1
-            xmsg = " (with {num_undos:.0%} probability at this level of mistake)" + xmsg
+            xmsg = f" (with {num_undos:.0%} probability at this level of mistake)" + xmsg
         else:
             undo = len(node.parent.children) <= num_undos
             if len(node.parent.children) == num_undos:
@@ -283,14 +281,14 @@ class Game:
             return
 
         if mode == "extra":
-            visits = cn.analysis["root"]["visits"] + self.engine.config["visits"]
+            visits = cn.analysis["root"]["visits"] + self.engine.config["max_visits"]
             self.katrain.controls.set_status(f"Performing additional analysis to {visits} visits")
-            cn.analyze(self.engine, visits=visits, priority=-1_000)
+            cn.analyze(self.engine, visits=visits, priority=-1_000, time_limit=False)
             return
         elif mode == "sweep":
             board_size_x, board_size_y = self.board_size
             analyze_moves = [Move(coords=(x, y), player=cn.next_player) for x in range(board_size_x) for y in range(board_size_y) if (x, y) not in stones]
-            visits = int(self.engine.config["visits"] * self.config["sweep_visits_frac"] + 0.5)
+            visits = int(self.engine.config["max_visits"] * self.config["sweep_visits_frac"] + 0.5)
             self.katrain.controls.set_status(f"Refining analysis of entire board to {visits} visits")
             priority = -1_000_000_000
         else:  # mode=='equalize':
@@ -299,4 +297,4 @@ class Game:
             self.katrain.controls.set_status(f"Equalizing analysis of candidate moves to {visits} visits")
             priority = -1_000
         for move in analyze_moves:
-            cn.analyze(self.engine, priority, visits=visits, refine_move=move)
+            cn.analyze(self.engine, priority, visits=visits, refine_move=move, time_limit=False)  # explicitly requested so take as long as you need

game_node.py

@@ -29,8 +29,8 @@ class GameNode(SGFNode):
         return properties
 
     # various analysis functions
-    def analyze(self, engine, priority=0, visits=None, refine_move=None):
-        engine.request_analysis(self, lambda result: self.set_analysis(result, refine_move), priority=priority, visits=visits, next_move=refine_move)
+    def analyze(self, engine, priority=0, visits=None, time_limit=True, refine_move=None):
+        engine.request_analysis(self, lambda result: self.set_analysis(result, refine_move), priority=priority, visits=visits, time_limit=time_limit, next_move=refine_move)
 
     def update_move_analysis(self, move_analysis, move_gtp):
         cur = self.analysis["moves"].get(move_gtp)
@@ -61,13 +61,12 @@ class GameNode(SGFNode):
         win_rate = win_rate or self.analysis["root"]["winrate"]
         return f"{'B' if win_rate > 0.5 else 'W'} {max(win_rate,1-win_rate):.1%}"
 
-    def comment(self, sgf=False, eval=False, hints=False):
+    def comment(self, sgf=False, teach=False, hints=False):
         single_move = self.single_move
         if not self.parent or not single_move:  # root
             return ""
 
         text = f"Move {self.depth}: {single_move.player} {single_move.gtp()}\n"
-
         if self.analysis_ready:
             score = self.score
             if sgf:
@@ -82,15 +81,15 @@ class GameNode(SGFNode):
                             text += f"Estimated point loss: {points_lost:.1f}\n"
                     else:
                         text += f"Move was predicted best move.\n"
-                if sgf or hints:
+                if sgf or hints or teach:
                     policy_ranking = self.parent.policy_ranking
                     policy_ix = [ix + 1 for (m, p), ix in zip(policy_ranking, range(len(policy_ranking))) if m == single_move]
                     if policy_ix:
                         text += f"Move was #{policy_ix[0]} according to policy.\n"
-                    if not policy_ix or policy_ix[0] != 1:
+                    if not policy_ix or policy_ix[0] != 1 and (sgf or hints):
                         text += f"Top policy move was {policy_ranking[0][0].gtp()}.\n"
-            if self.auto_undo:
-                text += "Move was automatically undone."
+            if self.auto_undo and sgf:
+                text += "Move was automatically undone in teaching mode."
         else:
             text = "No analysis available" if sgf else "Analyzing move..."
         return text

gui/badukpan.py

@@ -172,7 +172,7 @@ class BadukPanWidget(Widget):
                         drawn_stone[m.coords] = m.player
                         self.draw_stone(m.coords[0], m.coords[1], stone_color[m.player], outline_color[m.player], inner, evalcol, evalsize)
 
-            if katrain.game.current_node.is_root and katrain.config("debug/level") >= 3: # secret ;)
+            if katrain.game.current_node.is_root and katrain.config("debug/level") >= 3:  # secret ;)
                 for s in range(0, 19):
                     c = s
                     evalcol = self.eval_color(s)

gui/controls.py

@@ -49,22 +49,23 @@ class Controls(BoxLayout):
 
         if current_node:
             move = current_node.single_move
-            current_player_is_human_or_both_robots = (
-                not current_node.player or "ai" not in self.player_mode(current_node.player) or "ai" in self.player_mode(current_node.next_player)
-            )
-            if current_player_is_human_or_both_robots and not current_node.is_root and move:
-                info += current_node.comment(eval=True, hints=self.hints.active)
+            next_player_is_human_or_both_robots = current_node.player and ("ai" not in self.player_mode(current_node.player) or "ai" in self.player_mode(current_node.next_player))
+            current_player_is_ai_playing_human = current_node.player and "ai" in self.player_mode(current_node.player) and "ai" not in self.player_mode(current_node.next_player)
+            if next_player_is_human_or_both_robots and not current_node.is_root and move:
+                info += current_node.comment(teach="undo" in self.player_mode(current_node.player), hints=self.hints.active)
+            elif current_player_is_ai_playing_human:
+                info += current_node.parent.comment(teach="undo" in self.player_mode(current_node.next_player), hints=self.hints.active)
 
             if current_node.analysis_ready:
                 self.score.text = current_node.format_score()
                 self.win_rate.text = current_node.format_win_rate()
-                if move and current_player_is_human_or_both_robots:  # don't immediately hide this when an ai moves comes in
-                    self.points_lost.label = f"Point loss {move.player}{move.gtp()}"
+                if move and next_player_is_human_or_both_robots:  # don't immediately hide this when an ai moves comes in
+                    self.score_change.label = f"Score change"
                     points_lost = current_node.points_lost
-                    self.points_lost.text = f"{current_node.points_lost:.1f}" if points_lost else "..."
-            else:
-                self.points_lost.label = f"Point loss"
-                self.points_lost.text = ""
+                    self.score_change.text = f"{move.player}{-current_node.points_lost:+.1f}" if points_lost else "..."
+                elif not current_player_is_ai_playing_human:
+                    self.score_change.label = f"Score change"
+                    self.score_change.text = ""
 
             self.graph.update_value(current_node)
 

gui/kivyutils.py

@@ -215,8 +215,9 @@ class ScoreGraph(Label):
         self.nodes.extend([None] * max(0, index - (len(self.nodes) - 1)))
         self.nodes[index] = node
         if index + 1 < len(self.nodes) and (node is None or self.nodes[index + 1] not in node.children):
-            self.nodes = self.nodes[:index]  # on branch switching, don't show history from other branch
-            while node.children:  # but from this one if it exists
+            self.nodes = self.nodes[: index + 1]  # on branch switching, don't show history from other branch
+        if index == len(self.nodes) - 1:  # possibly just switched branch
+            while node.children:  # add children back
                 node = node.children[0]
                 self.nodes.append(node)
         self.on_size()

gui/popups.py

@@ -153,12 +153,15 @@ class ConfigPopup(QuickConfigGui):
         engine_updates = updated_cat["engine"]
         if "visits" in engine_updates:
             self.katrain.engine.visits = engine_updates["visits"]
-        if set(engine_updates) != {"visits"}:
+        if {key for key in engine_updates if key not in {"max_visits", "max_time"}}:
             self.katrain.log(f"Restarting Engine after {engine_updates} settings change")
             self.katrain.controls.set_status(f"Restarting Engine after {engine_updates} settings change")
             old_engine = self.katrain.engine
             self.katrain.engine = KataGoEngine(self.katrain, self.config["engine"])
             self.katrain.game.engine = self.katrain.engine
-            old_engine.shutdown(finish=True)
+            if getattr(old_engine, "katago_process"):
+                old_engine.shutdown(finish=True)
+            else:
+                self.katrain.game.analyze_all_nodes()  # old engine was broken, so make sure we redo any failures
 
         self.katrain.update_state(redraw_board=True)

katrain.kv

@@ -4,6 +4,7 @@
 
 #:set AI_MODES ['Default','Balance','Jigo','Policy','P+Noise']
 #:set PLAYER_MODES ['Human', 'Teach','AI:']
+#:set PLAYER_MODE_VALUES ['human','human+undo','ai']
 #:set BUTTON_COLOR [0.23, 0.30, 0.35, 1]
 #:set ICON_BUTTON_COLOR [0.71, 0.78, 0.81, 1]
 #:set BOARD_COLOR [0.85, 0.68, 0.40,1]
@@ -52,6 +53,11 @@
 
 <ToggleButtonContainer>:
 
+<LabelledTextInput>:
+    padding_y: 0.1*self.size[1]
+    padding_x: 7
+    font_size: 0.5*self.size[1]
+
 
 <StyledIconButton@StyledButton>:
     button_color: ICON_BUTTON_COLOR
@@ -110,7 +116,14 @@
 <LabelledCheckBox>:
     color: WHITE
 
-<CheckBoxHintLabel@Label+ButtonBehavior>:
+<CheckBoxHintLabel@Label>:
+    font_size: 0.4*self.size[1]
+    color: BLACK
+	halign: 'center'
+	valign: 'center'
+	size: self.texture_size
+
+<ClickableLabel@DarkLabel+ButtonBehavior>:
     font_size: dp(12)
     color: BLACK
 	halign: 'center'
@@ -121,7 +134,7 @@
 	orientation: 'horizontal'
 	text: ''
 	label: ''
-	CheckBoxHintLabel:
+	ClickableLabel:
 		size_hint: 0.6,None
 		height: self.font_size * 1.5
 		id: desc
@@ -168,6 +181,13 @@
             circle: self.pos[0] + self.width/2, self.pos[1] + self.height/2, min(self.height,self.width) * root.radius - 1
             width: 1.1 if root.outline else 0.0001
 
+<GraphMarkerLabel@Label>:
+    valign: 'bottom'
+    halign: 'left'
+    text: '+0'
+    color:  GREY
+    size: self.texture_size
+
 <ScoreGraph>:
     range_label_bottom: range_label_bottom
     range_label_top: range_label_top
@@ -195,23 +215,14 @@
             id: dot
             pos: root.dot_pos
             size: self.highlight_size,self.highlight_size
-    CheckBoxHintLabel:
-        id: range_label_top
+    GraphMarkerLabel:
         font_size: 0.12 * root.height
+        id: range_label_top
         pos: root.pos[0] + root.marginx * root.width + 3, root.pos[1] + root.height * (1 - 2 * root.marginy) - self.font_size - 1
-        valign: 'bottom'
-        halign: 'left'
-        text: 'B+0'
-        color:  GREY
-    CheckBoxHintLabel:
+    GraphMarkerLabel:
         id: range_label_bottom
         font_size: 0.12 * root.height
         pos: root.pos[0] + root.marginx * root.width + 3, root.pos[1] + 1
-        valign: 'bottom'
-        halign: 'left'
-        text: 'W+0'
-        color:  GREY
-
 
 
 <BadukPanWidget>:
@@ -297,7 +308,7 @@
 	analyze_tab_button: analyze_tab_button
     info: info
 	win_rate: win_rate
-	points_lost: points_lost
+	score_change: score_change
 	score: score
 	show_children: show_children
 	hints: hints
@@ -359,7 +370,7 @@
                 CheckBoxHintLabel:
                     text: 'top\nmoves'
                 CheckBoxHintLabel:
-                    text: 'owner'
+                    text: 'show\nowner'
                 CheckBoxHintLabel:
                     text: 'NN\npolicy'
                 CheckBox:
@@ -426,6 +437,7 @@
                 cols: 4
                 rows: 2
                 spacing: 1
+                padding: 1
                 size_hint: 1,0.125
                 BlackCircleWithText:
                     text: 'B'
@@ -433,13 +445,14 @@
                 ToggleButtonContainer:
                     size_hint: 0.6, 1
                     id: B_player_mode
-                    options: ['human','human+undo','ai']
+                    options: PLAYER_MODE_VALUES
                     labels: PLAYER_MODES
                     on_selection: root.katrain.update_state()
                 StyledSpinner:
                     id: B_AI_mode
                     size_hint: 0.3, 1
                     values: AI_MODES
+                    on_text: B_player_mode.children[0].trigger_action(duration=0)
                 Label:
                     size_hint: None,1
                     width: 3
@@ -449,13 +462,14 @@
                 ToggleButtonContainer:
                     size_hint: 0.6, 1
                     id: W_player_mode
-                    options: ['human','human+undo','ai']
+                    options: PLAYER_MODE_VALUES
                     labels: PLAYER_MODES
                     on_selection: root.katrain.update_state()
                 StyledSpinner:
                     id: W_AI_mode
                     size_hint: 0.3, 1
                     values: AI_MODES
+                    on_text: W_player_mode.children[0].trigger_action(duration=0)
                 Label:
                     size_hint: None,1
                     width: 3
@@ -463,9 +477,9 @@
                 text: ''
                 size_hint: 1,0.3
     CensorableLabel:
-        id: points_lost
+        id: score_change
         size_hint: 1, 0.0225
-        label: 'Point loss'
+        label: 'Score change'
         text: ''
 	TextInput:
 		id: info

katrain.py

@@ -79,11 +79,11 @@ class KaTrainGui(BoxLayout):
         threading.Thread(target=self._message_loop_thread, daemon=True).start()
         self._do_new_game()
 
-    def update_state(self, redraw_board=False):
+    def update_state(self, redraw_board=False):  # is called after every message and on receiving analyses and config changes
         # AI and Trainer/auto-undo handlers
         cn = self.game.current_node
         auto_undo = cn.player and "undo" in self.controls.player_mode(cn.player)
-        if auto_undo and cn.analysis_ready:
+        if auto_undo and cn.analysis_ready and cn.parent and cn.parent.analysis_ready:
             self.game.analyze_undo(cn, self.config("trainer"))  # not via message loop
 
         if cn.analysis_ready and "ai" in self.controls.player_mode(cn.next_player) and not cn.children and not self.game.game_ended and not (auto_undo and cn.auto_undo is None):
@@ -91,7 +91,7 @@ class KaTrainGui(BoxLayout):
 
         # Handle prisoners and next player display
         prisoners = self.game.prisoner_count
-        top, bot = self.board_controls.black_prisoners, self.board_controls.white_prisoners
+        top, bot = self.board_controls.black_prisoners.__self__, self.board_controls.white_prisoners.__self__  # no weakref
         if self.game.next_player == "W":
             top, bot = bot, top
         self.board_controls.mid_circles_container.clear_widgets()
@@ -99,8 +99,6 @@ class KaTrainGui(BoxLayout):
         self.board_controls.mid_circles_container.add_widget(top)
         self.board_controls.black_prisoners.text = str(prisoners[1])
         self.board_controls.white_prisoners.text = str(prisoners[0])
-
-        # Update board and status
         if redraw_board:
             Clock.schedule_once(self.board_gui.draw_board, -1)  # main thread needs to do this
         Clock.schedule_once(self.board_gui.draw_board_contents, -1)
@@ -114,10 +112,10 @@ class KaTrainGui(BoxLayout):
                 if game != self.game.game_id:
                     self.log(f"Message skipped as it is outdated (current game is {self.game.game_id}", OUTPUT_EXTRA_DEBUG)
                     continue
-                getattr(self, f"_do_{msg.replace('-','_')}")(*args)  # TODO update state?
+                getattr(self, f"_do_{msg.replace('-','_')}")(*args)
+                self.update_state()
             except Exception as e:
-                self.log(f"Exception in Engine thread: {e}", OUTPUT_ERROR)
-                raise
+                self.log(f"Exception in processing message {msg} {args}: {e}", OUTPUT_ERROR)
 
     def __call__(self, message, *args):
         if self.game:
@@ -127,31 +125,26 @@ class KaTrainGui(BoxLayout):
         self.game = Game(self, self.engine, self.config("game"), move_tree=move_tree)
         self.controls.select_mode("analyze" if move_tree and len(move_tree.nodes_in_tree) > 1 else "play")
         self.controls.graph.initialize_from_game(self.game.root)
-        self.update_state(redraw_board=True)  # TODO: just board here/redraw is in all anyway?
+        self.update_state(redraw_board=True)
 
     def _do_ai_move(self, node=None):
         if node is None or self.game.current_node == node:
             self.game.ai_move(self.config("trainer"))
-            self.update_state()
 
     def _do_undo(self, n_times=1):
         self.game.undo(n_times)
-        self.update_state()
 
     def _do_redo(self, n_times=1):
         self.game.redo(n_times)
-        self.update_state()
 
     def _do_switch_branch(self, direction):
         self.game.switch_branch(direction)
-        self.update_state()
 
     def _do_play(self, coords):
         try:
             self.game.play(Move(coords, player=self.game.next_player))
         except IllegalMoveException as e:
             self.controls.set_status(f"Illegal Move: {str(e)}")
-        self.update_state()
 
     def _do_analyze_extra(self, mode):
         self.game.analyze_extra(mode)
@@ -190,7 +183,7 @@ class KaTrainGui(BoxLayout):
         for pl in Move.PLAYERS:
             if not self.game.root.get_first(f"P{pl}"):
                 _, model_file = os.path.split(self.engine.config["model"])
-                self.game.root.properties[f"P{pl}"] = [f"KaTrain (KataGo {model_file})" if "ai" in self.controls.player_mode(pl) else "Player"]
+                self.game.root.properties[f"P{pl}"] = [f"KaTrain (KataGo {model_file})" if "ai" in self.controls.player_mode(pl) else "Player"]  # TODO: more dynamic?
         msg = self.game.write_sgf(self.config("files/sgf_save"))
         self.log(msg, OUTPUT_INFO)
         self.controls.set_status(msg)