Efficient Multi-Task Reinforcement Learning via Task-Specific Action Correction*

We extend the single-task Markov Decision Process (MDP) problem to multi-task MDP for a agent under the framework of Contextual MDP (CMDP)[28]. CMDP is defined by a tuple $\left\langle\mathcal{C},\mathcal{S},\mathcal{A},\gamma,\mathcal{M}\right\rangle$. Here, $\mathcal{C}$ can be viewd as a task set $\mathcal{C}=\left\{\mathcal{T}_{i}\right\}_{i=1}^{N}$, where $\mathcal{T}_{i}$ denotes the task $i$ and $N$ is the number of tasks. $\mathcal{S}$ represents the shared state space , $\mathcal{A}$ denotes the shared action space, and $\gamma$ is the discount factor. State $s\in\mathcal{S}$ and action $a\in\mathcal{A}$. $\mathcal{M}$ is a fuction that maps a task $\mathcal{T}_{i}\in\mathcal{C}$ to MDP parameters, such that $\mathcal{M}(\mathcal{T}_{i})=\left\{P_{i},R_{i},\rho_{i}\right\}$. The transition probability $P_{i}$, reward function $R_{i}$ and initial state distribution $\rho_{i}$ vary by each task. During training, tasks are sampled from a uniform distribution $p(\mathcal{T})$. The agent’s policy takes state $s$ as input and outputs action $a$. The objective of the agent’s policy $\pi$ is to maximize the expected return $\mathbb{E}_{\mathcal{T}_{i}\sim p(\mathcal{T})}[\mathbb{E}_{\pi}[\sum_{t}\gamma^{t}R_{i}(s_{t},a_{t})]]$, where $s_{t}$ and $a_{t}$ represent the state and the action at timestep $t$.

In this paper, we adopt Soft Actor-Critic (SAC) [26] as the fundamental policy. As observed in [27], SAC, being an off-policy actor-critic algorithm with a strong exploration ability based on maximum entropy, exhibits superior performance.

The concept of SAC goes beyond merely maximizing cumulative rewards. It also introduces additional stochasticity to the policy. Thus, a regularization term that incorporates entropy is included in the reinforcement learning objective. The correction term added with entropy can be defined as follows:

In this context, $\alpha$ represents a regularization coefficient that controls the significance of the entropy term. $H$ denotes the entropy.

As illustrated in Fig. 2, TSAC consists of a pair of cooperative policies. The first policy, called the shared policy (SP) and denoted as $\operatorname{\pi}_{\phi}$, optimizes the guiding dense rewards by proposing a preliminary action $\hat{a}\sim\pi_{\phi}(\cdot|s)$. However, this preliminary action may be shortsighted. The second policy, referred to as the action correction policy (ACP) and denoted as $\pi_{\psi}$, corrects the preliminary action by providing an action correction $\Delta a\sim\pi_{\psi}(\cdot|s,\hat{a})$ to execute an effective action. The result action, denoted as $a=h(\hat{a},\Delta a)$, is then output to the environment, where $h$ represents an editing function. ACP is conditioned on SP’s output $\hat{a}$. Together, these two policies cooperate to improve sample efficiency and performance. For simplicity, we denote the overall composed policy as $\pi_{\psi\circ\phi}(a|s)$.

Motivation: TSAC decomposes policy learning into two subtasks that focus on guiding dense rewards or goal-oriented sparse rewards. This decomposition is motivated by the following considerations:

Manually designed dense rewards incorporate prior knowledge and effectively guide policy learning. However, the magnitude of reward values does not directly indicate the ability of a policy to accomplish tasks. Therefore, we have introduced goal-oriented sparse rewards, which are correlated with the completion of the objective. The goal-oriented sparse rewards of a task $\mathcal{T}_{i}$ is characterized by an "$\epsilon$-region" in state space, represented by:

