Rethinking Optimal Transport in Offline Reinforcement Learning

**Q1. If I understood the paper correctly, you are viewing as a transport map from the state distribution d(s) to the behavior policy $\beta(\cdot|s')$ for some s'. The transport map will depend on the choice of s'. How do you choose this s'?** Sorry for the misunderstanding. We are randomly sampling $s$ from the given dataset $\mathcal{D}$, please see derivations at lines 131 and Eq.9. For the sampled $s$ we choose the conditional distribution of the action space provided by the dataset (expert) $\beta(\cdot|s)$ as the target. We will make this explicit in the final version. **Q2. Lines 190 and 191 state that the agent has to "navigate through 50 intermediate steps." Doesn't this mean that the episode length is fixed to 51 regardless of the agent's path?** Yes, the length of the episode is fixed. This length of 50, corresponds to the number of steps provided by each expert policy in the dataset. **Q5. The value learning part of IQL can be completely separated from the policy learning part, which uses AWR, as you mentioned. IQL Q function can be trained without weighted behaviour cloning and thus can be used together with the proposed algorithm.** Thanks for the clarification, this experiment is really interesting. Following your question, we conducted a series of experiments on MuJoCo using the CORL[1] IQL implementation. In these tests, we found that the value function trained using expectile regression from the IQL method is inappropriate for off-policy evaluation. - First, we conducted an experiment *without* using our method. We trained a $Q$-function using expectile regression loss and attempted to extract a policy through direct optimization: $\min_{\pi}\mathbb{E}_{s \sim \mathcal{D}, a\sim\pi(s)} \big[-Q^{\text{IQL}}(s, a)\big]$. This approach resulted in zero rewards. - Second, we added a BC objective to avoid distribution shift: $\min_{\pi}\mathbb{E}_{s \sim \mathcal{D}, a\sim\pi(s)} \big[-Q^{\text{IQL}}(s, a)\big]+(a-\beta(s))^2$. The results were the same — $Q$-functions trained via expectile regression dramatically overestimate actions sampled by the learned policy $\pi$. - Finally, we tested our method: $\min_{\pi}\mathbb{E}_{s \sim \mathcal{D}, a\sim\pi(s)} \big[-Q^{\text{IQL}}(s, a)-f(s,a)\big]$ and observed improvements in the scores! Even with the ill-suited cost function, policy optimization with respect to the potential $-f(s,a)$ yielded the highest scores. We also tested the advantage $A(s,a) $and exponential advantage functions exp$A(s,a)$ from IQL, but did not observe any improvements. The table below summarizes our analyses. --- **Table: Averaged normalized scores on MuJoCo tasks. Reported scores are the results of the final 10 evaluations and 3 random seeds.** | Dataset | $-Q^{IQL}(s, a)$ | $-Q^{IQL}(s, a)$ + BC | $-Q^{IQL}(s, a) - f(s,a)$ (Ours) | |-------------------------|------------------|-----------------------|---------------------------------| | HalfCheetah-medium | -2.53 ± 0.1 | -2.54 ± 0.1 | 48.7 ± 0.3 | | HalfCheetah-medium-expert | -2.53 ± 0.2 | -2.54 ± 0.2 | 39.9 ± 1.4 | | Hopper-medium | 0.6 ± 0.1 | 0.6 ± 0.1 | 27.9 ± 15.4 | | Hopper-medium-expert | 0.7 ± 0.1 | 0.6 ± 0.1 | 8.7 ± 4.6 | | Walker-medium | -0.16 ± 0.1 | -0.16 ± 0.0 | 37.6 ± 3.6 | | Walker-medium-expert | -0.23 ± 0.2 | -0.21 ± 0.1 | 17.5 ± 14.3 | --- While the IQL method can be formally decoupled, we see that both components, expectile-based in-sample value learning and weighted behaviour cloning, are important parts of each other to achieve strong results. Adapting IQL value learning for better off-policy evaluation is a promising direction, but beyond the scope of our contribution. We believe that our strong performance on most tasks, combined with different types of a cost function, justifies our formulation and makes it valuable to the community. We will include these analyses in the final version of the paper. **Reference:** [1] JAX-CORL: https://github.com/nissymori/JAX-CORL

**Q1.** Say we are going to use the log-likelihood as our regularizer, which means the objective function would be something like $\mathbb{E}_{s\sim \mathcal{D}}\left[Q(s, \pi(s))-\alpha\log\beta(\pi(s)\mid s)\right]$. For each $s\in S$, the regularizer $\log\beta(\pi(s)\mid s)$ is only affected by the value of $\pi(s)$ and nothing else. If we ignore the expressivity of a neural network, the resulting $\pi$ is straightforward to interpret: $\pi(s)=\text{arg}\max_a Q(s, a)-\alpha\log\beta(a\mid s)$. However, in the case of OT, the value of $T_s(s')$ for all $s'\in S$ matters, where $T_s$ is the transport map between the state distribution $d(\cdot)$ and $\beta(\cdot\mid s)$. I'm concerned that trying to force the same $T$ for all $\beta(\cdot\mid s)$ would cause conflicts between different $s$. **Q2.** Does it mean that an agent's actions in the intermediate 49 steps wouldn't affect the return?

**Q1.** I'm confused now. Let's consider the simplest case where $\mathcal{D}$ has two elements $(s_1, a_1)$ and $(s_2, a_2)$. Then the state distribution $d(s)$ is $\frac{1}{2}\delta(s=s_1)+\frac{1}{2}\delta(s=s_2)$ and the behavior policies are $\beta(a\mid s_1)=\delta(a=a_1)$ and $\beta(a\mid s_2)=\delta(a=a_2)$. The optimal transport map $T_{s_1}$ from $d$ to $\beta(\cdot\mid s_1)$ should be $T_{s_1}\equiv a_1$ and the map $T_{s_2}$ from $d$ to $\beta(\cdot\mid s_2)$ should be $T_{s_2}\equiv a_2$. If you try to use the same $T$ in both cases, conflict would occur since $a_1\neq a_2$, wouldn't it? Am I missing something? **Q2.** What is the transition function? From the description in the paper, it looks like there are states $s_1, s_2, \cdot, s_T$ uniformly spaced across the x-axis and the agent will transition from $s_i$ to $s_{i+1}$ regardless of the action it takes. After 50 steps, the agent will arrive at state $s_T$ no matter what and get a reward of 1. What would encourage the agents to follow the shortest path?