Help for DQN with nested action and observation space

[armusone]

Hello everybody,
I have been trying to create a custom DQN for a project but couldn't find any specification on networks with multiple inputs and multiple actions.
So far I created a custom environment with the following action_spec
self._action_spec = { 'Action1': array_spec.BoundedArraySpec( shape=(), dtype=np.int32, minimum=0, maximum=5, name='action1' ), 'Action2': array_spec.BoundedArraySpec( shape=(), dtype=np.float32, minimum=0, maximum=1000, name='action2' ) }
and observation_spec
self._observation_spec = { 'Obs1': array_spec.ArraySpec( shape=(3, 82), dtype=np.float32, name='obs1' ), 'Obs2': array_spec.ArraySpec( shape=(6,), dtype=np.float32, name='obs2' ) }

I tested the environment with utils.validate_py_environment(environment, episodes=5) and I got no error.
The model I am trying to build is as follow:
input_one = tf.keras.layers.Input((3,83), name='Obs1')
flatten = tf.keras.layers.Flatten()(input_one )
dense_one = tf.keras.layers.Dense(60, name='Dense1_Obs1')(flatten)
dense_two = tf.keras.layers.Dense(30, name='Dense2_Obs1')(dense_one )
action_one = tf.keras.layers.Dense(5, name='Action1')(dense_two )
input_two = tf.keras.layers.Input((6,), name='Obs2')
dense_three = tf.keras.layers.Dense(3, name='Dense1_Obs2')(input_two )
concatenate = tf.keras.layers.concatenate([dense_three , action_one ])
action_two= tf.keras.layers.Dense(1, name='Action2')(concatenate)
my_model= tf.keras.Model(inputs=[input_one , input_two ], outputs=[action_one , action_two])

I then used q_net = tf_agents.networks.Sequential([my_model]) as described in #457 and initialize the agent as:
optimizer = tf.keras.optimizers.Adam(learning_rate=learning_rate)
train_step_counter = tf.Variable(0)
agent = dqn_agent.DqnAgent( train_env.time_step_spec(), train_env.action_spec(), q_network=q_net, optimizer=optimizer, td_errors_loss_fn=common.element_wise_squared_loss, train_step_counter=train_step_counter)
agent.initialize()

But got the following error:
Only scalar actions are supported now, but action spec is: {'Action1': BoundedTensorSpec(shape=(), dtype=tf.int32, name='action1', minimum=array(0, dtype=int32), maximum=array(5, dtype=int32)), 'Action2': BoundedTensorSpec(shape=(), dtype=tf.float32, name='action2', minimum=array(0., dtype=float32), maximum=array(1000., dtype=float32))} In call to configurable 'DqnAgent' (<class 'tf_agents.agents.dqn.dqn_agent.DqnAgent'>)

Could you please help me with this issue?
Thank you very much,
Antonio