def find_next_action(self, obs, agents, i, full_agent_list=None, active_agent_mask=None):
    agent_index = i
    before_index = np.array(active_agent_mask)[:agent_index]
    agent_index = (agent_index - len(before_index[(before_index == False)]))
    agents = list(compress(full_agent_list, active_agent_mask))
    observation_array = []
    if (not self.is_init):
        self.n_agents = len(agents)
        self.init(agents)
        for a in range(self.n_agents):
            pos_difference = (agents[a].goal_global_frame - agents[a].pos_global_frame)
            dist_next_waypoint = ((pos_difference / (np.linalg.norm(pos_difference, ord=1) + 1e-07)) * agents[a].pref_speed)
            vel_next_waypoint = dist_next_waypoint
            observation_array.append([agents[a].pos_global_frame[0], agents[a].pos_global_frame[1], vel_next_waypoint[0], vel_next_waypoint[1], agents[a].goal_global_frame[0], agents[a].goal_global_frame[1]])
    else:
        self.n_agents = len(agents)
        self.init(agents)
        for a in range(self.n_agents):
            if (agents[a].speed_global_frame <= (agents[a].pref_speed / 3)):
                pos_difference = (agents[a].goal_global_frame - agents[a].pos_global_frame)
                dist_next_waypoint = ((pos_difference / (np.linalg.norm(pos_difference, ord=1) + 1e-07)) * agents[a].pref_speed)
                vel_next_waypoint = dist_next_waypoint
                observation_array.append([agents[a].pos_global_frame[0], agents[a].pos_global_frame[1], vel_next_waypoint[0], vel_next_waypoint[1], agents[a].goal_global_frame[0], agents[a].goal_global_frame[1]])
            else:
                observation_array.append([agents[a].pos_global_frame[0], agents[a].pos_global_frame[1], agents[a].vel_global_frame[0], agents[a].vel_global_frame[1], agents[a].goal_global_frame[0], agents[a].goal_global_frame[1]])
    initial_state = np.array(observation_array)
    s = None
    s = socialforce.Simulator(initial_state, delta_t=0.1)
    states = np.stack([s.step().state.copy() for _ in range(1)])
    next_waypoint_x = states[:, agent_index, 0][0]
    next_waypoint_y = states[:, agent_index, 1][0]
    next_waypoint_vel_x = states[:, agent_index, 2][0]
    next_waypoint_vel_y = states[:, agent_index, 3][0]
    self.next_waypoint = np.array([next_waypoint_x, next_waypoint_y])
    goal_direction = (self.next_waypoint - agents[agent_index].pos_global_frame)
    self.dist_to_goal = math.sqrt(((goal_direction[0] ** 2) + (goal_direction[1] ** 2)))
    if (self.dist_to_goal > 1e-08):
        ref_prll = (goal_direction / agents[agent_index].dist_to_goal)
    else:
        ref_prll = goal_direction
    ref_orth = np.array([(- ref_prll[1]), ref_prll[0]])
    ref_prll_angle_global_frame = np.arctan2(ref_prll[1], ref_prll[0])
    heading_ego_frame = wrap((agents[agent_index].heading_global_frame - ref_prll_angle_global_frame))
    vel_global_frame = np.array([next_waypoint_vel_x, next_waypoint_vel_y])
    speed_global_frame = np.linalg.norm(vel_global_frame)
    if (speed_global_frame > agents[agent_index].pref_speed):
        speed_global_frame = agents[agent_index].pref_speed
    action = np.array([speed_global_frame, (- heading_ego_frame)])
    return action
def find_next_action(self, obs, agents, i, full_agent_list=None, active_agent_mask=None):
    agent_index = i
    before_index = np.array(active_agent_mask)[:agent_index]
    agent_index = (agent_index - len(before_index[(before_index == False)]))
    agents = list(compress(full_agent_list, active_agent_mask))
    observation_array = []
    if (not self.is_init):
        self.n_agents = len(agents)
        self.init(agents)
        for a in range(self.n_agents):
            pos_difference = (agents[a].goal_global_frame - agents[a].pos_global_frame)
            dist_next_waypoint = ((pos_difference / (np.linalg.norm(pos_difference, ord=1) + 1e-07)) * agents[a].pref_speed)
            vel_next_waypoint = dist_next_waypoint
            observation_array.append([agents[a].pos_global_frame[0], agents[a].pos_global_frame[1], vel_next_waypoint[0], vel_next_waypoint[1], agents[a].goal_global_frame[0], agents[a].goal_global_frame[1]])
    else:
        self.n_agents = len(agents)
        self.init(agents)
        for a in range(self.n_agents):
            if (agents[a].speed_global_frame <= (agents[a].pref_speed / 3)):
                pos_difference = (agents[a].goal_global_frame - agents[a].pos_global_frame)
                dist_next_waypoint = ((pos_difference / (np.linalg.norm(pos_difference, ord=1) + 1e-07)) * agents[a].pref_speed)
                vel_next_waypoint = dist_next_waypoint
                observation_array.append([agents[a].pos_global_frame[0], agents[a].pos_global_frame[1], vel_next_waypoint[0], vel_next_waypoint[1], agents[a].goal_global_frame[0], agents[a].goal_global_frame[1]])
            else:
                observation_array.append([agents[a].pos_global_frame[0], agents[a].pos_global_frame[1], agents[a].vel_global_frame[0], agents[a].vel_global_frame[1], agents[a].goal_global_frame[0], agents[a].goal_global_frame[1]])
    initial_state = np.array(observation_array)
    s = None
    s = socialforce.Simulator(initial_state, delta_t=0.1)
    states = np.stack([s.step().state.copy() for _ in range(1)])
    next_waypoint_x = states[:, agent_index, 0][0]
    next_waypoint_y = states[:, agent_index, 1][0]
    next_waypoint_vel_x = states[:, agent_index, 2][0]
    next_waypoint_vel_y = states[:, agent_index, 3][0]
    self.next_waypoint = np.array([next_waypoint_x, next_waypoint_y])
    goal_direction = (self.next_waypoint - agents[agent_index].pos_global_frame)
    self.dist_to_goal = math.sqrt(((goal_direction[0] ** 2) + (goal_direction[1] ** 2)))
    if (self.dist_to_goal > 1e-08):
        ref_prll = (goal_direction / agents[agent_index].dist_to_goal)
    else:
        ref_prll = goal_direction
    ref_orth = np.array([(- ref_prll[1]), ref_prll[0]])
    ref_prll_angle_global_frame = np.arctan2(ref_prll[1], ref_prll[0])
    heading_ego_frame = wrap((agents[agent_index].heading_global_frame - ref_prll_angle_global_frame))
    vel_global_frame = np.array([next_waypoint_vel_x, next_waypoint_vel_y])
    speed_global_frame = np.linalg.norm(vel_global_frame)
    if (speed_global_frame > agents[agent_index].pref_speed):
        speed_global_frame = agents[agent_index].pref_speed
    action = np.array([speed_global_frame, (- heading_ego_frame)])
    return action