In this equation, $R_{i}^{s}$ denotes the goal-oriented sparse rewards of task $\mathcal{T}_{i}$, $s$ is the current state, $s_{g}$ denotes the goal state, $f(s,s_{g})$ is a function that maps the goal state and current state to a latent space, computing the distance between them. $\delta_{s_{g}}$ defines the reward value and set $\delta_{s_{g}}=1$. $\epsilon$ represents a small distance threshold.

Following the description of CMDP (section II. A), the initial state distribution $\rho_{i}$ determines the probability density of an episode starting at state $s_{0}$. For each transition $(s_{t},a_{t},s_{t+1})$ from task $\mathcal{T}_{i}$ at timestep $t$, the environment produces a scalar $R_{i}(s_{t},a_{t})$. It is worth mentioning that the reward function $R_{i}$ is artificially designed and intensive, which we refer to as guiding dense rewards. Similarly, the environment produces a scalar $R_{i}^{s}(s_{t},a_{t})$ as a goal-oriented sparse reward. Higher values for both the guiding dense rewards and the goal-oriented sparse rewards indicate better performance.

For each state $s_{t}$, the guiding dense rewards state value of policy $\pi$ is denoted as $V^{\pi}(s_{t})=\mathbb{E}_{\mathcal{T}_{i}\sim p(\mathcal{T})}[\mathbb{E}_{\pi}\sum_{t^{\prime}=t}^{\infty}\gamma^{t^{\prime}-t}R_{i}(s_{t^{\prime}},a_{t^{\prime}})]$. The guiding dense rewards state-action value is denoted as $Q^{\pi}(s_{t},a_{t})=\mathbb{E}_{\mathcal{T}_{i}\sim p(\mathcal{T})}[R_{i}(s_{t},a_{t})+\gamma\mathbb{E}_{s_{t+1}\thicksim P_{i}}V^{\pi}(s_{t+1})]$. Similarly, We define $V_{s}^{\pi}$ and $Q_{s}^{\pi}$ for the goal-oriented sparse rewards.

We consider the MTRL objective from the perspective of guiding dense rewards and goal-oriented sparse rewards:

MTRL can be viewed as a multi-objective optimization problem. The introduction of sparse rewards adds an additional objective to the MTRL framework, which increases the problem’s complexity. To convert multi-objective MTRL into single-objective MTRL, we assign a virtual expected budget $C$ to the sparse rewards. This allows us to write the MTRL objective as follows:

In this equation, $C$ represents a virtual expected budget. Specifically for each time step, Eq.4 can be rewritten as:

Here, $c$ denotes the expected budget specific to each step and it relates to $C$ through the equation $\sum_{t=0}^{\infty}\gamma^{t}c=\frac{c}{1-\gamma}=C$. Notably, the expected budget represents an average target to be achieved rather than a strict enforcement.

To simplify the problem, we transform the multi-objective MTRL into a single-objective MTRL using the Lagrangian method. This method converts the constrained optimization problem Eq.(4) into an unconstrained one by introducing a multiplier $\lambda$:

In this formulation, the weight of the goal-oriented sparse rewards combined with the guiding dense rewards is represented by $\lambda$. We solve this objective by using model-free MTRL algorithms.

We utilize an off-policy actor-critic approach to train the two policies. Given the overall policy $\pi_{\psi\circ\phi}(a|s)$, we use the typical Temporal Difference (TD)[29] backup to learn $Q^{\pi_{\psi\circ\phi}}(s,a;\theta)$ and $Q_{s}^{\pi_{\psi\circ\phi}}(s,a;\theta_{s})$, which are parameterized as $Q(s,a;\theta)$ and $Q_{s}(s,a;\theta_{s})$ respectively. Here, $\theta$ represents the network parameters collectively for the two state-action values. When provided with $s_{t+1}$ and $a_{t+1}$, the Bellman backup operator for the guiding dense rewards state-action value is expressed as:

