code: def test_generate(tmp_path, generate_env):
      model, policy, env_name, n_env, n_episodes = generate_env
  
      if n_env > 1:
          env = make_atari_env(env_name, num_env=n_env, seed=0)
          model = model(policy, env, verbose=0)
      else:
          model = model(policy, env_name, verbose=0)
  
      dataset = generate_expert_traj(
          model,
          str(tmp_path / "expert"),
          n_timesteps=300,
          n_episodes=n_episodes,
          image_folder=str(tmp_path / "test_recorded_images"),
      )
  
      assert set(dataset.keys()).issuperset(
          [
              "actions",
              "obs",
              "rewards",
              "episode_returns",
              "episode_starts",
          ]
      )
      assert sum(dataset["episode_starts"]) == n_episodes
      assert len(dataset["episode_returns"]) == n_episodes
      n_timesteps = len(dataset["episode_starts"])
      for key, val in dataset.items():
          if key != "episode_returns":
              assert (
                  val.shape[0] == n_timesteps
              ), "inconsistent number of timesteps at '{}'".format(key)
  
      dataset_loaded = np.load(
          str(tmp_path / "expert.npz"), allow_pickle=True
      )
      assert dataset.keys() == dataset_loaded.keys()
      for key in dataset.keys():
          assert (
              dataset[key] == dataset_loaded[key]
          ).all(), "different data at '{}'".format(key)
      # Cleanup folder
      if os.path.isdir(str(tmp_path / "test_recorded_images")):
          shutil.rmtree(str(tmp_path / "test_recorded_images"))
  
imports: import os
  import shutil
  import gym
  import numpy
  import pytest
  from stable_baselines import A2C, ACER, ACKTR, GAIL, DDPG, DQN, PPO1, PPO2, TD3, TRPO, SAC
  from stable_baselines.common.cmd_util import make_atari_env
  from stable_baselines.common.vec_env import VecFrameStack, DummyVecEnv
  from stable_baselines.common.evaluation import evaluate_policy
  from stable_baselines.common.callbacks import CheckpointCallback
  from stable_baselines.gail import ExpertDataset, generate_expert_traj
label_desc: Use image node_id 2 for calling a global function with example usage: test_generate(tmp_path, generate_env) without return types