Read less bytes than requested

[HRItdy]

Hi Sebbyjp,

Really thanks for your excellent work! I deploy your code on my desktop. It worked well with rt1main. But when I try to change the model to rt1multirobot and rt1realsim, there is an error raised:

Traceback (most recent call last):
  File "/home/lab_cheem/robo_transformers/UR5_test.py", line 175, in <module>
    node = InferenceNode()
  File "/home/lab_cheem/robo_transformers/UR5_test.py", line 34, in __init__
    self.inference_server = InferenceServer()
  File "<@beartype(robo_transformers.inference_server.InferenceServer.__init__) at 0x7f4929a15430>", line 88, in __init__
  File "/home/lab_cheem/robo_transformers/robo_transformers/inference_server.py", line 63, in __init__
    self.agent: Agent = REGISTRY[model_type]['agent'](weights_key)
  File "<@beartype(robo_transformers.models.rt1.agent.RT1Agent.__init__) at 0x7f4934af78b0>", line 31, in __init__
  File "/home/lab_cheem/robo_transformers/robo_transformers/models/rt1/agent.py", line 16, in __init__
    self.model: PyPolicy | TFPolicy = load_rt1(model_key=weights_key)
  File "/home/lab_cheem/robo_transformers/robo_transformers/models/rt1/inference.py", line 118, in load_rt1
    policy: LoadedPolicy = LoadedPolicy(
  File "/home/lab_cheem/miniforge3/envs/RT1-3.9/lib/python3.9/site-packages/gin/config.py", line 1605, in gin_wrapper
    utils.augment_exception_message_and_reraise(e, err_str)
  File "/home/lab_cheem/miniforge3/envs/RT1-3.9/lib/python3.9/site-packages/gin/utils.py", line 41, in augment_exception_message_and_reraise
    raise proxy.with_traceback(exception.__traceback__) from None
  File "/home/lab_cheem/miniforge3/envs/RT1-3.9/lib/python3.9/site-packages/gin/config.py", line 1582, in gin_wrapper
    return fn(*new_args, **new_kwargs)
  File "/home/lab_cheem/miniforge3/envs/RT1-3.9/lib/python3.9/site-packages/tf_agents/policies/py_tf_eager_policy.py", line 193, in __init__
    policy = tf.compat.v2.saved_model.load(model_path)
  File "/home/lab_cheem/miniforge3/envs/RT1-3.9/lib/python3.9/site-packages/tensorflow/python/saved_model/load.py", line 912, in load
    result = load_partial(export_dir, None, tags, options)["root"]
  File "/home/lab_cheem/miniforge3/envs/RT1-3.9/lib/python3.9/site-packages/tensorflow/python/saved_model/load.py", line 1043, in load_partial
    loader = Loader(object_graph_proto, saved_model_proto, export_dir,
  File "/home/lab_cheem/miniforge3/envs/RT1-3.9/lib/python3.9/site-packages/tensorflow/python/saved_model/load.py", line 226, in __init__
    self._restore_checkpoint()
  File "/home/lab_cheem/miniforge3/envs/RT1-3.9/lib/python3.9/site-packages/tensorflow/python/saved_model/load.py", line 561, in _restore_checkpoint
    load_status = saver.restore(variables_path, self._checkpoint_options)
  File "/home/lab_cheem/miniforge3/envs/RT1-3.9/lib/python3.9/site-packages/tensorflow/python/checkpoint/checkpoint.py", line 1439, in restore
    object_graph_string = reader.get_tensor(base.OBJECT_GRAPH_PROTO_KEY)
  File "/home/lab_cheem/miniforge3/envs/RT1-3.9/lib/python3.9/site-packages/tensorflow/python/training/py_checkpoint_reader.py", line 66, in get_tensor
    return CheckpointReader.CheckpointReader_GetTensor(
IndexError: Read less bytes than requested
  In call to configurable 'SavedModelPyTFEagerPolicy' (<class 'tf_agents.policies.py_tf_eager_policy.SavedModelPyTFEagerPolicy'>)

The models are downloaded from https://github.com/google-research/robotics_transformer.git. It seems to be a mismatch of the model parameters. Do you have any idea about this? Any suggestion is appreciated. Thanks!