Dict space from gym support

[rndD]

I'm trying to use my own custom OpenAI gym with tf-agents. So I load it with suite_gym from tf-agent: env = suite_gym.load(env_name, max_episode_steps=max_episode_steps) train_env = tf_py_environment.TFPyEnvironment(env)

But my observation space is a dict type space form OpenAI gym lib and it throws an error on agent initialization:

File "learn_tf_dqn.py", line 60, in <module>
    tf_agent.initialize()
  File "/Users/alex-k/.local/lib/python3.6/site-packages/tf_agents/agents/tf_agent.py", line 107, in initialize
    return self._initialize_fn()
  File "/Users/alex-k/.local/lib/python3.6/site-packages/tf_agents/agents/dqn/dqn_agent.py", line 236, in _initialize
    self._q_network.variables, self._target_q_network.variables, tau=1.0)
  File "/Users/alex-k/.local/lib/python3.6/site-packages/tf_agents/networks/network.py", line 129, in variables
    self._build()
  File "/Users/alex-k/.local/lib/python3.6/site-packages/tf_agents/networks/network.py", line 111, in _build
    self.__call__(random_input, step_type, None)
  File "/Users/alex-k/.local/lib/python3.6/site-packages/tf_agents/networks/network.py", line 153, in __call__
    return super(Network, self).__call__(inputs, *args, **kwargs)
  File "/Users/alex-k/code/hobby/ml/minicodna/envs/card_game/lib/python3.6/site-packages/tensorflow/python/keras/engine/base_layer.py", line 712, in __call__
    outputs = self.call(inputs, *args, **kwargs)
  File "/Users/alex-k/.local/lib/python3.6/site-packages/tf_agents/networks/q_network.py", line 134, in call
    observation, step_type=step_type, network_state=network_state)
  File "/Users/alex-k/.local/lib/python3.6/site-packages/tf_agents/networks/network.py", line 153, in __call__
    return super(Network, self).__call__(inputs, *args, **kwargs)
  File "/Users/alex-k/code/hobby/ml/minicodna/envs/card_game/lib/python3.6/site-packages/tensorflow/python/keras/engine/base_layer.py", line 712, in __call__
    outputs = self.call(inputs, *args, **kwargs)
  File "/Users/alex-k/.local/lib/python3.6/site-packages/tf_agents/networks/encoding_network.py", line 261, in call
    states = layer(states)
  File "/Users/alex-k/code/hobby/ml/minicodna/envs/card_game/lib/python3.6/site-packages/tensorflow/python/keras/engine/base_layer.py", line 712, in __call__
    outputs = self.call(inputs, *args, **kwargs)
  File "/Users/alex-k/code/hobby/ml/minicodna/envs/card_game/lib/python3.6/site-packages/tensorflow/python/keras/layers/core.py", line 580, in call
    inputs, (tensor_shape.dimension_value(inputs.shape[0]) or
AttributeError: 'dict' object has no attribute 'shape'

It looks like a problem with a dict type input. Is it supported now in tf-agents? Maybe I need to do some workaround with observation space in my gym?