Merge pull request dnv-opensource#2 from aleksandarbabicdnv/refactor/split-test-files

Refactor/split test files

Author: eisDNV

README.rst

@@ -70,14 +70,91 @@ Installation
 Running
 -------
 
-Training and evaluation are driven through the test suite (run with ``pytest``) or by executing
-the test files directly as scripts. The main test file is ``tests/test_crane_pendulum.py``.
+Install dependencies and run the test suite with ``uv``:
 
 .. code-block:: shell
 
-   pytest tests/test_crane_pendulum.py -v
+   uv run pytest tests/ -v
 
-Saved models (PPO ``.zip``) and Q-tables (``.json``) are stored in the ``tests/`` directory.
+Test files are organised by algorithm:
+
+- ``tests/test_crane_pendulum.py`` — environment, Q-learning, and algorithm tests
+- ``tests/test_ppo.py`` — PPO pipeline smoke test (``test_monitor``)
+
+Tests are suitable for CI/CD — no plot windows are produced.
+
+Training
+--------
+
+**PPO:**
+
+.. code-block:: shell
+
+   uv run python scripts/train_ppo.py
+
+Key options:
+
+- ``--steps N`` — total training timesteps (default: 100 000)
+- ``--n-envs N`` — number of parallel environments (default: 4)
+- ``--save-path PATH`` — where to write the trained model (default: ``models/ppo_AntiPendulumEnv.zip``)
+- ``--dry-run`` — run 1 000 steps with a live reward-tracking plot and no model saved
+
+**Q-learning:**
+
+.. code-block:: shell
+
+   uv run python scripts/train_q.py
+
+Key options:
+
+- ``--episodes N`` — total training episodes (default: 10 000)
+- ``--v0 F`` — initial crane speed; negative = stop mode, positive = start mode (default: ``-1.0``)
+- ``--reward-limit F`` — per-episode termination threshold (default: ``-0.05``)
+- ``--save-path PATH`` — where to write the Q-table (default: ``models/q_AntiPendulumEnv.json``)
+- ``--trained PATH`` — continue training from an existing Q-table JSON
+- ``--intervals N`` — run interval training: N rounds of 10 episodes each
+- ``--dry-run`` — run 50 episodes with a reward plot and no model saved
+
+Playing
+-------
+
+Run a trained agent visually. Both scripts accept ``--render-mode`` with the following options:
+
+- ``plot`` — 4-panel figure per episode (load angle, crane position/speed, rewards)
+- ``play-back`` — animated crane trajectory after each episode
+- ``reward-tracking`` — live reward line plot updating every step
+
+**PPO** (default render-mode: ``play-back``):
+
+.. code-block:: shell
+
+   uv run python scripts/play_ppo.py --model-path models/ppo_AntiPendulumEnv.zip
+   uv run python scripts/play_ppo.py --model-path models/ppo_AntiPendulumEnv.zip --render-mode plot --episodes 3
+
+**Q-learning** (default render-mode: ``plot``):
+
+.. code-block:: shell
+
+   uv run python scripts/play_q.py --model-path models/q_AntiPendulumEnv.json
+   uv run python scripts/play_q.py --model-path tests/anti-pendulum.json --render-mode play-back --episodes 3
+
+Analysing
+---------
+
+Inspect a trained Q-table without running the environment:
+
+.. code-block:: shell
+
+   uv run python scripts/analyse_q.py --model-path tests/anti-pendulum.json
+
+Prints per-pos/speed average Q-values for a quick sanity check. To drill into
+specific states, use ``--obs`` with 5 integers (use ``-1`` as a wildcard):
+
+.. code-block:: shell
+
+   uv run python scripts/analyse_q.py --model-path tests/anti-pendulum.json --obs -1 0 0 -1 -1
+
+The five observation dimensions are: ``[energy, pos, speed, distance, sector]``.
 
 Contributing
 ------------

models/.gitkeep

