araffin / rl-tutorial-jnrr19

Stable-Baselines tutorial for Journées Nationales de la Recherche en Robotique 2019
MIT License
591 stars 113 forks source link

1_getting_started.ipynb: record_video() uses hardcoded environment string #4

Closed rbahumi closed 4 years ago

rbahumi commented 4 years ago

Hi, Sorry for not making a proper pull request but I think it is better this way in the case of ipython notebooks.

Issue: in the function record_video() a parameter 'env_id' is given but the function uses the hardcoded string 'CartPole-v1'. Here is my small modification:

gym.make('CartPole-v1') => gym.make(env_id)

Below is the full function code with this change.

def record_video(env_id, model, video_length=500, prefix='', video_folder='videos/'):
  """
  :param env_id: (str)
  :param model: (RL model)
  :param video_length: (int)
  :param prefix: (str)
  :param video_folder: (str)
  """
  eval_env = DummyVecEnv([lambda: gym.make(env_id)])
  # Start the video at step=0 and record 500 steps
  eval_env = VecVideoRecorder(env, video_folder=video_folder,
                              record_video_trigger=lambda step: step == 0, video_length=video_length,
                              name_prefix=prefix)

  obs = eval_env.reset()
  for _ in range(video_length):
    action, _ = model.predict(obs)
    obs, _, _, _ = eval_env.step(action)

  # Close the video recorder
  eval_env.close()
araffin commented 4 years ago

Nice Catch, thanks =)