Merge pull request dnv-opensource#4 from dnv-opensource/eis

Fixed the tests so that pytest can be run on the package. Added the p…

Author: aleksandarbabicdnv

scripts/play_q.py

@@ -15,7 +15,9 @@
 def main():
     parser = argparse.ArgumentParser(description="Run a trained Q-learning agent on the crane anti-pendulum task.")
     parser.add_argument("--model-path", type=str, required=True, help="Path to a trained Q-table JSON")
-    parser.add_argument("--render-mode", type=str, default="plot", help="Render mode (plot, play-back, reward-tracking)")
+    parser.add_argument(
+        "--render-mode", type=str, default="plot", help="Render mode (plot, play-back, reward-tracking)"
+    )
     parser.add_argument("--episodes", type=int, default=1, help="Number of episodes to run")
     parser.add_argument("--v0", type=float, default=-1.0, help="Initial crane speed (negative = stop mode)")
     args = parser.parse_args()

src/crane_controller/algorithm.py

@@ -21,9 +21,6 @@ def __init__(
     ):
         self.env = env
         assert type(self.env).__name__ in AlgorithmAgent.envs, f"Environment {type(self.env).__name__} not listed."
-        # print("ACTION_SPACE.N", env.action_space.n, defaultdict( lambda: np.zeros(env.action_space.n))['xx'])
-        # Q-table: maps (state, action) to expected reward
-        # defaultdict automatically creates entries with zeros for new states
 
         # Track learning progress
         self.training_error: list[float] = []
@@ -75,8 +72,9 @@ def do_strategies(self, max_steps: int = 5000):
                 if steps > max_steps:
                     truncated = True
             res.append(reward)
-        for i, self.strategy in enumerate(product(range(3), range(3), range(3), range(3))):
-            print(f"{i}. strategy {self.strategy}: {res[i]}")
+        if not self.env.render_mode == 'none':
+            for i, self.strategy in enumerate(product(range(3), range(3), range(3), range(3))):
+                print(f"{i}. strategy {self.strategy}: {res[i]}")
 
     def do_episodes(self, n_episodes: int = 1000, show: int = 0, max_steps: int = 1000):
         """Run episodes."""

src/crane_controller/envs/controlled_crane_pendulum.py

@@ -73,7 +73,7 @@ def __init__(
         elif render_mode == "plot":
             self.traces: dict[str, list[float]] = {"c_x": [], "c_v": [], "l_x": [], "l_v": []}
 
-        self.obeservation_space : spaces.Box | spaces.Discrete
+        self.obeservation_space: spaces.Box | spaces.Discrete
         # Observations is a 4-dim np-array with
         # (crane-x, crane-v_x, load-polar-angle_x, load-v_x)
         self.min_speed = 0.1  # np.sqrt(2*reward_limit) # starting with less does not make sense (goal already reached)
@@ -98,7 +98,8 @@ def __init__(
         self.dt = dt
 
         # We have 1 acceleration action which can each be min, zero or max, corresponding to acceleration of crane
-        self.action_space = spaces.Discrete(3, start=0, seed=42, dtype=np.int64)
+        super().reset(seed=seed)  # make sure that the environment seed is set
+        self.action_space = spaces.Discrete(3, start=0, dtype=np.int64)
         self.action_to_acc = {0: -self.acc, 1: 0.0, 2: self.acc}
 
     def _init_discrete(self, spec: dict[str, tuple[float, ...]]):
@@ -211,7 +212,7 @@ def level(idx: int, val: float, categories: tuple[float, ...]):
                 # if the crane moves towards the origo we do not add 'energy'
         self.reward = reward
 
-        obs : tuple[int,...] | np.ndarray
+        obs: tuple[int, ...] | np.ndarray
         if len(self.discrete):
             obs = (
                 level(1, energy, self.discrete["energies"]),  # energy level