Accessing tf.Module object for RT-1-X checkpoint

Hi, Thanks for sharing the datasets and the models.

I am trying to get the RT-1-X model (rt_1_x_tf_trained_for_002272480_step.zip) working in PyTorch by transpiling the TensorFlow model using Ivy. I am able to load the model in a TF Agents object and make inferences, but I have inspected the TF Agents and policy objects and I can't find a way to extract a tf.Module object with the RT-1-X model from there, which is the object I need to transpile to PyTorch using Ivy. I can only see the layers and weights in tf.Variable and tensor format only, but not the tf.Module object. Is there a way to get this object?

Also, what is the reason to use TF Agents in this case? The model is trained using imitation learning, not reinforcement learning, so TF Agents wouldn't be needed, right?

Thank you very much in advance.

Kind regards,

Victor Montesinos