where $R_{i}$ represents the guiding dense rewards obtained from task $\mathcal{T}_{i}$. The backup operator for the goal-oriented sparse rewards state-action value $Q_{s}$ is defined similarly with $R^{s}$. Both $Q(s,a;\theta)$ and $Q_{s}(s,a;\theta_{s})$ can be learned from transitions $(s_{t},a_{t},s_{t+1})$ sampled from a replay buffer.

To achieve off-policy training of $\psi$ and $\phi$, we convert Eq.(6) into a bi-level optimization surrogate. The formulation is presented as follows:

Here, $\mathcal{D}$ represents a replay buffer containing a historical marginal state distribution used to train the policies. The initial state distribution $\rho$ is employed to train $\lambda$. This distinction arises from the idea that when fine-tuning $\lambda$, we should primarily consider how well the policy satisfies the virtual expected budget starting with $\rho$, rather than with some historical state distribution. Subsequently, We further transform the off-policy objective(Eq.(8),a) into two parts:

In this modified formulation, the distance function $d(a,\hat{a})$ quantifies the change from $\hat{a}$ to $a$. It is not necessarily proportional to $\Delta a$ as the editing function $h$ can introduce nonlinearity. ACP $\pi_{\psi}$ aims to maximize the goal-oriented sparse rewards while minimizing the distance between the actions before and after the correction. On the other hand, SP $\pi_{\phi}$ solely focus on maximizing the guiding dense rewards. Importantly, ACP modifies SP’s action, which aligns with the discussed motivation for efficient exploration. The training objective (Eq.(9),b) for ACP relies on a critic $Q_{s}$, which learns the expected future goal-oriented sparse rewards. Consequently, guided by $Q_{s}$, ACP explores actions with greater potential in long-term sequences.

In this section, we present the design for the action correction function $h(\hat{a},\Delta a)$ and the distance function $d(a,\hat{a})$.

To practically train the objectives, we employ stochastic gradient descent (SGD) simultaneously to Eq.(8)(b) and Eq.(9). The re-parameterization trick is utilized for both $\pi_{\psi}$ and $\pi_{\phi}$ to enable the application of SGD. To assess $\Lambda_{\pi_{\psi\circ\phi}}$, a batch of rollout experiences $\{(s_{n},a_{n})\}_{n=1}^{N}$ following $\pi_{\psi\circ\phi}$ is provided. The gradient of $\lambda$ (Eq.(8)(b)) is approximated as:

Here, $c$ represents the virtual expected budget as defined in Eq.(5). Subsequent to each rollout, a batch of goal-oriented sparse rewards is collected, each reward is compared to $-c$, and the mean of the differences is used to adjust $\lambda$. This approximation enables the updating of $\lambda$ using mini-batches of data, irather than waiting for complete episodes to conclude or relying on the often inaccurate estimated $V_{s}^{\pi\psi\circ\phi}$. Multiple parallel environments are employed to mitigate temporal correlation within the rollout batch data, which constitutes the data to be placed into the replay buffer.

The computational graph presenting Eq.(9) is illustrated in Fig. 3. Our approach is applicable to various goal-oriented sparse rewards and action distance functions. To encourage exploration, we integrate SAC[26], which incorporates the entropy terms of $\pi_{\psi}$ and $\pi_{\phi}$ in Eq.(9), dynamically adjusting their weights based on two entropy targets as described in [26].In our experiments, both SP and ACP are trained from scratch. The pseudocode for the overall TSAC is shown in the algorithm 1.

We will compare our method against the following baselines:

CARE: As a representations sharing method, representations are shared to learn how to compose them by leveraging additional information about each task.

MT–SAC: This approach directly applies the SAC algorithm to the multi-task setting. It utilizes a shared backbone with disentangled alphas.

Soft Modularization: As a parameters sharing method, Soft Modularization shares parameters and uses a routing network to softly combine all possible routes for each task.

PCGrad: It is a gradient manipulation method that projects a task’s gradient onto the normal plane of any other conflicting task. However, it has high time complexity and is not suitable for MT50 in the long horizon.

