class_docstr: Wrapper of the procgen env.
class_docstr_tok: ["Wrapper","of","the","procgen","env","."]
code: def __init__(
      self,
      env_name: str,
      num_envs: int = 1024,
      n_repeats: int = 1,
      num_levels: int = 200,
      distribution_mode: str = "easy",
      max_steps: int = 1000,
      num_threads: int = -1,
      test: bool = False,
  ):
      self.max_steps = max_steps
      self.test = test
  
      try:
          from procgen import ProcgenEnv
      except ModuleNotFoundError:
          print("You need to install procgen for this task:")
          print("  pip install procgen")
          sys.exit(1)
  
      procgen_venv = ProcgenEnv(
          num_envs=4,
          env_name=env_name,
          num_levels=0 if test else num_levels,
          start_level=num_levels if test else 0,
          distribution_mode=distribution_mode,
          num_threads=num_threads
          if num_threads > 0
          else mp.cpu_count(),
      )
      obs_space = procgen_venv.observation_space["rgb"]
      act_space = procgen_venv.action_space
      self.procgen_venv = procgen_venv
      self.obs_shape = obs_space.shape
      self.act_shape = act_space.shape
      self.num_acts = act_space.n
      assert num_envs % n_repeats == 0
      pop_size = num_envs // n_repeats
  
      def reset_fn(key):
          rand_seed = int(
              jax.random.randint(
                  key[0], shape=(), minval=0, maxval=2**31 - 1
              )
          )
          self.procgen_venv = ProcgenEnv(
              num_envs=num_envs,
              env_name=env_name,
              num_levels=0 if test else num_levels,
              start_level=num_levels if test else 0,
              distribution_mode=distribution_mode,
              num_threads=num_threads
              if num_threads > 0
              else mp.cpu_count(),
              rand_seed=rand_seed,
              use_backgrounds=True,
          )
          if not self.test:
              state = self.procgen_venv.env.get_state()
              assert len(state) == num_envs, "{} vs {}".format(
                  len(state), num_envs
              )
              ss = [
                  x for _ in range(pop_size) for x in state[:n_repeats]
              ]
              self.procgen_venv.env.set_state(ss)
  
          obs, _, _, _ = self.procgen_venv.step(np.zeros(num_envs))
          return State(obs=to_jnp_array(obs["rgb"]))
  
      self._reset_fn = reset_fn
  
      def step_fn(state, action):
          action_np = np.array(action)
          obs, reward, done, info = self.procgen_venv.step(action_np)
          return (
              State(obs=to_jnp_array(obs["rgb"])),
              to_jnp_array(reward),
              to_jnp_array(done),
          )
  
      self._step_fn = step_fn
  
file_path: evojax/evojax/task/procgen_task.py
imports: import sys
  import multiprocessing
  import numpy
  from typing import Tuple
  import jax
  import jax.numpy
  from flax.struct import dataclass
  from evojax.task.base import VectorizedTask
  from evojax.task.base import TaskState
label_desc: Use image node_id 1 to create a new ProcgenTask object from inherited base classes: VectorizedTask with example: obj = ProcgenTask(env_name, num_envs, n_repeats, num_levels, distribution_mode, max_steps, num_threads, test)