Change post-race retry to avoid ghost

When doing a retry, the subsequent race will be against the ghost character. Doing a course 'change' rather than a retry avoids this (see #25)

gym_mupen64plus/envs/MarioKart64/mario_kart_env.py

@@ -1,5 +1,6 @@
 import abc
 import inspect
+import itertools
 import os
 import yaml
 from termcolor import cprint
@@ -51,7 +52,9 @@ def _reset(self):
             if self.episode_over:
                 self._act(ControllerState.NO_OP, count=59)
                 self._navigate_post_race_menu()
-                self._act(ControllerState.NO_OP, count=61) # Wait for race to load
+                self._act(ControllerState.NO_OP, count=40) # Wait for the map select screen
+                self._navigate_map_select()
+                self._act(ControllerState.NO_OP, count=50) # Wait for race to load
                 self.episode_over = False
             else:
                 self.controller_server.send_controls(ControllerState.NO_OP, start_button=1)
@@ -212,32 +215,40 @@ def _navigate_map_select(self):
         while cur_col != self.MAP_SERIES:
             self._press_button(ControllerState.JOYSTICK_RIGHT)
             cur_col += 1
-
         self._press_button(ControllerState.A_BUTTON)
 
         # Select map choice
         while cur_row != self.MAP_CHOICE:
             self._press_button(ControllerState.JOYSTICK_DOWN)
             cur_row += 1
-
         self._press_button(ControllerState.A_BUTTON)
 
         # Press OK
         self._press_button(ControllerState.A_BUTTON)
 
     def _navigate_post_race_menu(self):
+        # Times screen
+        self._press_button(ControllerState.A_BUTTON)
+        self._act(ControllerState.NO_OP, count=13)
+
         # Post race menu (previous choice selected by default)
         # - Retry
         # - Course Change
         # - Driver Change
         # - Quit
         # - Replay
         # - Save Ghost
-        self._press_button(ControllerState.A_BUTTON)
-        self._act(ControllerState.NO_OP, count=13)
-        self._press_button(ControllerState.A_BUTTON)
-
 
+        # Because the previous choice is selected by default, we navigate to the top entry so our
+        # navigation is consistent. The menu doesn't cycle top to bottom or bottom to top, so we can
+        # just make sure we're at the top by hitting up a few times
+        for _ in itertools.repeat(None, 5):
+            self._press_button(ControllerState.JOYSTICK_UP)
+
+        # Now we are sure to have the top entry selected
+        # Go down to 'course change'
+        self._press_button(ControllerState.JOYSTICK_DOWN)
+        self._press_button(ControllerState.A_BUTTON)
 
     def _set_character(self, character):
         characters = {'mario'  : (0, 0),

gym_mupen64plus/envs/mupen64plus_env.py

@@ -3,6 +3,7 @@
 import abc
 import array
 import inspect
+import itertools
 import json
 import os
 import subprocess
@@ -79,7 +80,7 @@ def _step(self, action):
         return obs, reward, self.episode_over, {}
 
     def _act(self, action, count=1):
-        for i in range(count):
+        for _ in itertools.repeat(None, count):
             self.controller_server.send_controls(action)
 
     def _press_button(self, button):
@@ -255,8 +256,6 @@ def _start_emulator(self,
         self.mss_grabber = mss.mss()
         time.sleep(2) # Give mss a couple seconds to initialize; also may not be necessary
 
-        time.sleep(30) # TODO: Remove... added for testing (give time to connect with VNC client)
-
         # Restore the DISPLAY env var
         os.environ["DISPLAY"] = initial_disp
         cprint('Changed back to DISPLAY %s' % os.environ["DISPLAY"], 'red')