action_saver.py 1.13 KB
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
    @property
    def is_action_saved(self):
        return self.saved_action is not None
    
    def __repr__(self):
13
        return f"is_action_saved: {self.is_action_saved}, saved_action: {str(self.saved_action)}"
Dipam Chakraborty's avatar
Dipam Chakraborty committed
14
15
16


    def save_action_if_allowed(self, action, state):
17
18
19
        """
        Save the action if all conditions are met
            1. It is a movement based action -> Forward, Left, Right
20
21
            2. Action is not already saved 
            3. Agent is not already done
22
        """
23
        if action.is_moving_action() and not self.is_action_saved and not state == TrainState.DONE:
Dipam Chakraborty's avatar
Dipam Chakraborty committed
24
25
26
27
            self.saved_action = action

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

29
30
31
32
33
    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
34
35
36
    
    def __eq__(self, other):
        return self.saved_action == other.saved_action
37

Dipam Chakraborty's avatar
Dipam Chakraborty committed
38