code: def get_obs(
      agent: BubbleStatus, items: BubbleStatus, walls: jnp.ndarray
  ) -> jnp.ndarray:
      sensor_obs = []
      agent_xy = jnp.array([agent.pos_x, agent.pos_y]).ravel()
      for i in jnp.arange(NUM_RANGE_SENSORS):
          ang = i * DELTA_ANG
          range_xy = jnp.array(
              [
                  agent.pos_x + MAX_RANGE * jnp.cos(ang),
                  agent.pos_y + MAX_RANGE * jnp.sin(ang),
              ]
          )
          # Check for intersections with the 4 walls.
          (
              intersected_with_walls,
              wall_intersections,
          ) = get_line_seg_intersection(
              jnp.ones(4) * agent_xy[0],
              jnp.ones(4) * agent_xy[1],
              jnp.ones(4) * range_xy[0],
              jnp.ones(4) * range_xy[1],
              walls[:, 0],
              walls[:, 1],
              walls[:, 2],
              walls[:, 3],
          )
          dist_to_walls = jnp.where(
              intersected_with_walls,
              jnp.sqrt(
                  jnp.square(wall_intersections[:, 1] - agent_xy[1])
                  + jnp.square(wall_intersections[:, 0] - agent_xy[0])
              ),
              MAX_RANGE,
          )
          ix_walls = jnp.argmin(dist_to_walls)
          # Check for intersections with the items.
          n_items = len(items.valid)
          (
              intersected_with_items,
              item_intersections,
          ) = get_line_dot_intersection(
              jnp.ones(n_items) * agent_xy[0],
              jnp.ones(n_items) * agent_xy[1],
              jnp.ones(n_items) * range_xy[0],
              jnp.ones(n_items) * range_xy[1],
              items.pos_x,
              items.pos_y,
          )
          dist_to_items = jnp.where(
              jnp.bitwise_and(items.valid, intersected_with_items),
              jnp.sqrt(
                  jnp.square(item_intersections[:, 1] - agent_xy[1])
                  + jnp.square(item_intersections[:, 0] - agent_xy[0])
              ),
              MAX_RANGE,
          )
          ix_items = jnp.argmin(dist_to_items)
          # Fill in the sensor data.
          detected_xy = jnp.where(
              jnp.min(dist_to_walls) < jnp.min(dist_to_items),
              jnp.where(
                  intersected_with_walls[ix_walls],
                  jnp.array(
                      [
                          dist_to_walls[ix_walls],
                          MAX_RANGE,
                          MAX_RANGE,
                          0,
                          0,
                      ]
                  ),
                  jnp.array([MAX_RANGE, MAX_RANGE, MAX_RANGE, 0, 0]),
              ),
              jnp.where(
                  intersected_with_items[ix_items],
                  jnp.where(
                      items.bubble_type[ix_items] == TYPE_FOOD,
                      jnp.array(
                          [
                              MAX_RANGE,
                              dist_to_items[ix_items],
                              MAX_RANGE,
                              items.vel_x[ix_items],
                              items.vel_y[ix_items],
                          ]
                      ),
                      jnp.array(
                          [
                              MAX_RANGE,
                              MAX_RANGE,
                              dist_to_items[ix_items],
                              items.vel_x[ix_items],
                              items.vel_y[ix_items],
                          ]
                      ),
                  ),
                  jnp.array([MAX_RANGE, MAX_RANGE, MAX_RANGE, 0, 0]),
              ),
          )
          sensor_obs.append(detected_xy)
  
      sensor_obs = jnp.stack(sensor_obs)
      sensor_obs = sensor_obs.at[:, :3].divide(
          MAX_RANGE
      )  # Normalized distances.
      vel_xy = jnp.array([agent.vel_x, agent.vel_y]).ravel()
      return jnp.concatenate([sensor_obs.ravel(), vel_xy], axis=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 7 for calling a global function with example usage: get_obs(agent, items, walls) and returns: jnp