Custom environment using TFagents

Question

I am trying to learn a custom environment using the TFAgents package. I am following the Hands-on-ML book (Code in colab see cell 129). My aim is to use DQN agent on a custom-written grid world environment.

Grid-World environment:

class MyEnvironment(tf_agents.environments.py_environment.PyEnvironment):

def __init__(self, discount=1.0):
    super().__init__()
    self.discount = discount

    self._action_spec = tf_agents.specs.BoundedArraySpec(shape=(), dtype=np.int32, name="action", minimum=0, maximum=3)
    self._observation_spec = tf_agents.specs.BoundedArraySpec(shape=(4, 4), dtype=np.int32, name="observation", minimum=0, maximum=1)


def action_spec(self):
    return self._action_spec

def observation_spec(self):
    return self._observation_spec

def _reset(self):
    self._state = np.zeros(2, dtype=np.int32)
    obs = np.zeros((4, 4), dtype=np.int32)
    obs[self._state[0], self._state[1]] = 1
    return tf_agents.trajectories.time_step.restart(obs)

def _step(self, action):
    self._state += [(-1, 0), (+1, 0), (0, -1), (0, +1)][action]
    reward = 0
    obs = np.zeros((4, 4), dtype=np.int32)
    done = (self._state.min() < 0 or self._state.max() > 3)
    if not done:
        obs[self._state[0], self._state[1]] = 1
    if done or np.all(self._state == np.array([3, 3])):
        reward = -1 if done else +10
        return tf_agents.trajectories.time_step.termination(obs, reward)
    else:
        return tf_agents.trajectories.time_step.transition(obs, reward, self.discount)

And the Q network is:

tf_env = MyEnvironment()

preprocessing_layer = keras.layers.Lambda(lambda obs: tf.cast(obs, np.float32) / 255.)
conv_layer_params=[(32, (2, 2), 1)]
fc_layer_params=[512]
q_net = QNetwork(
    tf_env.observation_spec(),
    tf_env.action_spec(),
    preprocessing_layers=preprocessing_layer,
    conv_layer_params=conv_layer_params,
    fc_layer_params=fc_layer_params)

And finally, the DQN agent is

