Dictionary observation space Acme DQN agent

Question

I'm trying to add illegal action masking to my dqn agent using masked_epsilon_greedy. Does anyone know how can I update the policy network to use observation["your_key_for_observation"] rather than 'observation' since the observation space is a dictionary containing both the observations and legal actions?

Answer 1

the answer is adding lambda inputs: inputs["your_key_for_observation"] to the network in case someone encounters this issue in the future.