action_preprocessing.py 1.92 KB
Newer Older
1
from flatland.core.grid.grid_utils import position_to_coordinate
Dipam Chakraborty's avatar
Dipam Chakraborty committed
2
from flatland.envs.agent_utils import TrainState
3
4
5
from flatland.envs.rail_env_action import RailEnvActions
from flatland.envs.step_utils.transition_utils import check_valid_action

Dipam Chakraborty's avatar
Dipam Chakraborty committed
6
7

def process_illegal_action(action: RailEnvActions):
8
	if not RailEnvActions.is_action_valid(action): 
Dipam Chakraborty's avatar
Dipam Chakraborty committed
9
10
		return RailEnvActions.DO_NOTHING
	else:
Dipam Chakraborty's avatar
Dipam Chakraborty committed
11
		return RailEnvActions(action)
Dipam Chakraborty's avatar
Dipam Chakraborty committed
12
13


14
def process_do_nothing(state: TrainState, saved_action: RailEnvActions):
15
16
    if state == TrainState.MOVING:
        action = RailEnvActions.MOVE_FORWARD
17
18
    elif saved_action:
        action = saved_action
19
20
21
    else:
        action = RailEnvActions.STOP_MOVING
    return action
Dipam Chakraborty's avatar
Dipam Chakraborty committed
22
23


24
25
def process_left_right(action, rail, position, direction):
    if not check_valid_action(action, rail, position, direction):
26
27
28
29
30
31
32
33
34
35
36
37
38
        action = RailEnvActions.MOVE_FORWARD
    return action


def preprocess_action_when_waiting(action, state):
    """
    Set action to DO_NOTHING if in waiting state
    """
    if state == TrainState.WAITING:
        action = RailEnvActions.DO_NOTHING
    return action


39
def preprocess_raw_action(action, state, saved_action):
40
41
42
43
44
45
46
47
    """
    Preprocesses actions to handle different situations of usage of action based on context
        - DO_NOTHING is converted to FORWARD if train is moving
        - DO_NOTHING is converted to STOP_MOVING if train is moving
    """
    action = process_illegal_action(action)

    if action == RailEnvActions.DO_NOTHING:
48
        action = process_do_nothing(state, saved_action)
49
50
51

    return action

52
def preprocess_moving_action(action, rail, position, direction):
53
54
55
56
57
58
59
60
    """
    LEFT/RIGHT is converted to FORWARD if left/right is not available and train is moving
    FORWARD is converted to STOP_MOVING if leading to dead end?
    """
    if action in [RailEnvActions.MOVE_LEFT, RailEnvActions.MOVE_RIGHT]:
        action = process_left_right(action, rail, position, direction)

    return action