Invalid agent directions while visiting cells with ShortestPathPredictorForRailEnv

Sometimes, the agent's direction of the last visited cell generated by the ShortestPathPredictorForRailEnv marks an invalid direction - the cell does not permit that transition.

Possible source of the problem that the initial agent direction is appended to the list of visited cells when reaching the agent.target position (link: https://gitlab.aicrowd.com/flatland/flatland/blob/master/flatland/envs/predictions.py#L163).

for index in range(1, self.max_depth + 1):
    # if we're at the target, stop moving until max_depth is reached
    if new_position == agent.target or not shortest_path:
        prediction[index] = [index, *new_position, new_direction, RailEnvActions.STOP_MOVING]
        visited.add((*new_position, agent.direction))
        continue

Suggestion: Change agent.direction to the new_direction variable, so the list of the visited cells would contain the directions of the agent's path.

Version: flatland-rl==2.2.1