class_docstr: Water world, multi-agents training version.
class_docstr_tok: ["Water","world",",","multi-agents","training","version","."]
code: def __init__(
      self,
      num_agents: int = 16,
      num_items: int = 100,
      max_steps: int = 1000,
      test: bool = False,
  ):
      self.multi_agent_training = True
      self.max_steps = max_steps
      self.test = test
      self.obs_shape = tuple(
          [
              num_agents,
              NUM_RANGE_SENSORS * SENSOR_DATA_DIM + 2,
          ]
      )
      self.act_shape = tuple([num_agents, 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, num_agents + num_items)
          agents = create_bubbles(ks[:num_agents], True)
          items = create_bubbles(ks[num_agents:], False)
          obs = get_obs(agents, agents, items, walls)
          return State(
              agent_state=agents,
              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)
          action_keys = random.split(key, num_agents)
          directions = select_direction(action_keys, action)
          agents = update_agent_state(state.agent_state, directions)
          items = update_item_state(state.item_state)
          agents, items, rewards = get_rewards(agents, items)
          # items_state.shape=(num_agents, num_items), merge to (num_items, ).
          items = BubbleStatus(
              pos_x=items.pos_x[0],
              pos_y=items.pos_y[0],
              vel_x=items.vel_x[0],
              vel_y=items.vel_y[0],
              bubble_type=items.bubble_type[0],
              poison_cnt=items.poison_cnt[0],
              valid=jnp.prod(items.valid, axis=0),
          )
          steps = state.steps + 1
          done = jnp.where(steps >= max_steps, 1, 0)
          obs = get_obs(agents, agents, items, walls)
          return (
              State(
                  agent_state=agents,
                  item_state=items,
                  obs=obs,
                  steps=steps,
                  key=next_key,
              ),
              rewards,
              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 MultiAgentWaterWorld object from inherited base classes: VectorizedTask with example: obj = MultiAgentWaterWorld(num_agents, num_items, max_steps, test)