Updated record_video.py to truncate recording upon episode end. File still saves based on steps specified, however...

agents/record_video.py

@@ -3,12 +3,13 @@ from stable_baselines3 import DQN
 from stable_baselines3.common.monitor import Monitor
 from stable_baselines3.common.vec_env import VecVideoRecorder, DummyVecEnv, VecEnv
 
+model_name = "dqn_v2-5/ALE-Pacman-v5" # path to model, should be an argument
 env_id = "ALE/Pacman-v5"
 video_folder = "videos/"
-video_length = 1000 #steps
+video_length = 100000 #steps
 
 vec_env = DummyVecEnv([lambda: gym.make(env_id, render_mode="rgb_array")])
-model = DQN.load("ALE-Pacman-v5")
+model = DQN.load(model_name)
 # output: <stable_baselines3.common.vec_env.dummy_vec_env.DummyVecEnv object at 0x0000029974DC6550>
 
 # vec_env = gym.make(env_id, render_mode="rgb_array")
@@ -34,8 +35,19 @@ vec_env = VecVideoRecorder(vec_env,
 #   I want to act according to the policy that has been trained
 obs = vec_env.reset()
 print(vec_env)
-for _ in range(video_length + 1):
+# for _ in range(video_length + 1):
+#   action, states = model.predict(obs)
+#   obs, _, _, _ = vec_env.step(action)
+
+# Instead of using the specified steps in a for loop
+# use a while loop to check if the episode has terminated
+# Stop recording when the episode ends
+end = True
+while end == True:
   action, states = model.predict(obs)
-  obs, _, _, _ = vec_env.step(action)
+  obs, _, done, _ = vec_env.step(action)
+  if done == True:
+    print("exiting loop")
+    end = False
 # # Save the video
 vec_env.close()

agents/videos/rl-video-step-0-to-step-10000.meta.json

@@ -0,0 +1 @@
+{"step_id": 0, "content_type": "video/mp4"}

agents/videos/rl-video-step-0-to-step-10000.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1ca60ce100e54e68edd9e76c84bb7f1bfd6730b2ddea8090b7ef0c1ba675e592
+size 149189

agents/videos/rl-video-step-0-to-step-100000.meta.json

@@ -0,0 +1 @@
+{"step_id": 0, "content_type": "video/mp4"}

agents/videos/rl-video-step-0-to-step-100000.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d433123f98e1de6186e713a1890a2ddfc9f2330ace545f5f38b73d3d03f3a0c8
+size 174076

agents/videos/rl-video-step-0-to-step-200.meta.json

@@ -0,0 +1 @@
+{"step_id": 0, "content_type": "video/mp4"}

agents/videos/rl-video-step-0-to-step-200.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:07ba166ed637c797aef7bc6017b724ae5162ffbb264b57e4f0515140d0e15463
+size 33296