train_step = tf.Variable(0)
update_period = 4 # train the model every 4 steps
optimizer = keras.optimizers.RMSprop(lr=2.5e-4, rho=0.95, momentum=0.0, epsilon=0.00001, centered=True)
epsilon_fn = keras.optimizers.schedules.PolynomialDecay(initial_learning_rate=1.0, decay_steps=250000 // update_period, end_learning_rate=0.01)


agent = DqnAgent(tf_env.time_step_spec(),
    tf_env.action_spec(),
    q_network=q_net,
    optimizer=optimizer,
    target_update_period=2000, # <=> 32,000 ALE frames
    td_errors_loss_fn=keras.losses.Huber(reduction="none"),
    gamma=0.99, # discount factor
    train_step_counter=train_step,
    epsilon_greedy=lambda: epsilon_fn(train_step))  
agent.initialize()

Directly, running the code gave me the following error trace:

/usr/local/lib/python3.6/dist-packages/gin/config.py in gin_wrapper(*args, **kwargs) 1067 scope_info = " in scope '{}'".format(scope_str) if scope_str else '' 1068 err_str = err_str.format(name, fn_or_cls, scope_info) -> 1069 utils.augment_exception_message_and_reraise(e, err_str) 1070 1071 return gin_wrapper

/usr/local/lib/python3.6/dist-packages/gin/utils.py in augment_exception_message_and_reraise(exception, message)
     39   proxy = ExceptionProxy()
     40   ExceptionProxy.__qualname__ = type(exception).__qualname__
---> 41   raise proxy.with_traceback(exception.__traceback__) from None
     42 
     43 

/usr/local/lib/python3.6/dist-packages/gin/config.py in gin_wrapper(*args, **kwargs)
   1044 
   1045     try:
-> 1046       return fn(*new_args, **new_kwargs)
   1047     except Exception as e:  # pylint: disable=broad-except
   1048       err_str = ''

/usr/local/lib/python3.6/dist-packages/tf_agents/agents/dqn/dqn_agent.py in __init__(self, time_step_spec, action_spec, q_network, optimizer, observation_and_action_constraint_splitter, epsilon_greedy, n_step_update, boltzmann_temperature, emit_log_probability, target_q_network, target_update_tau, target_update_period, td_errors_loss_fn, gamma, reward_scale_factor, gradient_clipping, debug_summaries, summarize_grads_and_vars, train_step_counter, name)
    216     tf.Module.__init__(self, name=name)
    217 
--> 218     self._check_action_spec(action_spec)
    219 
    220     if epsilon_greedy is not None and boltzmann_temperature is not None:

/usr/local/lib/python3.6/dist-packages/tf_agents/agents/dqn/dqn_agent.py in _check_action_spec(self, action_spec)
    293 
    294     # TODO(oars): Get DQN working with more than one dim in the actions.
--> 295     if len(flat_action_spec) > 1 or flat_action_spec[0].shape.rank > 0:
    296       raise ValueError(
    297           'Only scalar actions are supported now, but action spec is: {}'

AttributeError: 'tuple' object has no attribute 'rank'
  In call to configurable 'DqnAgent' (<class 'tf_agents.agents.dqn.dqn_agent.DqnAgent'>)

What I have tried: Following the suggestions here The modified

self._action_spec = tf_agents.specs.BoundedArraySpec(shape=(), dtype=np.int32, name="action", minimum=0, maximum=3)
self._observation_spec = tf_agents.specs.BoundedArraySpec(shape=(4, 4), dtype=np.int32, name="observation", minimum=0, maximum=1)

to:

self._action_spec = tf_agents.specs.BoundedTensorSpec(
    shape=(), dtype=np.int32, name="action", minimum=0, maximum=3)
self._observation_spec = tf_agents.specs.BoundedTensorSpec(
    shape=(4, 4), dtype=np.int32, name="observation", minimum=0, maximum=1)

However, this resulted in:

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-53-ce737b2b13fd> in <module>()
     21 
     22 
---> 23 agent = DqnAgent(tf_env.time_step_spec(),
     24     tf_env.action_spec(),
     25     q_network=q_net,

1 frames
/usr/local/lib/python3.6/dist-packages/tf_agents/environments/py_environment.py in time_step_spec(self)
    147       the step_type, reward, discount, and observation structure.
    148     """
--> 149     return ts.time_step_spec(self.observation_spec(), self.reward_spec())
    150 
    151   def current_time_step(self) -> ts.TimeStep:

/usr/local/lib/python3.6/dist-packages/tf_agents/trajectories/time_step.py in time_step_spec(observation_spec, reward_spec)
    388           'Expected observation and reward specs to both be either tensor or '
    389           'array specs, but saw spec values {} vs. {}'
--> 390           .format(first_observation_spec, first_reward_spec))
    391   if isinstance(first_observation_spec, tf.TypeSpec):
    392     return TimeStep(

TypeError: Expected observation and reward specs to both be either tensor or array specs, but saw spec values BoundedTensorSpec(shape=(4, 4), dtype=tf.int32, name='observation', minimum=array(0, dtype=int32), maximum=array(1, dtype=int32)) vs. ArraySpec(shape=(), dtype=dtype('float32'), name='reward')

I understand the reward is the issue: So, added an extra line

self._reward_spec = tf_agents.specs.TensorSpec((1,), np.dtype('float32'), 'reward')

but still resulted in the same error. Is there anyway I can solve this:

Answer 1

You cannot use TensorSpec with PyEnvironment class objects, this is why your attempted solution does not work. A simple fix should be to use the original code

self._action_spec = tf_agents.specs.BoundedArraySpec(shape=(), dtype=np.int32, name="action", minimum=0, maximum=3)
self._observation_spec = tf_agents.specs.BoundedArraySpec(shape=(4, 4), dtype=np.int32, name="observation", minimum=0, maximum=1)

And then wrap your environment like so:

env= MyEnvironment()
tf_env = tf_agents.environments.tf_py_environment.TFPyEnvironment(env)

This is the simplest thing. Alternatively you could define your environment as a TFEnvironment class object, use TensorSpec and change all your environment code to operate on tensors.I do not recommend this to a beginner...