class_docstr: Water world.
  ref: https://cs.stanford.edu/people/karpathy/reinforcejs/waterworld.html
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.
      agent = state.agent_state
      x, y = agent.pos_x, agent.pos_y
      sensor_data = np.array(state.obs)[:-2].reshape(
          NUM_RANGE_SENSORS, SENSOR_DATA_DIM
      )
      for i, obs in enumerate(sensor_data):
          ang = i * DELTA_ANG
          dist = np.min(obs[:3])
          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 WaterWorld obj's underlying member method code with example usage: obj.render(state, task_id) and returns: img