class_docstr: Water world, multi-agents training version.
class_docstr_tok: ["Water","world",",","multi-agents","training","version","."]
code: def render(state: State, task_id: int = 0) -> Image:
      img = Image.new("RGB", (SCREEN_W, SCREEN_H), (255, 255, 255))
      draw = ImageDraw.Draw(img)
      state = tree_util.tree_map(lambda s: s[task_id], state)
      # Draw the agent.
      agents = state.agent_state
      sensor_data = np.array(
          state.obs[:, :-2].reshape(
              -1, NUM_RANGE_SENSORS, SENSOR_DATA_DIM
          )
      )
      for i, (x, y) in enumerate(zip(agents.pos_x, agents.pos_y)):
          for j, obs in enumerate(sensor_data[i]):
              ang = j * DELTA_ANG
              dist = np.min(obs[:4])
              x_end = x + dist * MAX_RANGE * np.cos(ang)
              y_end = y + dist * MAX_RANGE * np.sin(ang)
              draw.line((x, y, x_end, y_end), fill=(0, 0, 0), width=1)
          draw.ellipse(
              (
                  x - BUBBLE_RADIUS,
                  y - BUBBLE_RADIUS,
                  x + BUBBLE_RADIUS,
                  y + BUBBLE_RADIUS,
              ),
              fill=(255, 255, 0),
              outline=(0, 0, 0),
          )
      # Draw the items.
      items = state.item_state
      for v, t, x, y in zip(
          np.array(items.valid, dtype=bool),
          np.array(items.bubble_type, dtype=int),
          np.array(items.pos_x),
          np.array(items.pos_y),
      ):
          if v:
              color = (0, 255, 0) if t == TYPE_FOOD else (255, 0, 0)
              draw.ellipse(
                  (
                      x - BUBBLE_RADIUS,
                      y - BUBBLE_RADIUS,
                      x + BUBBLE_RADIUS,
                      y + BUBBLE_RADIUS,
                  ),
                  fill=color,
                  outline=(0, 0, 0),
              )
      return img
  
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 4 for calling the MultiAgentWaterWorld obj's underlying member method code with example usage: obj.render(state, task_id) and returns: img