code: def create_bubbles(key: jnp.ndarray, is_agent: bool) -> BubbleStatus:
      k_pos_r, k_pos_theta, k_vel, k_bubble_type = random.split(key, 4)
      if is_agent:
          bubble_type = TYPE_AGENT
          vel_x = vel_y = 0.0
          r = 270
      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
          )
          r = random.uniform(k_pos_r, shape=(), minval=0, maxval=200)
      theta = random.uniform(
          k_pos_theta, shape=(), minval=0, maxval=2 * jnp.pi
      )
      pos_x = r * jnp.cos(theta) + SCREEN_W // 2
      pos_y = r * jnp.sin(theta) + SCREEN_H // 2
      return BubbleStatus(
          pos_x=pos_x,
          pos_y=pos_y,
          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