code: def get_obs(
      agent: BubbleStatus,
      agents: 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 and agents.
          items_and_agents = tree_util.tree_map(
              lambda x, y: jnp.concatenate([x, y], axis=0),
              items,
              agents,
          )
          n_items = len(items_and_agents.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_and_agents.pos_x,
              items_and_agents.pos_y,
          )
          dist_to_item = jnp.where(
              jnp.bitwise_and(
                  items_and_agents.valid, intersected_with_items
              ),
              jnp.sqrt(
                  jnp.square(item_intersections[:, 1] - agent_xy[1])
                  + jnp.square(item_intersections[:, 0] - agent_xy[0])
              ),
              MAX_RANGE,
          )
          dist_to_item = jnp.where(
              dist_to_item == 0.0, MAX_RANGE, dist_to_item
          )
          ix_items = jnp.argmin(dist_to_item)
          # Fill in the sensor data.
          detected_xy = jnp.where(
              jnp.min(dist_to_walls) < jnp.min(dist_to_item),
              jnp.where(
                  intersected_with_walls[ix_walls],
                  jnp.array(
                      [
                          dist_to_walls[ix_walls],
                          MAX_RANGE,
                          MAX_RANGE,
                          MAX_RANGE,
                          0,
                          0,
                      ]
                  ),
                  jnp.array(
                      [MAX_RANGE, 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_item[ix_items],
                              MAX_RANGE,
                              MAX_RANGE,
                              items.vel_x[ix_items],
                              items.vel_y[ix_items],
                          ]
                      ),
                      jnp.where(
                          items.bubble_type[ix_items] == TYPE_POISON,
                          jnp.array(
                              [
                                  MAX_RANGE,
                                  MAX_RANGE,
                                  dist_to_item[ix_items],
                                  MAX_RANGE,
                                  items.vel_x[ix_items],
                                  items.vel_y[ix_items],
                              ]
                          ),
                          jnp.array(
                              [
                                  MAX_RANGE,
                                  MAX_RANGE,
                                  MAX_RANGE,
                                  dist_to_item[ix_items],
                                  items.vel_x[ix_items],
                                  items.vel_y[ix_items],
                              ]
                          ),
                      ),
                  ),
                  jnp.array(
                      [MAX_RANGE, MAX_RANGE, MAX_RANGE, MAX_RANGE, 0, 0]
                  ),
              ),
          )
          sensor_obs.append(detected_xy)
  
      sensor_obs = jnp.stack(sensor_obs)
      sensor_obs = sensor_obs.at[:, :4].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 8 for calling a global function with example usage: get_obs(agent, agents, items, walls) and returns: jnp