Merge pull request dnv-opensource#3 from aleksandarbabicdnv/fix/ppo-observation-normalization

Fix PPO observation dtype bug and add VecNormalize

Author: eisDNV

scripts/train_ppo.py

@@ -57,7 +57,9 @@ def main():
             trained=(args.save_path, True),
         )
         agent.do_training(args.steps)
+        vecnorm_path = Path(args.save_path).parent / f"{Path(args.save_path).stem}_vecnorm.pkl"
         print(f"Model saved to {args.save_path}")
+        print(f"VecNormalize stats saved to {vecnorm_path}")
 
 
 if __name__ == "__main__":

src/crane_controller/envs/controlled_crane_pendulum.py

@@ -84,7 +84,7 @@ def __init__(
             self.discrete = {}
             self.spaces_min = np.array((-size, -max_speed, -np.pi / 2, -max_speed), float)
             self.spaces_max = np.array((size, max_speed, np.pi / 2, max_speed), float)
-            self.observation_space = spaces.Box(self.spaces_min, self.spaces_max, shape=(4,), dtype=np.int64)
+            self.observation_space = spaces.Box(self.spaces_min, self.spaces_max, shape=(4,), dtype=np.float64)
 
         self.nresets: int = 0
         # self.reset(seed)

src/crane_controller/ppo_agent.py

@@ -1,12 +1,14 @@
 from __future__ import annotations
 
+from pathlib import Path
 from typing import Any
 
 import gymnasium as gym
 import matplotlib.pyplot as plt
 from stable_baselines3 import PPO
 from stable_baselines3.common.env_util import make_vec_env
 from stable_baselines3.common.evaluation import evaluate_policy
+from stable_baselines3.common.vec_env import VecNormalize
 
 plt.rcParams["figure.figsize"] = (10, 5)
 
@@ -17,6 +19,7 @@ class ProximalPolicyOptimizationAgent:
     Initializes an agent that learns a policy via PPO algorithm to solve the task at hand.
 
     PPO agents can be saved as zip file and re-loaded to avoid re-training.
+    VecNormalize statistics are saved alongside the model as `<name>_vecnorm.pkl`.
 
     Args:
         env (gym.Env): the environment the agent is acting on.
@@ -34,43 +37,62 @@ def __init__(
         trained: tuple[str, bool] | None = None,
     ):
         self.trained = trained
-        if env_kwargs is None:
-            self.env = env() # type: ignore[operator]  ## the object is callable! (__init__())
-        else:
-            self.env = env(**env_kwargs) # type: ignore[operator]  ## the object is callable! (__init__())
-        _n_envs = n_envs = 1 if n_envs <= 0 else n_envs
-        self.vec_env = make_vec_env(env_id=env, n_envs=_n_envs, env_kwargs=env_kwargs)  # type: ignore ## should be correct
-        if n_envs <= 0:
-            assert self.trained is not None, "When no training is specified a saved model should be provided"
-            self.model = PPO.load(self.trained[0])
-        elif n_envs == 1:
-            self.model = PPO("MlpPolicy", self.env, verbose=1)
-            if trained is not None:
-                self.trained = (trained[0], trained[1])
+        inference_only = n_envs <= 0
+        _n_envs = 1 if inference_only else n_envs
+
+        raw_vec_env = make_vec_env(env_id=env, n_envs=_n_envs, env_kwargs=env_kwargs)  # type: ignore
+
+        if inference_only:
+            assert trained is not None, "When no training is specified a saved model should be provided"
+            stats_path = self._stats_path(trained[0])
+            if stats_path.exists():
+                self.vec_env = VecNormalize.load(str(stats_path), raw_vec_env)
+            else:
+                self.vec_env = VecNormalize(raw_vec_env, norm_obs=True, norm_reward=False)
+            self.vec_env.training = False
+            self.vec_env.norm_reward = False
+            self.model = PPO.load(trained[0], env=self.vec_env)
         else:
-            self.model = PPO("MlpPolicy", self.vec_env)
+            self.vec_env = VecNormalize(raw_vec_env, norm_obs=True, norm_reward=True)
+            if _n_envs == 1:
+                self.model = PPO("MlpPolicy", self.vec_env, verbose=1)
+            else:
+                self.model = PPO("MlpPolicy", self.vec_env)
             self.trained = (
-                trained[0] if trained is not None else f"ppo_{env.__name__}",  # type: ignore ## has name
+                trained[0] if trained is not None else f"ppo_{env.__name__}",  # type: ignore[attr-defined]
                 False if trained is None else trained[1],
             )
 
+        # Single unwrapped env for do_one_episode/evaluate without reconstructing a new crane.
+        self.env = self.vec_env.venv.envs[0]  # type: ignore[attr-defined]
+
+    @staticmethod
+    def _stats_path(model_path: str) -> Path:
+        p = Path(model_path)
+        return p.parent / f"{p.stem}_vecnorm.pkl"
+
     def do_training(self, total_timesteps: int = 25000, progress_bar: bool = True):
         self.model.learn(total_timesteps, progress_bar=progress_bar)
-        if self.trained is not None and self.trained[1] and self.env.render_mode not in ("play-back"):
+        if self.trained is not None and self.trained[1] and self.env.render_mode not in ("play-back",):
             self.model.save(self.trained[0])
+            self.vec_env.save(str(self._stats_path(self.trained[0])))
 
     def evaluate(self, n_episodes: int = 10):
-        mean_reward, std_reward = evaluate_policy(self.model, self.env, n_eval_episodes=n_episodes)
+        self.vec_env.training = False
+        self.vec_env.norm_reward = False
+        mean_reward, std_reward = evaluate_policy(self.model, self.vec_env, n_eval_episodes=n_episodes)
+        self.vec_env.training = True
+        self.vec_env.norm_reward = True
         print(f"Mean:{mean_reward}, stdev:{std_reward}")
 
     def do_one_episode(self, seed: int = 1):
-        """Do one episode on the non-vectorized, trained environment."""
+        """Do one episode using the trained normalizer for observations."""
         obs, info = self.env.reset(seed=seed)
         terminated = truncated = False
         while not terminated and not truncated:
-            action, _states = self.model.predict(obs)
+            norm_obs = self.vec_env.normalize_obs(obs)
+            action, _states = self.model.predict(norm_obs, deterministic=True)
             obs, rewards, terminated, truncated, info = self.env.step(action)
-            # print("Action", obs, rewards, terminated, truncated, info)
         self.env.render()
 
 

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)