Add tests for dtype fix and VecNormalize

test_observation_space_dtype: asserts the continuous observation space
uses float64, directly documenting the bug fixed in this PR (was int64,
causing float observations to be silently truncated to integers).

test_observations_are_float: confirms sub-integer precision is preserved
end-to-end after a physics step — the case that was broken before.

test_ppo_saves_vecnorm: verifies that do_training() persists the
VecNormalize statistics as <model>_vecnorm.pkl alongside the model zip,
so inference can reload the same normalisation.

test_ppo_vecnorm_updates: confirms the running mean is updated during
training (normaliser is actually learning from observations, not stuck
at zero).

Author: aleksandarbabicdnv

tests/test_environment.py

@@ -76,6 +76,19 @@ def test_environment(crane: Callable, show: bool, v0: float = 1.0, reward_limit=
     assert q_values[obs2.tobytes()][2] == -98.1
 
 
+def test_observation_space_dtype(crane: Callable):
+    env = AntiPendulumEnv(crane)
+    assert env.observation_space.dtype == np.float64
+
+
+def test_observations_are_float(crane: Callable):
+    env = AntiPendulumEnv(crane)
+    env.reset()
+    obs, _, _, _, _ = env.step(1)  # one physics step produces fractional values
+    assert obs.dtype == np.float64
+    assert not np.all(obs == obs.astype(int))  # sub-integer precision is preserved
+
+
 def test_init(crane: Crane, show: bool = False):
     """Test the initialization of the environment."""
     env = AntiPendulumEnv(crane, seed=1, start_speed=1.0, render_mode="play-back" if show else "data")

tests/test_ppo.py

@@ -1,5 +1,7 @@
 import logging
+from pathlib import Path
 
+import numpy as np
 from py_crane.crane import Crane
 
 from crane_controller.envs.controlled_crane_pendulum import AntiPendulumEnv
@@ -20,3 +22,25 @@ def test_monitor(crane: Crane, show: bool):
         },
     )
     agent.do_training(1000)
+
+
+def test_ppo_saves_vecnorm(crane, tmp_path):
+    save_path = str(tmp_path / "model.zip")
+    agent = ProximalPolicyOptimizationAgent(
+        AntiPendulumEnv,  # type: ignore[arg-type]
+        n_envs=1,
+        env_kwargs={"crane": crane, "start_speed": 1.0},
+        trained=(save_path, True),
+    )
+    agent.do_training(500, progress_bar=False)
+    assert (tmp_path / "model_vecnorm.pkl").exists()
+
+
+def test_ppo_vecnorm_updates(crane):
+    agent = ProximalPolicyOptimizationAgent(
+        AntiPendulumEnv,  # type: ignore[arg-type]
+        n_envs=1,
+        env_kwargs={"crane": crane, "start_speed": 1.0},
+    )
+    agent.do_training(500, progress_bar=False)
+    assert not np.allclose(agent.vec_env.obs_rms.mean, 0.0)