The issue in loading the trajectory dataset

I am trying to load the trajectory dataset using the [Open_X_Embodiment_Datasets.ipynb](https://colab.research.google.com/github/google-deepmind/open_x_embodiment/blob/main/colabs/Open_X_Embodiment_Datasets.ipynb#scrollTo=t2V0xrIVMWNc)

Unfortunately, I get unexpected data when trying to use the **Demonstration of transformation from an episode to a trajectory** to load the **berkeley_gnm_cory_hall**, **berkeley_gnm_recon**, **berkeley_gnm_sac_son** datasets.
For example, when loading the **berkeley_gnm_sac_son**, I got:
```
{'is_first': ,
'is_last': ,
'action': ,
'is_terminal': ,
'observation': {'yaw': ,
'position': ,
'state': ,
'image': }}
```
,where the image shape should be **(3, 120, 160, 3)** instead of **(3,)**

The code for this example listed below:
```
import tensorflow_datasets as tfds
mt_opt_builder = tfds.builder_from_directory(builder_dir='gs://gresearch/robotics/berkeley_gnm_sac_son/0.1.0/')

mt_opt_episodic_dataset = mt_opt_builder.as_dataset(split='train[:10]')
mt_opt_rlds_spec = RLDSSpec(
observation_info=b.info.features['steps']['observation'],
action_info=b.info.features['steps']['action'],
)

# The following will create a trajectories of length 3.
trajectory_length = 3
trajectory_transform = TrajectoryTransformBuilder(mt_opt_rlds_spec, pattern_fn=n_step_pattern_builder(trajectory_length)).build(validate_expected_tensor_spec=False)
trajectory_dataset = trajectory_transform.transform_episodic_rlds_dataset(mt_opt_episodic_dataset)

trajectory_iter = iter(trajectory_dataset)
trajectory = next(trajectory_iter)
trajectory
```
How can I solve this problem?