@@ -0,0 +1,62 @@
+"""Inspect a trained Q-table for the AntiPendulumEnv.
+
+Prints a per-pos/speed average summary by default. Use --obs to drill
+into specific states; negative values act as wildcards (match any).
+
+The observation tuple has 5 dimensions:
+  [energy, pos, speed, distance, sector]
+
+Example:
+    uv run python scripts/analyse_q.py --model-path tests/anti-pendulum.json
+    uv run python scripts/analyse_q.py --model-path tests/anti-pendulum.json --obs -1 0 0 -1 -1
+    uv run python scripts/analyse_q.py --model-path tests/anti-pendulum.json --obs -1 1 1 -1 -1
+"""
+
+import argparse
+
+import numpy as np
+
+from crane_controller.crane_factory import build_crane
+from crane_controller.envs.controlled_crane_pendulum import AntiPendulumEnv
+from crane_controller.q_agent import QLearningAgent
+
+
+def _build_dummy_env():
+    """Minimal env needed to satisfy QLearningAgent constructor (action_space.n)."""
+    return AntiPendulumEnv(build_crane, discrete=QLearningAgent.DEFAULT_DISCRETE.copy())
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Inspect a trained Q-table for 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(
+        "--obs",
+        type=int,
+        nargs=5,
+        metavar=("ENERGY", "POS", "SPEED", "DISTANCE", "SECTOR"),
+        help="Filter Q-values for a specific observation (use -1 as wildcard)",
+    )
+    args = parser.parse_args()
+
+    env = _build_dummy_env()
+    agent = QLearningAgent(env, trained=(args.model_path, True))
+
+    print(f"Q-table: {len(agent.q_values)} states  ({args.model_path})")
+    print()
+
+    if args.obs:
+        print(f"Q-values matching obs {args.obs}  (columns: state, q-values, best-action, mean, cv)")
+        print("-" * 72)
+        agent.analyse_q(tuple(args.obs))
+    else:
+        print("Per-pos/speed average Q-values  (actions: left=0, coast=1, right=2)")
+        print("-" * 52)
+        for pos in (0, 1):
+            for speed in (0, 1):
+                res = {k: v for k, v in agent.q_values.items() if k[1] == pos and k[2] == speed}
+                avgs = [np.average([x[i] for x in res.values()]) for i in range(3)]
+                print(f"  pos={pos}  speed={speed}  ->  {[f'{a:.4f}' for a in avgs]}")
+
+
+if __name__ == "__main__":
+    main()

scripts/analyse_q.py

@@ -0,0 +1,39 @@
+"""Run a trained PPO agent on the AntiPendulumEnv.
+
+Example:
+    uv run python scripts/play_ppo.py --model-path models/ppo_AntiPendulumEnv.zip
+    uv run python scripts/play_ppo.py --model-path models/ppo.zip --render-mode plot --episodes 3
+"""
+
+import argparse
+
+from crane_controller.crane_factory import build_crane
+from crane_controller.envs.controlled_crane_pendulum import AntiPendulumEnv
+from crane_controller.ppo_agent import ProximalPolicyOptimizationAgent
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Run a trained PPO agent on the crane anti-pendulum task.")
+    parser.add_argument("--model-path", type=str, required=True, help="Path to a trained .zip model")
+    parser.add_argument("--render-mode", type=str, default="play-back", help="Render mode for playback")
+    parser.add_argument("--episodes", type=int, default=1, help="Number of episodes to run")
+    args = parser.parse_args()
+
+    agent = ProximalPolicyOptimizationAgent(
+        AntiPendulumEnv,  # type: ignore[arg-type]
+        n_envs=0,  # load-from-file mode
+        env_kwargs={
+            "crane": build_crane,
+            "start_speed": 1.0,
+            "render_mode": args.render_mode,
+        },
+        trained=(args.model_path, True),
+    )
+
+    for episode in range(args.episodes):
+        print(f"Episode {episode + 1}/{args.episodes}")
+        agent.do_one_episode(seed=episode + 1)
+
+
+if __name__ == "__main__":
+    main()

scripts/play_ppo.py

