Repository: simpler-env/SimplerEnv
Issue: Benchmarking `octo-small-v1.5` with SimplerEnv
Root cause: Interface Contract Mismatch

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](https://github.com/octo-models/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):

```diff
- 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!

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.

> 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.