| 24 | return all_args |
| 25 | |
| 26 | def test_videos(args): |
| 27 | all_args = parse(args) |
| 28 | |
| 29 | random.seed(all_args.seed) |
| 30 | |
| 31 | run_dir = Path("../results") / all_args.project_name / all_args.experiment_name |
| 32 | if not run_dir.exists(): |
| 33 | os.makedirs(str(run_dir)) |
| 34 | |
| 35 | wandb.init(config=all_args, |
| 36 | project=all_args.project_name, |
| 37 | entity=all_args.team_name, |
| 38 | notes=socket.gethostname(), |
| 39 | name=all_args.experiment_name+"_"+str(all_args.seed), |
| 40 | group=all_args.scenario_name, |
| 41 | dir=str(run_dir), |
| 42 | job_type="training", |
| 43 | reinit=True) |
| 44 | |
| 45 | env = gym.make("PongNoFrameskip-v4") |
| 46 | for episode in range(3): |
| 47 | env.reset() |
| 48 | done = False |
| 49 | frames = [] |
| 50 | while not done: |
| 51 | for _ in range(4): |
| 52 | obs,r,done,_=env.step(env.action_space.sample()) |
| 53 | if done: |
| 54 | break |
| 55 | frames.append(obs) |
| 56 | sequence = np.stack(frames, -1).transpose(3,2,0,1) # time, channels, height, width |
| 57 | print(sequence.shape) |
| 58 | video = wandb.Video(sequence, fps=10, format="gif",caption="Pong") |
| 59 | wandb.log({"video": video},step=episode) |
| 60 | |
| 61 | |
| 62 | if __name__ == '__main__': |