I'm trying to figure out how action masking works and the closest workaround i get is following the hanabi environment example and then write my own custom version to learn more:
it is my very basic choosing numbers game working with action_mask
""" based on RL environment for Hanabi, using an API similar to OpenAI Gym.
when obs return [0,0] available actions are [0,1,2]
when obs return [1,0] available actions are [0,2,3]
when obs return [0,1] available actions are [0,1,4]
when obs return [1,1] available actions are [0,3,4]
"""
import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
import numpy as np
import gin.tf
from tf_agents.environments import py_environment
from tf_agents.specs import array_spec
from tf_agents.trajectories import time_step as ts
class CustomMaskedEnvironment(py_environment.PyEnvironment):
def __init__(self, gamma=0.95):
super().__init__()
self.gamma = gamma
self.one = 0
self.two = 0
self.invalid_action = 0
self._action_spec = array_spec.BoundedArraySpec(shape=(), dtype=np.int_, minimum=0, maximum=1)
self._observation_spec = {
'observation': array_spec.ArraySpec(shape=(2,), dtype=np.float32),
'legal_moves': array_spec.ArraySpec(shape=(), dtype=np.bool_),
}
def observation_spec(self):
return self._observation_spec
def action_spec(self):
return self._action_spec
def _reset(self):
self.state = self._make_observation()
obs = self.state
legal_moves = self.num_moves()
observations_and_legal_moves = {
'observation': obs,
'legal_moves':np.logical_not(legal_moves),
}
return ts.restart(observations_and_legal_moves)
def _step(self, action):
if self._current_time_step.is_last():
return self.reset()
# Apply the action to the state.
self.apply_move(action)
obs = self._make_observation()
legal_moves = self.num_moves()
done = self.state.is_terminal()
reward = self.invalid_action
self.invalid_action = 0
observations_and_legal_moves = {
'observation': obs,
'legal_moves': np.logical_not(legal_moves),
}
if done:
return ts.termination(observations_and_legal_moves, reward)
else:
return ts.transition(observations_and_legal_moves, reward, self.gamma)
def apply_move(self, action):
if action == 0:
pass
elif action == 1:
if self.one == 1:
self.invalid_action = -1
else:
self.one = 1
self.invalid_action = 1
elif action == 2:
if self.two == 1:
self.invalid_action = -1
else:
self.two = 1
self.invalid_action = 1
elif action == 3:
if self.one == 1:
self.one = 0
self.invalid_action = 1
else:
self.invalid_action = -1
elif action == 4:
if self.two == 1:
self.two = 0
self.invalid_action = 1
else:
self.invalid_action = -1
def num_moves(self):
action_mask = np.ones(5, dtype=np.uint8)
# IF one and two are present in the observation
# make available three and four
if self.one == 1 and self.two == 1:
action_mask[0] = 0 # wait
action_mask[3] = 0 # make three available
action_mask[4] = 0 # make four available
# IF one is present in the observation make three available
elif self.one == 1:
action_mask[0] = 0 # wait
action_mask[2] = 0 # two is available
action_mask[3] = 0 # three is available
# IF two is present in the observation make four available
elif self.one == 1:
action_mask[0] = 0 # wait
action_mask[1] = 0 # one is available
action_mask[4] = 0 # make four available
# IF isnt one or two present in the observation
else:
action_mask[0] = 0 # wait
action_mask[1] = 0 # one is available
action_mask[2] = 0 # two is available
return action_mask
def _make_observation(self):
one = self.one
two = self.two
obs = np.array([one, two], dtype=np.float32)
return obs
from tf_agents.agents.reinforce import reinforce_agent
from tf_agents.agents.dqn import dqn_agent
from tf_agents.agents.categorical_dqn import categorical_dqn_agent
from tf_agents.networks import categorical_q_network
from tf_agents.networks import q_network
from tf_agents.replay_buffers import tf_uniform_replay_buffer
from tf_agents.utils import common
import tensorflow as tf
from tf_agents.environments import tf_py_environment
import reverb
from tf_agents.networks.mask_splitter_network import MaskSplitterNetwork
from tf_agents.networks.actor_distribution_network import ActorDistributionNetwork
from tf_agents.specs import tensor_spec
from tf_agents.replay_buffers import reverb_replay_buffer
from tf_agents.replay_buffers import reverb_utils
optimizer = tf.keras.optimizers.Adam(learning_rate=0.00025)
train_step_counter = tf.Variable(0)
fc_layer_params = (256, 256)
def filter_fun(observation):
return observation['observation'], observation['legal_moves']
train_py_env = CustomMaskedEnvironment()
train_env = tf_py_environment.TFPyEnvironment(train_py_env)
eval_py_env = CustomMaskedEnvironment()
eval_env = tf_py_environment.TFPyEnvironment(eval_py_env)
num_eval_episodes = 10
masked_actor_network = MaskSplitterNetwork(
splitter_fn=filter_fun,
wrapped_network= ActorDistributionNetwork(
train_env.observation_spec(['observation'],
train_env.action_spec(),
fc_layer_params=fc_layer_params
),
passthrough_mask=True
)
agent = reinforce_agent.ReinforceAgent(
train_env.time_step_spec(),
train_env.action_spec(),
actor_network=masked_actor_network,
optimizer=optimizer,
normalize_returns=True,
train_step_counter=train_step_counter,
)
agent.initialize()
eval_policy = agent.policy
collect_policy = agent.collect_policy
def compute_avg_return(environment, policy, num_episodes=10):
total_return = 0.0
for _ in range(num_episodes):
time_step = environment.reset()
episode_return = 0.0
while not time_step.is_last():
action_step = policy.action(time_step)
time_step = environment.step(action_step.action)
episode_return += time_step.reward
total_return += episode_return
avg_return = total_return / num_episodes
return avg_return.numpy()[0]
replay_buffer_capacity = 90000
table_name = 'uniform_table'
replay_buffer_signature = tensor_spec.from_spec(
agent.collect_data_spec)
replay_buffer_signature = tensor_spec.add_outer_dim(
replay_buffer_signature)
table = reverb.Table(
table_name,
max_size=replay_buffer_capacity,
sampler=reverb.selectors.Uniform(),
remover=reverb.selectors.Fifo(),
rate_limiter=reverb.rate_limiters.MinSize(1),
signature=replay_buffer_signature)
reverb_server = reverb.Server([table])
replay_buffer = reverb_replay_buffer.ReverbReplayBuffer(
agent.collect_data_spec,
table_name=table_name,
sequence_length=None,
local_server=reverb_server)
rb_observer = reverb_utils.ReverbAddEpisodeObserver(
replay_buffer.py_client,
table_name,
replay_buffer_capacity
)
def collect_episode(environment, policy, num_episodes):
driver = py_driver.PyDriver(
environment,
py_tf_eager_policy.PyTFEagerPolicy(
policy, use_tf_function=True),
[rb_observer],
max_episodes=num_episodes)
initial_time_step = environment.reset()
driver.run(initial_time_step)
# (Optional) Optimize by wrapping some of the code in a graph using TF function.
agent.train = common.function(agent.train)
# Reset the train step
agent.train_step_counter.assign(0)
# Evaluate the agent's policy once before training.
avg_return = compute_avg_return(eval_env, agent.policy, num_eval_episodes)
returns = [avg_return]
for _ in range(num_iterations):
# Collect a few episodes using collect_policy and save to the replay buffer.
collect_episode(
train_py_env, agent.collect_policy, collect_episodes_per_iteration)
# Use data from the buffer and update the agent's network.
iterator = iter(replay_buffer.as_dataset(sample_batch_size=1))
trajectories, _ = next(iterator)
train_loss = agent.train(experience=trajectories)
replay_buffer.clear()
step = tf_agent.train_step_counter.numpy()
if step % log_interval == 0:
print('step = {0}: loss = {1}'.format(step, train_loss.loss))
if step % eval_interval == 0:
avg_return = compute_avg_return(eval_env, agent.policy, num_eval_episodes)
print('step = {0}: Average Return = {1}'.format(step, avg_return))
returns.append(avg_return)
everything seems to works fine but i get the following error messages:
logits = tf.compat.v2.where( tensorflow.python.framework.errors_impl.InvalidArgumentError: Exception encountered when calling layer "CategoricalProjectionNetwork" (type CategoricalProjectionNetwork).
condition [1,5], then [1,2], and else [] must be broadcastable [Op:SelectV2]
Call arguments received by layer "CategoricalProjectionNetwork" (type CategoricalProjectionNetwork): • inputs=tf.Tensor(shape=(1, 256), dtype=float32) • outer_rank=1 • training=False • mask=tf.Tensor(shape=(1, 5), dtype=bool) [reverb/cc/platform/default/server.cc:84] Shutting down replay server
but acctually without a graphic card maybe reverb is not in my "usable options" just a workaround, i would migrate my code to pytorch, rllib, whatever that gives me a usable solution
