action_saver.py 549 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
5
6
7
8
9

class ActionSaver:
	def __init__(self):
		self.saved_action = None
	
	@property
	def is_action_saved(self):
10
		return self.saved_action is not None
Dipam Chakraborty's avatar
Dipam Chakraborty committed
11

12
13
14
15
	def save_action_if_allowed(self, action, state):
		if not self.is_action_saved and \
           RailEnvActions.is_moving_action(action) and \
           not TrainState.is_malfunction_state(state):
Dipam Chakraborty's avatar
Dipam Chakraborty committed
16
17
18
19
20
			self.saved_action = action

	def clear_saved_action(self):
		self.saved_action = None