TSAC(ours): Our proposed method builds upon CARE, leveraging its representation-sharing module. In contrast, our approach is based on behavior sharing and utilizes goal-oriented sparse rewards.

Fig. 5a shows the average success rate on the 10 tasks of the MT10 benchmark form Meta-world for TSAC, CARE, MT–SAC, Soft Modularization, and PCGrad. Since the success rate is a binary variable, it is noisy; therefore, the results were averaged across multiple seeds and the curves were smoothed. Mean and standard error are reported for each value.

We consider 1 million steps as a long horizon, and 150 thousand steps as a short horizon. The short horizon is utilized to observe the exploration ability of different methods, while the long horizon is used to visualize the performance of the method at various time points. It is worth noting that all methods are trained using SAC with disentangled alphas.

Table LABEL:table1 and Fig. 5a demonstrate that our method outperforms all the baselines on the short horizon in MT10. Additionally, Fig. 5c illustrates that even in MT50 our method still outperform all baselines. For comparison, it takes around 1 million steps for the Multi-task SAC agent to reach the accuracy that our TSAC agent achieves around 300 thousand steps, suggesting that our method is highly sample-efficient and exploration-efficient.

Table LABEL:table2 and LABEL:table3 as well as Fig. 5b and Fig. 5d depict that TSAC is able to learn a good policy on the long horizon. For MT10, TSAC performs best and reaches a top success rate of 0.827 during the training. Furthermore, sampling the success rate around 0.8 million and 1 million steps reveals that TSAC has the best performance. For MT50, conflicts between tasks become more acute. CARE stops learning around 600 thousand steps and performance begins to decline in the following training steps. Despite suffering from the conflicts between tasks, TSAC’s performance declines, but it still performs better than CARE. MT-SAC achieves smooth learning and has similar performance to that of CARE and TSAC at 0.8 million and 1 million steps. However, in terms of the best performance throughout training, MT–SAC is far inferior to TSAC, which outperforms all methods. Since the success metric is a binary variable and very noisy, the best performance is obtained by smoothing over the curve. Importantly, TSAC achieves average success rates of 0.450 and 0.445 on MT50 at 0.8 million and 1 million steps, surpassing the reported results from CARE, MT–SAC and Soft Modularization.

The result in the previous section suggests that TSAC is both sample efficient and exploration efficient. Furthermore, TSAC yields a good policy on long horizons. In this section, we further investigate the impact of the various action correction functions employed by TSAC and discuss the potential for improved performance.

In ablation study, We will consider four different functions $h$:

SP dominated (ours): $h=\operatorname*{min}(\operatorname*{max}(2\hat{a}+\Delta a,-A),A)$, SP’s action dominates the final action.

ACP dominated: $h=\operatorname*{min}(\operatorname*{max}(\hat{a}+2\Delta a,-A),A)$, ACP’s action dominates the final action.

equal contribution: $h=\operatorname*{min}(\operatorname*{max}(\hat{a}+\Delta a,-A),A)$, SP and ACP contribute equally to the final action.

Softclip: $h=Softclip(2\hat{a}+\Delta a)$, $h$ use softclip to smooth out the output action and bring in nonlinearity.

Fig. 6 illustrates that SP dominated action correction function outperforms other functions in producing the final action. However, when SP and ACP contribute equally to the final action, the agent encounters diffculties in learning a policy. This conflict relationship between SP and ACP presents challenges in optimizing each respective goal. Conversely, when either SP or ACP dominates the action correction funcion, the agent can successfully learn a well-performing policy. Comparatively, SP outperforms ACP due to ACP’s focus on optimizing the goal-oriented sparse rewards, which is challenging to train due to its sparsity. In contrast, the Softclip function exhibits a slight decrease in performance compared to our action correction function. This decline can be attributed to the introduction of nonlinearity, which further complicates and hampers the learning process. Consequently, we opt for a simple and linear operation instead of training a learnable network.