code: def unpack_obs(obs: jnp.ndarray) -> jnp.ndarray:
      position, theta = obs[..., :2], obs[..., 2:]
      return position, theta
  
imports: from PIL import Image, ImageDraw
  from functools import partial
  from typing import Tuple
  import jax
  from jax import vmap
  import jax.numpy
  from flax.struct import dataclass
  from evojax.task.base import VectorizedTask
  from evojax.task.base import TaskState
label_desc: Use image node_id 3 for calling a global function with example usage: unpack_obs(obs) and returns: position,  theta