class_docstr: Cart-pole swing up task.
class_docstr_tok: ["Cart-pole","swing","up","task","."]
code: def __init__(
      self,
      max_steps: int = 1000,
      harder: bool = False,
      test: bool = False,
  ):
      self.max_steps = max_steps
      self.obs_shape = tuple(
          [
              5,
          ]
      )
      self.act_shape = tuple(
          [
              1,
          ]
      )
      self.test = test
      if harder:
          get_init_state_fn = get_init_state_hard
      else:
          get_init_state_fn = get_init_state_easy
  
      def reset_fn(key):
          next_key, key = random.split(key)
          state = get_init_state_fn(key)
          return State(
              state=state,
              obs=get_obs(state),
              steps=jnp.zeros((), dtype=int),
              key=next_key,
          )
  
      self._reset_fn = jax.jit(jax.vmap(reset_fn))
  
      def step_fn(state, action):
          cur_state = update_state(action=action, state=state.state)
          reward = get_reward(state=cur_state)
          steps = state.steps + 1
          done = jnp.bitwise_or(
              out_of_screen(cur_state), steps >= max_steps
          )
          steps = jnp.where(done, jnp.zeros((), jnp.int32), steps)
          next_key, key = random.split(state.key)
          cur_state = jax.lax.cond(
              done,
              lambda x: get_init_state_fn(key),
              lambda x: x,
              cur_state,
          )
          return (
              State(
                  state=cur_state,
                  obs=get_obs(state=cur_state),
                  steps=steps,
                  key=next_key,
              ),
              reward,
              done,
          )
  
      self._step_fn = jax.jit(jax.vmap(step_fn))
  
imports: from typing import Tuple
  from PIL import Image
  from PIL import ImageDraw
  import numpy
  import jax
  import jax.numpy
  from jax import random
  from flax.struct import dataclass
  from evojax.task.base import TaskState
  from evojax.task.base import VectorizedTask
label_desc: Use image node_id 1 to create a new CartPoleSwingUp object from inherited base classes: VectorizedTask with example: obj = CartPoleSwingUp(max_steps, harder, test)