class_docstr: Water world.
  ref: https://cs.stanford.edu/people/karpathy/reinforcejs/waterworld.html
code: def step(
      self, state: State, action: jnp.ndarray
  ) -> Tuple[State, jnp.ndarray, jnp.ndarray]:
      return self._step_fn(state, action)
  
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 3 for calling the WaterWorld obj's underlying member method code with example usage: obj.step(state, action) and returns: self