From 9133ecd61beb8cf4ed7036a1a7e86041d41ee4aa Mon Sep 17 00:00:00 2001
From: Dominik Roth <dominik.roth.dev@gmail.com>
Date: Sun, 17 Jul 2022 00:48:17 +0200
Subject: [PATCH] Show confidence-ellipsoid for supported envs

---
 replay.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/replay.py b/replay.py
index f4554fc..c135b9f 100755
--- a/replay.py
+++ b/replay.py
@@ -28,6 +28,8 @@ def main(load_path, n_eval_episodes=0):
 
     model = PPO.load(load_path, env=env)
 
+    show_chol = env_name.startswith('Columbus')
+
     if n_eval_episodes:
         mean_reward, std_reward = evaluate_policy(
             model, env, n_eval_episodes=n_eval_episodes, deterministic=False)
@@ -41,7 +43,10 @@ def main(load_path, n_eval_episodes=0):
         time.sleep(1/30)
         action, _ = model.predict(obs, deterministic=False)
         obs, reward, done, info = env.step(action)
-        env.render()
+        if show_chol:
+            env.render(chol=model.policy.chol)
+        else:
+            env.render()
         episode_reward += reward
         if done:
             episode_reward = 0.0