Failed to import from tf_agents.agents.categorical_dqn import categorical_dqn_agent

Im getting this error when importing this command on python from tf_agents.agents.categorical_dqn import categorical_dqn_agent

Traceback (most recent call last):
  File "<pyshell#19>", line 1, in <module>
    from tf_agents.agents.categorical_dqn import categorical_dqn_agent
  File "C:\Users\kasun\AppData\Local\Programs\Python\Python310\lib\site-packages\tf_agents\agents\__init__.py", line 17, in <module>
    from tf_agents.agents import tf_agent
  File "C:\Users\kasun\AppData\Local\Programs\Python\Python310\lib\site-packages\tf_agents\agents\tf_agent.py", line 26, in <module>
    from tf_agents.specs import tensor_spec
  File "C:\Users\kasun\AppData\Local\Programs\Python\Python310\lib\site-packages\tf_agents\specs\__init__.py", line 20, in <module>
    from tf_agents.specs.distribution_spec import DistributionSpec
  File "C:\Users\kasun\AppData\Local\Programs\Python\Python310\lib\site-packages\tf_agents\specs\distribution_spec.py", line 22, in <module>
    import tensorflow_probability as tfp
  File "C:\Users\kasun\AppData\Local\Programs\Python\Python310\lib\site-packages\tensorflow_probability\__init__.py", line 23, in <module>
    from tensorflow_probability.python import *  # pylint: disable=wildcard-import
  File "C:\Users\kasun\AppData\Local\Programs\Python\Python310\lib\site-packages\tensorflow_probability\python\__init__.py", line 138, in <module>
    dir(globals()[pkg_name])  # Forces loading the package from its lazy loader.
  File "C:\Users\kasun\AppData\Local\Programs\Python\Python310\lib\site-packages\tensorflow_probability\python\internal\lazy_loader.py", line 57, in __dir__
    module = self._load()
  File "C:\Users\kasun\AppData\Local\Programs\Python\Python310\lib\site-packages\tensorflow_probability\python\internal\lazy_loader.py", line 40, in _load
    module = importlib.import_module(self.__name__)
  File "C:\Users\kasun\AppData\Local\Programs\Python\Python310\lib\importlib\__init__.py", line 126, in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
  File "C:\Users\kasun\AppData\Local\Programs\Python\Python310\lib\site-packages\tensorflow_probability\python\experimental\__init__.py", line 31, in <module>
    from tensorflow_probability.python.experimental import bijectors
  File "C:\Users\kasun\AppData\Local\Programs\Python\Python310\lib\site-packages\tensorflow_probability\python\experimental\bijectors\__init__.py", line 17, in <module>
    from tensorflow_probability.python.bijectors.ldj_ratio import forward_log_det_jacobian_ratio
  File "C:\Users\kasun\AppData\Local\Programs\Python\Python310\lib\site-packages\tensorflow_probability\python\bijectors\__init__.py", line 19, in <module>
    from tensorflow_probability.python.bijectors.absolute_value import AbsoluteValue
  File "C:\Users\kasun\AppData\Local\Programs\Python\Python310\lib\site-packages\tensorflow_probability\python\bijectors\absolute_value.py", line 19, in <module>
    from tensorflow_probability.python.bijectors import bijector
  File "C:\Users\kasun\AppData\Local\Programs\Python\Python310\lib\site-packages\tensorflow_probability\python\bijectors\bijector.py", line 26, in <module>
    from tensorflow_probability.python.internal import batch_shape_lib
  File "C:\Users\kasun\AppData\Local\Programs\Python\Python310\lib\site-packages\tensorflow_probability\python\internal\batch_shape_lib.py", line 23, in <module>
    from tensorflow_probability.python.internal import prefer_static as ps
  File "C:\Users\kasun\AppData\Local\Programs\Python\Python310\lib\site-packages\tensorflow_probability\python\internal\prefer_static.py", line 26, in <module>
    from tensorflow_probability.python.internal.backend import numpy as nptf
  File "C:\Users\kasun\AppData\Local\Programs\Python\Python310\lib\site-packages\tensorflow_probability\python\internal\backend\numpy\__init__.py", line 18, in <module>
    from tensorflow_probability.python.internal.backend.numpy import bitwise
  File "C:\Users\kasun\AppData\Local\Programs\Python\Python310\lib\site-packages\tensorflow_probability\python\internal\backend\numpy\bitwise.py", line 19, in <module>
    from tensorflow_probability.python.internal.backend.numpy import _utils as utils
  File "C:\Users\kasun\AppData\Local\Programs\Python\Python310\lib\site-packages\tensorflow_probability\python\internal\backend\numpy\_utils.py", line 22, in <module>
    from tensorflow_probability.python.internal.backend.numpy import nest
  File "C:\Users\kasun\AppData\Local\Programs\Python\Python310\lib\site-packages\tensorflow_probability\python\internal\backend\numpy\nest.py", line 68, in <module>
    assert_same_structure = dm_tree.assert_same_structure
AttributeError: module 'tree' has no attribute 'assert_same_structure'

I tried this:

pip3 install dm-tree but it also didnt work

Try:

pip install --upgrade dm-tree