code: def get_reward(
      agent: BubbleStatus, items: BubbleStatus
  ) -> Tuple[BubbleStatus, BubbleStatus, jnp.float32]:
      dist = jnp.sqrt(
          jnp.square(agent.pos_x - items.pos_x)
          + jnp.square(agent.pos_y - items.pos_y)
      )
      rewards = (
          jnp.where(items.bubble_type == TYPE_FOOD, 1.0, -1.0)
          * items.valid
          * jnp.where(dist < MIN_DIST, 1, 0)
      )
      poison_cnt = (
          jnp.sum(jnp.where(rewards == -1.0, 1, 0)) + agent.poison_cnt
      )
      reward = jnp.sum(rewards)
      items_valid = (dist >= MIN_DIST) * items.valid
      agent_state = BubbleStatus(
          pos_x=agent.pos_x,
          pos_y=agent.pos_y,
          vel_x=agent.vel_x,
          vel_y=agent.vel_y,
          bubble_type=agent.bubble_type,
          valid=agent.valid,
          poison_cnt=poison_cnt,
      )
      items_state = BubbleStatus(
          pos_x=items.pos_x,
          pos_y=items.pos_y,
          vel_x=items.vel_x,
          vel_y=items.vel_y,
          bubble_type=items.bubble_type,
          valid=items_valid,
          poison_cnt=items.poison_cnt,
      )
      return agent_state, items_state, reward
  
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 2 for calling a global function with example usage: get_reward(agent, items) and returns: agent_state,  items_state,  reward