code: def create_bubbles(key: jnp.ndarray, is_agent: bool) -> BubbleStatus:
      k_pos_x, k_pos_y, k_vel, k_bubble_type = random.split(key, 4)
      if is_agent:
          bubble_type = TYPE_AGENT
          vel_x = vel_y = 0.0
      else:
          bubble_type = jnp.where(
              random.uniform(k_bubble_type) > 0.5,
              TYPE_FOOD,
              TYPE_POISON,
          )
          vel_x, vel_y = random.uniform(
              k_vel, shape=(2,), minval=-2.5, maxval=2.5
          )
      return BubbleStatus(
          pos_x=random.uniform(
              k_pos_x,
              shape=(),
              dtype=jnp.float32,
              minval=MIN_DIST,
              maxval=SCREEN_W - MIN_DIST,
          ),
          pos_y=random.uniform(
              k_pos_y,
              shape=(),
              dtype=jnp.float32,
              minval=MIN_DIST,
              maxval=SCREEN_H - MIN_DIST,
          ),
          vel_x=vel_x,
          vel_y=vel_y,
          bubble_type=bubble_type,
          valid=1,
          poison_cnt=0,
      )
  
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 for calling a global function with example usage: create_bubbles(key, is_agent) and returns: BubbleStatus