Paper: SEABO: A Simple Search-Based Method for Offline Imitation Learning (ICLR-2024)

Reviewer comment (Theory):
The paper lacks theoretical justifications, which makes understanding some parts a little difficult.

Author response:
**Concern 2: the paper lacks theoretical justifications**

We would like to clarify that the main contribution of this paper lies in the empirical side. Many previous offline IL papers like OTR [3] also do not include theoretical analysis. Our concern lies in the fact that theoretical analysis often relies heavily on assumptions that may be hard to satisfy in practice, resulting in a gap between theory and empirical results.

[3] Luo, Y., Jiang, Z., Cohen, S., Grefenstette, E., Deisenroth, M. P. Optimal Transport for Offline Imitation Learning.

Reviewer comment (Theory):
Is there an inherent assumption being made between the distribution of state-action pairs in the expert demonstrations and the unlabeled dataset?

Author response:
**Concern 4: Is there an inherent assumption being made between the distribution of state-action pairs in the expert demonstrations and the unlabeled dataset?**

SEABO does not involve a distribution matching procedure and hence we believe no inherent assumption is needed. We decide the optimality of the single transition by querying its nearest neighbor in the expert trajectory. A large reward is assigned if their deviation is small, and a smaller reward otherwise. The rewards can be computed anyway. We do not make assumptions on the state-action distribution of offline datasets, and the offline dataset can contain similar trajectories, small amounts of demonstrations or a mixture of trajectories with diversified returns, etc.

Reviewer comment (Experiments):
Does the stochasticity of environments affect the performance of SEBO?

Author response:
**Concern 1: evaluations restricted to deterministic environments**

We agree that our evaluations are conducted in deterministic environments. We respectfully argue that most offline IL/offline RL algorithms are evaluated on deterministic environments, and the dynamics in many scenarios (e.g., robotic manipulation) are completely deterministic. Meanwhile, we do not claim that our method is all-around and can handle all scenarios (it is difficult for one single algorithm to do so). With only one expert trajectory, SEABO may fail to produce satisfying rewards for learning in stochastic environments. However, we do not think the *potential* incapability of handling stochastic environments downgrades the contribution of our work.

We agree that, intuitively, there may exist an inherent correlation between the number of expert trajectories needed and the stochasticity of the environment. We deem that it is interesting to check whether the stochasticity of environments affects the performance of SEABO. Nevertheless, as far as we can tell, there are few standard, widely adopted benchmarks with stochastic continuous control environments. Some prior work (e.g., stochastic D4RL Mujoco suite in Appendix C.2 in [1]) introduce stochasticity to the environment by modifying the reward in the dataset, which is somewhat incompatible with SEABO since there are no rewards in the dataset in our setting. We then introduce stochasticity by injecting Gaussian noises (with mean 0 and std $\sigma$) to observations during evaluation, following prior work [2]. We try std=0.05 and std=0.1, respectively, and conduct experiments on MuJoCo datasets. We use IQL as the base algorithm and follow the same hyperparameter setup of SEABO in MuJoCo tasks. We summarize the results below.

| Task Name | IQL (oracle, std=0.05) | IQL+SEABO (std=0.05) | IQL (oracle, std=0.1) | IQL+SEABO (std=0.1) |
| ---- | :---: | :---: | :---: | :---: |
| halfcheetah-medium-v2 | 40.9$\pm$1.1 | 40.6$\pm$0.6 | 31.8$\pm$0.8 | 30.4$\pm$0.2 |
| hopper-medium-v2 | 45.2$\pm$1.8 | 52.4$\pm$3.1 | 33.5$\pm$2.6 | 31.1$\pm$2.8 |
| walker2d-medium-v2 | 76.8$\pm$4.0 | 80.2$\pm$3.2 | 66.6$\pm$2.2 | 74.3$\pm$5.0 |
| halfcheetah-medium-replay-v2 | 38.6$\pm$0.9 | 39.5$\pm$1.2 | 30.0$\pm$1.6 | 26.4$\pm$3.1 |
| hopper-medium-replay-v2 | 51.8$\pm$8.2 | 71.5$\pm$7.9 | 31.7$\pm$7.0 | 33.1$\pm$4.1 |
| walker2d-medium-replay-v2 | 61.4$\pm$7.5 | 69.2$\pm$4.3 | 51.8$\pm$5.5 | 53.5$\pm$7.3 |
| halfcheetah-medium-expert-v2 | 36.5$\pm$1.0 | 34.2$\pm$2.9 | 26.8$\pm$1.5 | 24.7$\pm$1.0 |
| hopper-medium-expert-v2 | 30.4$\pm$5.3 | 35.8$\pm$7.1 | 34.0$\pm$4.8 | 29.5$\pm$3.3 | 
| walker2d-medium-expert-v2 | 101.4$\pm$4.6 | 107.6$\pm$0.7 | 87.1$\pm$9.1 | 90.5$\pm$6.5 |

