Benchmarking octo-small-v1.5 with SimplerEnv

[PommesPeter]

Hi, thank for your great work on benchmarking robotic foundation model on real-to-sim setting.

We encountered an issue when inferring the Octo model. The current branch only supports inference with Octo v1.0, but we want to compare the differences between Octo v1.0 and Octo v1.5, using the octo-small model as our test model. We make some changes to let current branch supported inference on Octo v1.5 code (the latest branch main in Octo repo).

After running inference, we found that using the Octo v1.0 code produced results very close to those reported in the SimpleEenv paper. However, we discovered that the performance was significantly worse. What could be the reason for this?

Here are the parts of the code we changed in octo_model.py (octo v1.5 code):

- input_observation = {"image_primary": images, "pad_mask": pad_mask}
+ input_observation = {"image_primary": images, "timestep_pad_mask": pad_mask}
norm_raw_actions = self.model.sample_actions(
    input_observation,
    self.task,
+   unnormalization_statistics=self.model.dataset_statistics["bridge_dataset"]["action"], 
    rng=key,
)
- raw_actions = norm_raw_actions * self.action_std[None] + self.action_mean[None]
+ # raw_actions = norm_raw_actions * self.action_std[None] + self.action_mean[None]
raw_actions = raw_actions[0]  # remove batch, becoming (action_pred_horizon, action_dim)

And here are the comparative experiments I ran:

Running scripts: octo_bridge.sh
Environment: same as README.md
Python version: 3.10

Policy	Put Carrot On Plate	Put Spoon On Table Cloth	Stack Green Cube On Yellow Cube BakedTex	Put Eggplant In Basket
Octo-small v1.0 (original code w/o modified)	8.30%	48.61%	0.00%	54.17%
Octo-small v1.5 (using octo v1.5 code above)	0.00%	6.25%	0.00%	4.17%

Why are there two different results? Is there a problem with the code modifications we made above?

Thank you again for your contributions to the community!

Looking forward to your reply!

[xuanlinli17]

I've not tried Octo v1.5 but looks like there might be issues with either normalization or seed; note that octo requires a different seed at each inference step since it uses diffusion action head, and this is implemented by using jax split seed; otherwise its performance will be significantly worse.

[PommesPeter]

I've not tried Octo v1.5 but looks like there might be issues with either normalization or seed; note that octo requires a different seed at each inference step since it uses diffusion action head, and this is implemented by using jax split seed; otherwise its performance will be significantly worse.

Thank you for your response. we also do jax.random.split(self.rng) and make it use different seed at each inference step. I think we should check the normalization if it is correct.

For normalization, I found that unnormalization_statistics argument will be passed to sample_actions function in Octo v1.5 and normalize the action, however Octo v1.0 is not.

Furthermore, I noticed there are some differences in ActionHead.predict_action function.