action_saver.py 704 Bytes
Newer Older
1
from flatland.envs.rail_env_action import RailEnvActions
2
from flatland.envs.step_utils.states import TrainState
Dipam Chakraborty's avatar
Dipam Chakraborty committed
3
4

class ActionSaver:
Dipam Chakraborty's avatar
Dipam Chakraborty committed
5
6
    def __init__(self):
        self.saved_action = None
Dipam Chakraborty's avatar
Dipam Chakraborty committed
7

Dipam Chakraborty's avatar
Dipam Chakraborty committed
8
9
10
11
12
13
14
15
16
17
    @property
    def is_action_saved(self):
        return self.saved_action is not None
    
    def __repr__(self):
        return f"is_action_saved: {self.is_action_saved}, saved_action: {self.saved_action}"


    def save_action_if_allowed(self, action, state):
        if not self.is_action_saved and \
Dipam Chakraborty's avatar
Dipam Chakraborty committed
18
19
               action.is_moving_action() and \
               not state.is_malfunction_state():
Dipam Chakraborty's avatar
Dipam Chakraborty committed
20
21
22
23
            self.saved_action = action

    def clear_saved_action(self):
        self.saved_action = None
Dipam Chakraborty's avatar
Dipam Chakraborty committed
24
25