class_docstr: Water world.
  ref: https://cs.stanford.edu/people/karpathy/reinforcejs/waterworld.html
code: def __init__(
      self,
      num_items: int = 50,
      max_steps: int = 1000,
      test: bool = False,
  ):
      self.max_steps = max_steps
      self.test = test
      self.obs_shape = tuple(
          [
              NUM_RANGE_SENSORS * SENSOR_DATA_DIM + 2,
          ]
      )
      self.act_shape = tuple(
          [
              4,
          ]
      )
      walls = jnp.array(
          [
              [0, 0, 0, SCREEN_H],
              [0, SCREEN_H, SCREEN_W, SCREEN_H],
              [SCREEN_W, SCREEN_H, SCREEN_W, 0],
              [SCREEN_W, 0, 0, 0],
          ]
      )
  
      def reset_fn(key):
          next_key, key = random.split(key)
          ks = random.split(key, 1 + num_items)
          agent = create_bubbles(ks[0][None, :], True)
          items = create_bubbles(ks[1:], False)
          obs = get_obs(agent, items, walls)
          return State(
              agent_state=agent,
              item_state=items,
              obs=obs,
              steps=jnp.zeros((), dtype=jnp.int32),
              key=next_key,
          )
  
      self._reset_fn = jax.jit(jax.vmap(reset_fn))
  
      def step_fn(state, action):
          next_key, key = random.split(state.key)
          direction = random.choice(
              key, 4, (), replace=False, p=action.ravel()
          )
          agent = update_agent_state(state.agent_state, direction)
          items = update_item_state(state.item_state)
          agent, items, reward = get_reward(agent, items)
          steps = state.steps + 1
          done = jnp.where(steps >= max_steps, 1, 0)
          obs = get_obs(agent, items, walls)
          return (
              State(
                  agent_state=agent,
                  item_state=items,
                  obs=obs,
                  steps=steps,
                  key=next_key,
              ),
              reward,
              done,
          )
  
      self._step_fn = jax.jit(jax.vmap(step_fn))
  
imports: from typing import Tuple
  from functools import partial
  from PIL import Image
  from PIL import ImageDraw
  import numpy
  import jax
  import jax.numpy
  from jax import random
  from jax import tree_util
  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 WaterWorld object from inherited base classes: VectorizedTask with example: obj = WaterWorld(num_items, max_steps, test)