Table 1. Experimental results in stochastic environments. Oracle means that raw rewards are used for training. The results are averaged over 5 runs.

Based on the results, we find that under stochastic environments, the base algorithm itself fails to behave well (we observe performance collapse on many datasets like halfcheetah-medium-expert-v2). The performance of both IQL and IQL+SEABO decreases with larger noise std. Note that SEABO still outperforms vanilla rewards on many datasets under stochastic environments.

[1] Ma, Y., Jayaraman, D., Bastani, O. Conservative offline distributional reinforcement learning. 

[2] Guo, Y., Oh, J., Singh, S., Lee, H. Generative adversarial self-imitation learning.

Reviewer comment (Theory):
There is only empirical analysis of the proposed method. I believe there are certain tasks where SEABO would perform poorly, such as a cliff-walking type of task where there is a precise boundary between what is acceptable and what is a failure.

Author response:
We have to admit that our paper mainly lies on the empirical side, mainly due to the fact that theoretical analysis often incurs a gap between theory and practical implementation or empirical results, and some assumptions made for the benefit of theoretical derivation may not necessarily hold in practice. We agree that a theoretical analysis can be a nice addition to our work, and can leave that for future exploration.

It is interesting to explore how our method behaves when there is a precise boundary between what is acceptable and what is a failure. To that end, we use the risky pointmass environment from [1] (please see the environment code in https://github.com/JasonMa2016/CODAC/blob/main/env/risky_pointmass.py and the details of risky pointmass environment in Appendix C.1 of the CODAC paper) and make some slight modifications. In the original pointmass environment, the agent gets a penalty with some probability. To make the boundary more precise, we directly assign the agent a penalty if it steps into the risky region (line 107-111). That is,
```
if not self.is_safe(self.state):
    u = np.random.uniform(0, 1)
    if u > self.risk_prob:
        cost = 1
        reward -= self.risk_penalty
```
is modified to
```
if not self.is_safe(self.state):
    cost = 1
    reward -= self.risk_penalty
```

Other codes are kept unchanged. We collect the offline dataset by training an SAC agent for 1M interactions. We choose the trajectory with the highest return as the expert trajectory. We train an IQL agent upon the collected dataset, and set the reward scale $\alpha=1$ and coefficient $\beta=0.5$ for SEABO. As the environment is simple, we train the IQL agent for 500K gradient steps (hyperparameters are set to be identical to MuJoCo tasks) and summarize the experimental results below. In the risky pointmass environment, some of the data in the expert trajectory can be close to the risky region, while it turns out that IQL+SEABO still beats IQL with vanilla reward.


| Task name | $\mu\_{\rm{max}}$ | $\mu\_{\rm{min}}$ | IQL (oracle) | IQL+SEABO |
| ---- | :---: | :---: | :---: | :---: |
| risky pointmass | -1.1| -11083.3 | -149.3$\pm$62.1 | **-98.3$\pm$13.2** |

Table 1. Results of SEABO in risky pointmass across 5 different random seeds. $\mu\_{\rm{max}}$ denotes the highest return in the dataset and $\mu\_{\rm{min}}$ denotes the smallest return in the dataset.

Nevertheless, we agree that under some extreme cases (e.g., the expert trajectory is very close to the risky boundary), our method may fail. We actually do not expect that our method can handle all possible situations (and it is hard to do so). If one is concerned with the existence of a precise boundary of failure, one can add some rules to post-process the rewards calculated by SEABO (e.g., penalize those query states that are too close to the boundary).


[1] Ma, Y., Jayaraman, D., Bastani, O. Conservative offline distributional reinforcement learning.

Reviewer comment (Theory):
I hypothesize that the approach will also only work in lower dimensional control environments. This is because the method relies heavily on a distance function, and this could suffer from the curse of dimensionality in more complex environments.

Author response:
It is true that MuJoCo tasks are low-dimensional control tasks. However, in the main text, we also evaluate SEABO on Adroit tasks, which controls a 24-DoF simulated Shadow Hand robot to accomplish different tasks like hammering a nail. The observation dimension of `pen-human-v0` gives 45, and the observation dimension of `hammer-human-v0` gives 46. We further include experiments on kitchen datasets, whose observation space gives 60 (which we think can be regarded as high dimensional control tasks), and the results are shown below. Based on the results in Adroit tasks (Table 3 in the paper) and kitchen tasks, we believe SEABO can also work in high dimensional control environments (the curse of dimensionality does not seem to be an issue here).

| Task name | BC | CQL | IQL (oracle) | IQL+SEABO |
| ---- | :---: | :---: | :---: | :---: |
| kitchen-complete-v0 | 65.0 | 43.8 | 62.5 | **67.5$\pm$4.2** |
| kitchen-partial-v0 | 38.0 | 49.8 | 46.3 | **71.0$\pm$4.1** |
| kitchen-mixed-v0 | 51.5 | 51.0 | 51.0 | **55.0$\pm$3.5**|
| average score | 51.5 | 48.2 | 53.3 | **64.5** |

Table 2. Results on kitchen environments across 5 different random seeds.

However, it may be challenging to apply SEABO on pixel-based tasks (which are of high dimension and SEABO may suffer from the curse of dimensionality), which we leave as future work since we mainly focus on state-based tasks in this paper. As we stated in Appendix E, one possible solution is to leverage a pre-trained encoder to convert the high-dimensional input into low-dimensional representations. Then, we can apply SEABO to acquire rewards.

Reviewer comment (Theory):
"we hypothesize that the transition is near-optimal if it lies close to the expert trajectory" -> Is this strategy always good or does it have weaknesses? I would imagine that certain environments with discontinuities or discrete events (e.g. a car crash) would not favor this strategy.

Author response:
We agree that our hypothesis may have limitations under some certain environments. As we stated above, we do not expect our method to handle all conditions. Under those cases, we can add some rules to manually correct the reward output from SEABO (e.g., penalize those states that are too close to the failure boundary). This concern shares some similarities with **Concern 1**, please refer to our responses above.

Reviewer comment (Experiments):
There is another simple baseline commonly used in offline RL, which is sometimes referred to as percent-BC, which is to imitate the top N% of trajectories in the offline dataset (such as N=10 or N=25). As it is somewhat similar in spirit to SEABO, it would be good to see comparisons to this approach.

Author response:
Thanks for the advice. In the main text, we have included 10\%-BC as a baseline in Table 2 of the manuscript. It is easy to find that SEABO outperforms 10\%-BC on many datasets.

Reviewer comment (Experiments):
It's unclear how well the proposed method would work with visual observations where computing distances and nearest neighbors is a lot more challenging. The authors mention this as a direction for future work and this is a limitation shared by many similar works, but it should be mentioned that the proposed approach with its reliance on an analytic distance metric may be more troublesome in such scenarios than methods that use discriminator-based "distances".

Author response:
**Concern 3: commentary on limitations**

Thanks for the additional comments on the limitation part.

- For visual observations, as stated in Appendix E, it is challenging to directly query the nearest neighbor and acquire rewards with SEABO. A possible solution is to leverage pre-trained image encoders to produce low-dimensional representations. Then, we can query nearest neighbors and compute rewards in the representation space with SEABO. As we mainly focus on state-based tasks in this paper, we leave it as future work. It is possible that discriminator-based methods may produce better rewards in such scenarios.

- We include experimental results on the IKEA benchmark and kitchen environment above. Both of them are long-horizon manipulation tasks and require precise manipulation. Our results show that SEABO aids baseline methods on kitchen datasets while not seeming to have advantages over baselines on IKEA tasks (and match trajectory segments instead of individual transitions in SEABO do not bring better performance here).

Hopefully, these can resolve the concerns. If there is still something unclear, please let us know!

Reviewer comment (Presentation):
I believe that "(oracle)" should be omitted from Table 1-3 and 6. For instance, consider "IQL (oracle)": it utilizes the ground truth reward but doesn't rely on expert demonstrations. Removing the ground truth reward and integrating an additional expert demonstration does not necessarily make the task more challenging.

Author response:
**Concern 2: (oracle) should be omitted from Table 1-3 and 6**

By specifying *oracle*, we would like to emphasize that the results are obtained by running the offline RL algorithms on datasets with ground-truth rewards. We can remove these if the reviewer deems it necessary.

We thank the reviewer for raising the score to 6! We will remove the use of *(orcale)* in our next revision, before the end of the rebuttal. Thanks for your time and efforts in making our paper better.

Dear Reviewer KA4c, we have uploaded a new revision and removed *(oracle)* from all of the tables in the main text and the appendix. We also revised the corresponding descriptions in the table captions and texts. Thank you for your comments and feedback.

Reviewer comment (Experiments):
The experiments currently compare with only two Offline RL methods (IQL and TD3_BC). It would be better to include more recent baselines such as Trajectory Transformer [[1]], Diffuser [[2]], or other methods.

Author response:
**Concern 3: results on other offline RL methods**

To see the performance of SEABO upon some recent baselines, we include two additional baselines, Decision Transformer [2], and TD7 [3]. These methods are chosen due to their simplicity, good reproducibility, and high training efficiency. We use their official GitHub implementations. We adopt the identical hyperparameter setup as TD3\_BC+SEABO for all of the evaluated tasks, and it should not be difficult to reproduce our reported results below. We take the results of DT and TD7 with vanilla rewards from their original papers and summarize the results in Table 2. It can be seen that DT+SEABO achieves competitive or even better performance on numerous datasets. We also find that the performance of TD7+SEABO on many datasets is on par with TD7 trained on raw rewards. We believe the evidence indicates that our method can also work when built upon recent strong methods.

| Task Name | DT (oracle) | DT+SEABO | TD7 (oracle) | TD7+SEABO |
| ---- | :---: | :---: | :---: | :---: |
| halfcheetah-medium-v2 | 42.6$\pm$0.1 | 42.4$\pm$0.4 | 58.0$\pm$0.4 | 47.1$\pm$0.1 |
| hopper-medium-v2 | 67.6$\pm$1.0 | 66.5$\pm$5.0 | 76.1$\pm$5.1 | **85.1$\pm$7.3** |
| walker2d-medium-v2 | 74.0$\pm$1.4 | **74.1$\pm$4.7** | 91.1$\pm$7.8 | 79.0$\pm$0.4 |
| halfcheetah-medium-replay-v2 | 36.6$\pm$0.8 | **37.4$\pm$0.6** | 53.8$\pm$0.8 | 53.6$\pm$0.4 |
| hopper-medium-replay-v2 | 82.7$\pm$7.0 | 51.2$\pm$18.3 | 91.1$\pm$8.0 | **98.0$\pm$0.6** |
| walker2d-medium-replay-v2 | 66.6$\pm$3.0 | 50.6$\pm$15.6 | 89.7$\pm$4.7 | 70.1$\pm$16.8 |
| halfcheetah-medium-expert-v2 | 86.8$\pm$1.3 | **89.8$\pm$0.6** | 104.6$\pm$1.6 | **105.9$\pm$1.8** |
| hopper-medium-expert-v2 | 107.6$\pm$1.8 | 94.7$\pm$8.2 | 108.2$\pm$4.8 | 104.7$\pm$0.7 |
| walker2d-medium-expert-v2 | 108.1$\pm$0.2 | 107.8$\pm$0.5 | 111.8$\pm$0.6 | 110.5$\pm$0.3 |

Table 2.Results of SEABO upon other offline RL baselines across 5 runs. We bold the cell if SEABO outperforms the baseline.

[2] Chen, L., Lu, K., Rajeswaran, A., Lee, K., etc. Decision transformer: Reinforcement learning via sequence modeling. NeurIPS 2021.

[3] Fujimoto, S., Chang, W. D., Smith, E. etc. For SALE: State-Action Representation Learning for Deep Reinforcement Learning. NeurIPS 2023.

Reviewer comment (Theory):
Can SEABO utilize alternative distance metrics in place of the Euclidean distance? If so, how much modification is required?

Author response:
It is interesting to check how SEABO behaves with other distance measurements. We choose Manhattan distance ($D(x,y) = \sum\_{i}|x\_i - y\_i|$) and cosine distance ($D(x,y) = \dfrac{x\cdot y}{\||x\|| \||y\||}$) for evaluation and summarize the results below. Note that it is quite easy to replace Euclidean distance with other distance measurements (the modification is minor as we only need to replace the function that calculates Euclidean distance with another distance measurement function). Note that we keep hyperparameters of SEABO and IQL unchanged on the evaluated datasets across different distance measurements.

| Task Name | IQL (orcale) | +SEABO(Euclidean) | +SEABO(Manhattan) | +SEABO(cosine) |
| ---- | :---: | :---: | :---: | :---: |
| halfcheetah-medium-v2 | **47.4$\pm$0.2** | 44.8$\pm$0.3 | 44.5$\pm$0.1 | 42.6$\pm$0.1 |
| hopper-medium-v2 | 66.2$\pm$5.7 | **80.9$\pm$3.2** | 74.9$\pm$2.7 | 80.9$\pm$1.1 |
| walker2d-medium-v2 | 78.3$\pm$8.7 | 80.9$\pm$0.6 | **81.1$\pm$0.7** | 77.3$\pm$0.7 |
| halfcheetah-medium-replay-v2 | **44.2$\pm$1.2** | 42.3$\pm$0.1 | 42.1$\pm$0.8 | 36.6$\pm$1.3 |
| hopper-medium-replay-v2 | 94.7$\pm$8.6 | 92.7$\pm$2.9 | **96.7$\pm$3.4** | 69.3$\pm$4.8 |
| walker2d-medium-replay-v2 | 73.8$\pm$7.1 | **74.0$\pm$2.7** | 61.6$\pm$22.0 | 63.0$\pm$8.4 |
| halfcheetah-medium-expert-v2 | 86.7$\pm$5.3 | 89.3$\pm$2.5 | **92.0$\pm$0.4** | 71.9$\pm$6.0 |
| hopper-medium-expert-v2 | 91.5$\pm$14.3 | 97.5$\pm$5.8 | **98.0$\pm$10.7** | 92.7$\pm$10.1 |
| walker2d-medium-expert-v2 | 109.6$\pm$1.0 | **110.9$\pm$0.2** | 109.1$\pm$0.5 | 105.7$\pm$5.2 |
| pen-human-v0 | 70.7$\pm$8.6 | 94.3$\pm$12.0 | **96.6$\pm$5.4** | 96.0$\pm$13.8 |
| pen-cloned-v0 | 37.2$\pm$7.3 | 48.7$\pm$15.3 | 52.3$\pm$20.7 | **52.6$\pm$17.7** |
| door-human-v0 | 3.3$\pm$1.3 | **5.1$\pm$2.0** | 5.0$\pm$3.5 | 2.7$\pm$0.9 |
| door-cloned-v0 | **1.6$\pm$0.5** | 0.4$\pm$0.8 | 0.0$\pm$0.0 | 0.0$\pm$0.1 |

Table 1. Comparison of SEABO under different distance measurements. The results are averaged over 5 different random seeds.

It can be seen that Manhattan distance can also incur good performance on some datasets, while cosine distance seems to have unsatisfying performance on numerous datasets. The optimal distance measurement for different dataset may vary, while simply using Euclidean distance can already ensure a good performance.

Reviewer comment (Experiments):
The comparison between ground-truth rewards and the rewards obtained by SEABO (as in Figures 2 and 12) is intriguing. Have you also conducted similar reward comparisons in additional environments, such as AntMaze-v0 and Adroit-v0?

Author response:
**Concern 4: reward comparisons in additional environments, such as AntMaze-v0 and Adroit-v0**

We include additional reward comparison between ground-truth rewards and the rewards acquired by using SEABO on hammer-human-v0, hammer-cloned-v0, door-human-v0, door-cloned-v0, and antmaze-umaze-v0, antmaze-medium-diverse-v0. Please see Appendix F in our updated manuscript for details. We find that the reward plots of SEABO resemble those of the vanilla reward signals.

Hopefully, these can resolve the concerns. If there is still something unclear, please let us know!