Improve performance

README.md

@@ -37,16 +37,16 @@ Testing Lunar Lander :
 
 ![lunarlander training preview](https://github.com/werner-duvaud/muzero-general/blob/master/docs/lunarlander_training_preview.png)
 
-## Code structure
-
-![code structure](https://github.com/werner-duvaud/muzero-general/blob/master/docs/how-it-works-werner-duvaud.png)
-
 ## Games already implemented with pretrained network available
 
 * Cartpole
 * Lunar Lander
 * Connect4
 
+## Code structure
+
+![code structure](https://github.com/werner-duvaud/muzero-general/blob/master/docs/how-it-works-werner-duvaud.png)
+
 ## Getting started
 ### Installation
 

games/lunarlander.py

@@ -67,16 +67,14 @@ class MuZeroConfig:
         Returns:
             Positive float.
         """
-        # if trained_steps < 0.2 * self.training_steps:
-        #     return float('inf')
-        # if trained_steps < 0.5 * self.training_steps:
-        #     return 0.8
-        # elif trained_steps < 0.75 * self.training_steps:
-        #     return 0.5
-        # else:
-        #     return 0.25
-
-        return 1
+        if trained_steps < 0.2 * self.training_steps:
+            return float('inf')
+        if trained_steps < 0.5 * self.training_steps:
+            return 0.8
+        elif trained_steps < 0.75 * self.training_steps:
+            return 0.5
+        else:
+            return 0.25
 
 
 class Game: