AttributeError: 'dict' object has no attribute 'policy'

I0000 00:00:1744804297.629139 1896187 gpu_device.cc:2019] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 1580 MB memory: -> device: 0, name: NVIDIA GeForce RTX 3090, pci bus id: 0000:09:00.0, compute capability: 8.6
(101,) (101,)
(101,) (101,)
(1, 14)
Traceback (most recent call last):
File "/home/jingze/learn/Robotics/opensource/Open_Duck_Playground/playground/open_duck_mini_v2/runner.py", line 64, in
main()
File "/home/jingze/learn/Robotics/opensource/Open_Duck_Playground/playground/open_duck_mini_v2/runner.py", line 60, in main
runner.train()
File "/home/jingze/learn/Robotics/opensource/Open_Duck_Playground/playground/common/runner.py", line 114, in train
_, params, _ = train_fn(
^^^^^^^^^
File "/home/jingze/learn/Robotics/opensource/Open_Duck_Playground/.venv/lib/python3.11/site-packages/brax/training/agents/ppo/train.py", line 713, in train
policy_params_fn(current_step, make_policy, params)
File "/home/jingze/learn/Robotics/opensource/Open_Duck_Playground/playground/common/runner.py", line 78, in policy_params_fn
export_onnx(
File "/home/jingze/learn/Robotics/opensource/Open_Duck_Playground/playground/common/export_onnx.py", line 150, in export_onnx
transfer_weights(params[1].policy["params"], tf_policy_network)
^^^^^^^^^^^^^^^^
AttributeError: 'dict' object has no attribute 'policy'