ledmands commited on
Commit
a2ff203
1 Parent(s): 1427095

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

Browse files
agents/record_video.py CHANGED
@@ -3,12 +3,13 @@ from stable_baselines3 import DQN
3
  from stable_baselines3.common.monitor import Monitor
4
  from stable_baselines3.common.vec_env import VecVideoRecorder, DummyVecEnv, VecEnv
5
 
 
6
  env_id = "ALE/Pacman-v5"
7
  video_folder = "videos/"
8
- video_length = 1000 #steps
9
 
10
  vec_env = DummyVecEnv([lambda: gym.make(env_id, render_mode="rgb_array")])
11
- model = DQN.load("ALE-Pacman-v5")
12
  # output: <stable_baselines3.common.vec_env.dummy_vec_env.DummyVecEnv object at 0x0000029974DC6550>
13
 
14
  # vec_env = gym.make(env_id, render_mode="rgb_array")
@@ -34,8 +35,19 @@ vec_env = VecVideoRecorder(vec_env,
34
  # I want to act according to the policy that has been trained
35
  obs = vec_env.reset()
36
  print(vec_env)
37
- for _ in range(video_length + 1):
 
 
 
 
 
 
 
 
38
  action, states = model.predict(obs)
39
- obs, _, _, _ = vec_env.step(action)
 
 
 
40
  # # Save the video
41
  vec_env.close()
 
3
  from stable_baselines3.common.monitor import Monitor
4
  from stable_baselines3.common.vec_env import VecVideoRecorder, DummyVecEnv, VecEnv
5
 
6
+ model_name = "dqn_v2-5/ALE-Pacman-v5" # path to model, should be an argument
7
  env_id = "ALE/Pacman-v5"
8
  video_folder = "videos/"
9
+ video_length = 100000 #steps
10
 
11
  vec_env = DummyVecEnv([lambda: gym.make(env_id, render_mode="rgb_array")])
12
+ model = DQN.load(model_name)
13
  # output: <stable_baselines3.common.vec_env.dummy_vec_env.DummyVecEnv object at 0x0000029974DC6550>
14
 
15
  # vec_env = gym.make(env_id, render_mode="rgb_array")
 
35
  # I want to act according to the policy that has been trained
36
  obs = vec_env.reset()
37
  print(vec_env)
38
+ # for _ in range(video_length + 1):
39
+ # action, states = model.predict(obs)
40
+ # obs, _, _, _ = vec_env.step(action)
41
+
42
+ # Instead of using the specified steps in a for loop
43
+ # use a while loop to check if the episode has terminated
44
+ # Stop recording when the episode ends
45
+ end = True
46
+ while end == True:
47
  action, states = model.predict(obs)
48
+ obs, _, done, _ = vec_env.step(action)
49
+ if done == True:
50
+ print("exiting loop")
51
+ end = False
52
  # # Save the video
53
  vec_env.close()
agents/videos/rl-video-step-0-to-step-10000.meta.json ADDED
@@ -0,0 +1 @@
 
 
1
+ {"step_id": 0, "content_type": "video/mp4"}
agents/videos/rl-video-step-0-to-step-10000.mp4 ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:1ca60ce100e54e68edd9e76c84bb7f1bfd6730b2ddea8090b7ef0c1ba675e592
3
+ size 149189
agents/videos/rl-video-step-0-to-step-100000.meta.json ADDED
@@ -0,0 +1 @@
 
 
1
+ {"step_id": 0, "content_type": "video/mp4"}
agents/videos/rl-video-step-0-to-step-100000.mp4 ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:d433123f98e1de6186e713a1890a2ddfc9f2330ace545f5f38b73d3d03f3a0c8
3
+ size 174076
agents/videos/rl-video-step-0-to-step-200.meta.json ADDED
@@ -0,0 +1 @@
 
 
1
+ {"step_id": 0, "content_type": "video/mp4"}
agents/videos/rl-video-step-0-to-step-200.mp4 ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:07ba166ed637c797aef7bc6017b724ae5162ffbb264b57e4f0515140d0e15463
3
+ size 33296