Overlaying inference results with training data using jax checkpoint

[zhiyuan-zhang0206]

Hi there,
Thanks for your incredible work and open-sourcing.
Could you guys give an example of running inference and overlaying inferece results using the jax checkpoint?

I recently ran into a problem trying to use jax checkpoint and reproduce the inference example using tensorflow. When I try to run the tf inference example on the first piece of data of bridge, and visualize only the [15:] frames, from Minimal_example_for_running_inference_using_RT_1_X_TF_using_tensorflow_datasets.ipynb, I get:

which is pretty good. But when I try to use rt1_inference_example.py and run on the same data, I get:

which is drastically different from the results using tf checkpoint. The tf checkpoint comes without code (I think) so it's really hard for me to debug. So I am wondering if you can release an example of running inference using the jax checkpoint and overlaying the results with training data? That would really help. Thanks in advance!