Language Embedding in Jax Example

Hello,

Thank you for making this code available. I'm currently experimenting with the jax example, and I am experiencing some issues with using an actual natural language embedding instead of the dummy data.

I tried the following:
```
embed = hub.load('https://tfhub.dev/google/universal-sentence-encoder-large/5')

language_instruction = "Pick up the yellow plush toy."
language_embedding = embed([self.language_instruction])[0]

observation = {
# with a seqlen of 15, to my understanding I need to use an array of this shape:
'natural_language_embedding': np.array([language_embedding for i in range(0,15)])
'image': img_array,
}

act = policy.action(observation)
```
This leads to the policy outputting constant actions, meaning the output action tokens are always the same, independent of the image input. I have tried this with actual images from the robot environment and with randomly generated images, as well as with a variety of language instructions, the output always stays the same. When instead of the language embedding I use `np.ones((15, 512))` in the observation, as given in the example, the action tokens are NOT constant, as I would expect them to behave also with a language embedding.

Any help on where this issue comes from and how I can resolve it would be greatly appreciated.

Thanks, Jonathan