@@ -0,0 +1,38 @@
+"""Run a trained Q-learning agent on the AntiPendulumEnv.
+
+Example:
+    uv run python scripts/play_q.py --model-path models/q_AntiPendulumEnv.json
+    uv run python scripts/play_q.py --model-path tests/anti-pendulum.json --render-mode plot --episodes 3
+"""
+
+import argparse
+
+from crane_controller.crane_factory import build_crane
+from crane_controller.envs.controlled_crane_pendulum import AntiPendulumEnv
+from crane_controller.q_agent import QLearningAgent
+
+
+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("--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()
+
+    env = AntiPendulumEnv(
+        build_crane,
+        start_speed=args.v0,
+        render_mode=args.render_mode,
+        discrete=QLearningAgent.DEFAULT_DISCRETE.copy(),
+    )
+    agent = QLearningAgent(env, trained=(args.model_path, True))
+
+    for episode in range(args.episodes):
+        print(f"Episode {episode + 1}/{args.episodes}")
+        env.reset(seed=episode + 1)
+        agent.do_episodes(n_episodes=1)
+
+
+if __name__ == "__main__":
+    main()

scripts/play_q.py

@@ -0,0 +1,64 @@
+"""Train a PPO agent on the AntiPendulumEnv.
+
+Example:
+    uv run python scripts/train_ppo.py
+    uv run python scripts/train_ppo.py --steps 500000 --n-envs 8 --save-path models/ppo.zip
+    uv run python scripts/train_ppo.py --dry-run
+"""
+
+import argparse
+from pathlib import Path
+
+from crane_controller.crane_factory import build_crane
+from crane_controller.envs.controlled_crane_pendulum import AntiPendulumEnv
+from crane_controller.ppo_agent import ProximalPolicyOptimizationAgent
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Train a PPO agent on the crane anti-pendulum task.")
+    parser.add_argument("--steps", type=int, default=100_000, help="Total training timesteps")
+    parser.add_argument("--n-envs", type=int, default=4, help="Number of parallel environments")
+    parser.add_argument("--render-mode", type=str, default="none", help="Render mode during training")
+    parser.add_argument(
+        "--save-path",
+        type=str,
+        default="models/ppo_AntiPendulumEnv.zip",
+        help="Where to save the trained model",
+    )
+    parser.add_argument(
+        "--dry-run",
+        action="store_true",
+        help="Run 1000 steps with live reward-tracking plot, without saving the model.",
+    )
+    args = parser.parse_args()
+
+    if args.dry_run:
+        agent = ProximalPolicyOptimizationAgent(
+            AntiPendulumEnv,  # type: ignore[arg-type]
+            n_envs=1,
+            env_kwargs={
+                "crane": build_crane,
+                "start_speed": -1.0,
+                "render_mode": "reward-tracking",
+            },
+            trained=None,
+        )
+        agent.do_training(1000, progress_bar=False)
+    else:
+        Path(args.save_path).parent.mkdir(parents=True, exist_ok=True)
+        agent = ProximalPolicyOptimizationAgent(
+            AntiPendulumEnv,  # type: ignore[arg-type]
+            n_envs=args.n_envs,
+            env_kwargs={
+                "crane": build_crane,
+                "start_speed": 1.0,
+                "render_mode": args.render_mode,
+            },
+            trained=(args.save_path, True),
+        )
+        agent.do_training(args.steps)
+        print(f"Model saved to {args.save_path}")
+
+
+if __name__ == "__main__":
+    main()

scripts/train_ppo.py

@@ -0,0 +1,75 @@
+"""Train a Q-learning agent on the AntiPendulumEnv.
+
+Example:
+    uv run python scripts/train_q.py
+    uv run python scripts/train_q.py --episodes 50000 --v0 1.0 --save-path models/q_start.json
+    uv run python scripts/train_q.py --trained models/q_AntiPendulumEnv.json
+    uv run python scripts/train_q.py --intervals 10
+    uv run python scripts/train_q.py --dry-run
+"""
+
+import argparse
+from pathlib import Path
+
+from crane_controller.crane_factory import build_crane
+from crane_controller.envs.controlled_crane_pendulum import AntiPendulumEnv
+from crane_controller.q_agent import QLearningAgent
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Train a Q-learning agent on the crane anti-pendulum task.")
+    parser.add_argument("--episodes", type=int, default=10_000, help="Total training episodes")
+    parser.add_argument("--v0", type=float, default=-1.0, help="Initial crane speed (negative = stop mode)")
+    parser.add_argument("--reward-limit", type=float, default=-0.05, help="Per-episode reward termination threshold")
+    parser.add_argument(
+        "--save-path",
+        type=str,
+        default="models/q_AntiPendulumEnv.json",
+        help="Where to save the trained Q-table",
+    )
+    parser.add_argument("--trained", type=str, default=None, help="Path to an existing Q-table JSON to continue from")
+    parser.add_argument(
+        "--intervals",
+        type=int,
+        default=0,
+        help="Run interval training: N intervals of 10 episodes each (0 = disabled)",
+    )
+    parser.add_argument(
+        "--dry-run",
+        action="store_true",
+        help="Run 50 episodes with a reward plot and no model saved, for a quick visual sanity check.",
+    )
+    args = parser.parse_args()
+
+    env = AntiPendulumEnv(
+        build_crane,
+        start_speed=args.v0,
+        render_mode="plot" if args.dry_run else "none",
+        reward_limit=args.reward_limit,
+        discrete=QLearningAgent.DEFAULT_DISCRETE.copy(),
+    )
+
+    if args.dry_run:
+        agent = QLearningAgent(env, trained=None)
+        agent.do_episodes(n_episodes=50, max_steps=1000)
+
+    elif args.intervals > 0:
+        Path(args.save_path).parent.mkdir(parents=True, exist_ok=True)
+        agent = QLearningAgent(env, trained=(args.save_path, False))
+        for i in range(args.intervals):
+            env.reset(seed=i + 1)
+            agent.do_episodes(n_episodes=10)
+            if i == 0:
+                agent = QLearningAgent(env, trained=(args.save_path, True))
+        print(f"Model saved to {args.save_path}")
+
+    else:
+        Path(args.save_path).parent.mkdir(parents=True, exist_ok=True)
+        trained = (args.trained, True) if args.trained else (args.save_path, False)
+        agent = QLearningAgent(env, trained=trained)
+        agent.do_episodes(n_episodes=args.episodes, max_steps=5000)
+        print(f"Model saved to {args.save_path}")
+
+
+if __name__ == "__main__":
+    main()