code: def update_agent_state(
      agent: BubbleStatus, direction: jnp.int32
  ) -> BubbleStatus:
      vel_x = agent.vel_x
      vel_x = jnp.where(direction == ACT_RIGHT, vel_x + 1, vel_x)
      vel_x = jnp.where(direction == ACT_LEFT, vel_x - 1, vel_x)
      vel_x = vel_x * 0.95
  
      vel_y = agent.vel_y
      vel_y = jnp.where(direction == ACT_UP, vel_y - 1, vel_y)
      vel_y = jnp.where(direction == ACT_DOWN, vel_y + 1, vel_y)
      vel_y = vel_y * 0.95
  
      pos_x = agent.pos_x + vel_x
      pos_y = agent.pos_y + vel_y
      # Collide with the west wall.
      vel_x = jnp.where(pos_x < 1, 0, vel_x)
      vel_y = jnp.where(pos_x < 1, 0, vel_y)
      pos_x = jnp.where(pos_x < 1, 1, pos_x)
      # Collide with the east wall.
      vel_x = jnp.where(pos_x > SCREEN_W - 1, 0, vel_x)
      vel_y = jnp.where(pos_x > SCREEN_W - 1, 0, vel_y)
      pos_x = jnp.where(pos_x > SCREEN_W - 1, SCREEN_W - 1, pos_x)
      # Collide with the north wall.
      vel_x = jnp.where(pos_y < 1, 0, vel_x)
      vel_y = jnp.where(pos_y < 1, 0, vel_y)
      pos_y = jnp.where(pos_y < 1, 1, pos_y)
      # Collide with the south wall.
      vel_x = jnp.where(pos_y > SCREEN_H - 1, 0, vel_x)
      vel_y = jnp.where(pos_y > SCREEN_H - 1, 0, vel_y)
      pos_y = jnp.where(pos_y > SCREEN_H - 1, SCREEN_H - 1, pos_y)
  
      return BubbleStatus(
          pos_x=pos_x,
          pos_y=pos_y,
          vel_x=vel_x,
          vel_y=vel_y,
          bubble_type=agent.bubble_type,
          valid=agent.valid,
          poison_cnt=agent.poison_cnt,
      )
  
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 4 for calling a global function with example usage: update_agent_state(agent, direction) and returns: BubbleStatus