action_preprocessing.py 2.03 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
8

def process_illegal_action(action: RailEnvActions):
	# TODO - Dipam : This check is kind of weird, change this
9
	if action is None or action not in RailEnvActions._value2member_map_: 
Dipam Chakraborty's avatar
Dipam Chakraborty committed
10
11
12
13
14
		return RailEnvActions.DO_NOTHING
	else:
		return action


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


23
24
def process_left_right(action, rail, position, direction):
    if not check_valid_action(action, rail, position, direction):
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
        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


def preprocess_raw_action(action, state):
    """
    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:
        action = process_do_nothing(state)

    return action

51
def preprocess_moving_action(action, rail, position, direction):
52
53
54
55
56
57
58
59
60
61
    """
    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)

    if not check_valid_action(action, rail, position, direction): # TODO: Dipam - Not sure if this is needed
        action = RailEnvActions.STOP_MOVING
    return action