code: def __init__(
      self,
      max_steps: int = 150,
      reward_type: int = 0,  # (0: as it is, 1: increase rewards for late step)
      action_type: int = 0,  # (0: theta, 1: theta/speed)
  ):
      self.max_steps = max_steps
      self.obs_shape = tuple([NEIGHBOR_NUM * 3, BOIDS_NUM])
      self.act_shape = tuple(
          [
              action_type + 1,
          ]
      )
  
      def reset_fn(key):
          next_key, key = jax.random.split(key)
          position = sample_position(key, BOIDS_NUM)
          theta = sample_theta(key, BOIDS_NUM)
          state = pack_obs(position, theta)
          return State(
              obs=choose_neighbor(state),
              state=state,
              steps=jnp.zeros((), dtype=jnp.int32),
              key=next_key,
          )
  
      self._reset_fn = jax.jit(jax.vmap(reset_fn))
  
      def step_fn(state, action):
          new_state = update_state(state, action, action_type)
          new_obs = choose_neighbor(new_state)
          new_steps = jnp.int32(state.steps + 1)
          next_key, _ = jax.random.split(state.key)
          reward = get_reward(state, max_steps, reward_type)
          done = jnp.where(max_steps <= new_steps, True, False)
          return (
              State(
                  obs=new_obs,
                  state=new_state,
                  steps=new_steps,
                  key=next_key,
              ),
              reward,
              done,
          )
  
      self._step_fn = jax.jit(jax.vmap(step_fn))
  
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 1 to create a new FlockingTask object from inherited base classes: VectorizedTask with example: obj = FlockingTask(max_steps, reward_type, action_type)