InvalidArgumentError: Exception encountered when calling layer "dynamic_unroll" (type DynamicUnroll)

[amarshivaram]

I am trying to run a QRNN Network with Neural LinUCB Agent.
My parameters are as follows:

AGENT_ALPHA = 1
GAMMA = 1
EPS_GREEDY = 0
LAYERS = (512,256,128)
LR = 0.005
NUM_TRAINING_STEPS = 1000
ENCODING_DIM = 7
batch_size = 64

I have defined my QRNN Network in the following way:

network = q_rnn_network.QRnnNetwork( input_tensor_spec=env.time_step_spec().observation, action_spec=env.action_spec(), lstm_size=(28,), # rnn_construction_fn=tf.keras.layers.LSTM, # rnn_construction_kwargs={"units":28}, input_fc_layer_params=LAYERS)
And my agent is defined in the following manner:

agent = neural_linucb_agent.NeuralLinUCBAgent( time_step_spec = env.time_step_spec(), action_spec = env.action_spec(), alpha = AGENT_ALPHA, gamma = GAMMA, epsilon_greedy = EPS_GREEDY, gradient_clipping = 1.0, dtype = tf.float32, optimizer=tf.compat.v1.train.AdamOptimizer(learning_rate=LR), error_loss_fn=tf.compat.v1.losses.mean_squared_error, encoding_network = network, encoding_network_num_train_steps = NUM_TRAINING_STEPS, encoding_dim = ENCODING_DIM)

I'm running the experiment in colab. And my tf-agent is of version 0.11.0
However, when performing agent.train(replay_buffer.gather_all()) , I encounter the following error

InvalidArgumentError: Exception encountered when calling layer "dynamic_unroll" (type DynamicUnroll)

The error stack trace is as follows:

---------------------------------------------------------------------------
InvalidArgumentError                      Traceback (most recent call last)
<ipython-input-68-a2ccbe26f502> in <module>()
     35   driver.run()
     36   print("Driver run done")
---> 37   loss_info = agent.train(replay_buffer.gather_all())
     38   # loss_info = agent.train(replay_buffer.as_dataset(sample_batch_size=batch_size, num_steps=steps_per_loop,single_deterministic_pass=False))
     39   # loss_info = agent.train(replay_buffer.get_next())

13 frames
/usr/local/lib/python3.7/dist-packages/tf_agents/agents/tf_agent.py in train(self, experience, weights, **kwargs)
    330     if self._enable_functions:
    331       loss_info = self._train_fn(
--> 332           experience=experience, weights=weights, **kwargs)
    333     else:
    334       loss_info = self._train(experience=experience, weights=weights, **kwargs)

/usr/local/lib/python3.7/dist-packages/tf_agents/utils/common.py in with_check_resource_vars(*fn_args, **fn_kwargs)
    186         # We're either in eager mode or in tf.function mode (no in-between); so
    187         # autodep-like behavior is already expected of fn.
--> 188         return fn(*fn_args, **fn_kwargs)
    189       if not resource_variables_enabled():
    190         raise RuntimeError(MISSING_RESOURCE_VARIABLES_ERROR)

/usr/local/lib/python3.7/dist-packages/tf_agents/bandits/agents/neural_linucb_agent.py in _train(self, experience, weights)
    607         self.actions_from_reward_layer,
    608         use_actions_from_reward_layer,
--> 609         no_actions_from_reward_layer)
    610     return loss_info

/usr/local/lib/python3.7/dist-packages/tensorflow/python/util/traceback_utils.py in error_handler(*args, **kwargs)
    151     except Exception as e:
    152       filtered_tb = _process_traceback_frames(e.__traceback__)
--> 153       raise e.with_traceback(filtered_tb) from None
    154     finally:
    155       del filtered_tb

/usr/local/lib/python3.7/dist-packages/tf_agents/bandits/agents/neural_linucb_agent.py in use_actions_from_reward_layer()
    598     def use_actions_from_reward_layer():
    599       return self.compute_loss_using_reward_layer(
--> 600           observation=observation, action=action, reward=reward, weights=weights, step_type=step_type, training=True)
    601 
    602     def no_actions_from_reward_layer():

