action_saver.py 924 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
    @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):
17
18
19
20
        if action.is_moving_action() and \
               not self.is_action_saved and \
               not state.is_malfunction_state() and \
               not state == TrainState.DONE:
Dipam Chakraborty's avatar
Dipam Chakraborty committed
21
22
23
24
            self.saved_action = action

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

26
27
28
29
30
31
    def to_dict(self):
        return {"saved_action": self.saved_action}
    
    def from_dict(self, load_dict):
        self.saved_action = load_dict['saved_action']

Dipam Chakraborty's avatar
Dipam Chakraborty committed
32