/usr/local/lib/python3.7/dist-packages/tf_agents/bandits/agents/neural_linucb_agent.py in compute_loss_using_reward_layer(self, observation, action, reward, step_type, weights, training)
    395     with tf.GradientTape() as tape:
    396       loss_info = self._loss_using_reward_layer(
--> 397           observations=observation, actions=action, rewards=reward, weights=weights, step_type=step_type, training=training)
    398     tf.debugging.check_numerics(loss_info[0], 'Loss is inf or nan')
    399     tf.compat.v2.summary.scalar(

/usr/local/lib/python3.7/dist-packages/tf_agents/bandits/agents/neural_linucb_agent.py in _loss_using_reward_layer(self, observations, actions, rewards, step_type, weights, training)
    342     with tf.name_scope('loss'):
    343       encoded_observation, _ = self._encoding_network(
--> 344           observations, step_type, training=training)
    345       encoded_observation = tf.reshape(
    346           encoded_observation,

/usr/local/lib/python3.7/dist-packages/tf_agents/networks/network.py in __call__(self, inputs, *args, **kwargs)
    424       normalized_kwargs.pop("network_state", None)
    425 
--> 426     outputs, new_state = super(Network, self).__call__(**normalized_kwargs)  # pytype: disable=attribute-error  # typed-keras
    427 
    428     nest_utils.assert_matching_dtypes_and_inner_shapes(

/usr/local/lib/python3.7/dist-packages/keras/utils/traceback_utils.py in error_handler(*args, **kwargs)
     65     except Exception as e:  # pylint: disable=broad-except
     66       filtered_tb = _process_traceback_frames(e.__traceback__)
---> 67       raise e.with_traceback(filtered_tb) from None
     68     finally:
     69       del filtered_tb

/usr/local/lib/python3.7/dist-packages/tf_agents/networks/lstm_encoding_network.py in call(self, observation, step_type, network_state, training)
    272         initial_state=network_state,
    273         training=training,
--> 274         **network_kwargs)
    275 
    276     if isinstance(self._lstm_network, dynamic_unroll_layer.DynamicUnroll):

/usr/local/lib/python3.7/dist-packages/tf_agents/keras_layers/dynamic_unroll_layer.py in call(self, inputs, initial_state, reset_mask, training)
    326           state=initial_state,
    327           zero_state=zero_state,
--> 328           training=training)
    329     else:
    330       outputs, new_state = _dynamic_unroll_multi_step(

/usr/local/lib/python3.7/dist-packages/tf_agents/keras_layers/dynamic_unroll_layer.py in _static_unroll_single_step(cell, inputs, reset_mask, state, zero_state, training)
    381     state = tf.nest.map_structure(
    382         lambda s, s_zero: _maybe_reset_state(reset_mask, s_zero, s), state,
--> 383         zero_state)
    384 
    385   outputs, final_state = cell(inputs, state, training=training)

/usr/local/lib/python3.7/dist-packages/tf_agents/keras_layers/dynamic_unroll_layer.py in <lambda>(s, s_zero)
    380     reset_mask = _squeeze(reset_mask)
    381     state = tf.nest.map_structure(
--> 382         lambda s, s_zero: _maybe_reset_state(reset_mask, s_zero, s), state,
    383         zero_state)
    384 

/usr/local/lib/python3.7/dist-packages/tf_agents/keras_layers/dynamic_unroll_layer.py in _maybe_reset_state(reset, s_zero, s)
    357   print(s)
    358   if not isinstance(s, tf.TensorArray) and s.shape.rank > 0:
--> 359     return tf.compat.v1.where(reset, s_zero, s)
    360   else:
    361     return s

InvalidArgumentError: Exception encountered when calling layer "dynamic_unroll" (type DynamicUnroll).

Inputs to operation Select of type Select must have the same size and shape.  Input 0: [1,64] != input 1: [64,28] [Op:Select]

Call arguments received:
  • inputs=tf.Tensor(shape=(64, 1, 128), dtype=float32)
  • initial_state=()
  • reset_mask=tf.Tensor(shape=(64, 1, 1), dtype=bool)
  • training=True

[n-chennakeshava]

Hey there! I am having the same issue. Did you manage to find a solution?

[lazZ33]

Same here as well