From ab19b8ca20b3a959b30364bfa77d7c94132e6241 Mon Sep 17 00:00:00 2001
From: HydrogenSulfate <490868991@qq.com>
Date: Tue, 12 Dec 2023 02:09:01 +0800
Subject: [PATCH 1/5] add MOO algorithm: ReloBRalo

---
 ppsci/loss/mtl/__init__.py  |  4 +-
 ppsci/loss/mtl/relobralo.py | 94 +++++++++++++++++++++++++++++++++++++
 2 files changed, 97 insertions(+), 1 deletion(-)
 create mode 100644 ppsci/loss/mtl/relobralo.py

diff --git a/ppsci/loss/mtl/__init__.py b/ppsci/loss/mtl/__init__.py
index 8b5aa0d5b..83e185a49 100644
--- a/ppsci/loss/mtl/__init__.py
+++ b/ppsci/loss/mtl/__init__.py
@@ -17,11 +17,13 @@
 from ppsci.loss.mtl.agda import AGDA
 from ppsci.loss.mtl.base import LossAggregator
 from ppsci.loss.mtl.pcgrad import PCGrad
+from ppsci.loss.mtl.relobralo import Relobralo
 
 __all__ = [
+    "AGDA",
     "LossAggregator",
     "PCGrad",
-    "AGDA",
+    "Relobralo",
 ]
 
 
diff --git a/ppsci/loss/mtl/relobralo.py b/ppsci/loss/mtl/relobralo.py
new file mode 100644
index 000000000..14d8dde54
--- /dev/null
+++ b/ppsci/loss/mtl/relobralo.py
@@ -0,0 +1,94 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from typing import List
+
+import paddle
+from paddle import nn
+
+from ppsci.loss.mtl import base
+
+
+class Relobralo(nn.Layer, base.LossAggregator):
+    """Base class of loss aggregator mainly for multitask learning.
+
+    Args:
+        model (nn.Layer): Training model.
+    """
+
+    def __init__(
+        self,
+        model: nn.Layer,
+        num_losses: int,
+        alpha: float = 0.95,
+        beta: float = 0.99,
+        tau: float = 1.0,
+        eps: float = 1e-8,
+    ) -> None:
+        super().__init__(model)
+        self.num_losses: int = num_losses
+        self.alpha: float = alpha
+        self.beta: float = beta
+        self.tau: float = tau
+        self.eps: float = eps
+        self.register_buffer("losses_init", paddle.zeros([self.num_losses]))
+        self.register_buffer("losses_prev", paddle.zeros([self.num_losses]))
+        self.register_buffer("lmbda", paddle.ones([self.num_losses]))
+
+    def _softmax(self, vec: paddle.Tensor) -> paddle.Tensor:
+        max_item = vec.max()
+        result = paddle.exp(vec - max_item) / paddle.exp(vec - max_item).sum()
+        return result
+
+    def _compute_bal(
+        self, losses_vec1: paddle.Tensor, losses_vec2: paddle.Tensor
+    ) -> paddle.Tensor:
+        return self.num_losses * (
+            self._softmax(losses_vec1 / (self.tau * losses_vec2 + self.eps))
+        )
+
+    def __call__(self, losses: List[paddle.Tensor], step: int = 0) -> "Relobralo":
+        self.step = step
+        losses_stacked = paddle.stack(losses)  # [num_losses, ]
+
+        if self.step == 0:
+            self.loss = losses_stacked.sum()
+            paddle.assign(losses_stacked.detach(), self.losses_init)
+        else:
+            with paddle.no_grad():
+                # 1. update lambda_hist
+                rho = paddle.bernoulli(paddle.tensor(self.beta))
+                lmbda_hist = rho * self.lmbda + (1 - rho) * self._compute_bal(
+                    losses_stacked, self.losses_init
+                )
+
+                # 2. update lambda
+                paddle.assign(
+                    self.alpha * lmbda_hist
+                    + (1 - self.alpha)
+                    * self._compute_bal(losses_stacked, self.losses_prev),
+                    self.lmbda,
+                )
+
+                # 3. compute reweighted total loss with lambda
+                self.loss = losses_stacked * self.lmbda
+
+        # update losses_prev at the end of each step
+        paddle.assign(losses_stacked.detach(), self.losses_prev)
+        return self
+
+    def backward(self) -> None:
+        self.loss.backward()

From b8e5691139f987c31f735e448e5df2dc457544d2 Mon Sep 17 00:00:00 2001
From: HydrogenSulfate <490868991@qq.com>
Date: Tue, 12 Dec 2023 09:05:00 +0800
Subject: [PATCH 2/5] update code

---
 ppsci/loss/mtl/relobralo.py | 24 +++++++++++++++---------
 1 file changed, 15 insertions(+), 9 deletions(-)

diff --git a/ppsci/loss/mtl/relobralo.py b/ppsci/loss/mtl/relobralo.py
index 14d8dde54..d52b92234 100644
--- a/ppsci/loss/mtl/relobralo.py
+++ b/ppsci/loss/mtl/relobralo.py
@@ -19,10 +19,8 @@
 import paddle
 from paddle import nn
 
-from ppsci.loss.mtl import base
 
-
-class Relobralo(nn.Layer, base.LossAggregator):
+class Relobralo(nn.Layer):
     """Base class of loss aggregator mainly for multitask learning.
 
     Args:
@@ -38,7 +36,13 @@ def __init__(
         tau: float = 1.0,
         eps: float = 1e-8,
     ) -> None:
-        super().__init__(model)
+        super().__init__()
+        self.model = model
+        self.step = 0
+        self.param_num = 0
+        for param in self.model.parameters():
+            if not param.stop_gradient:
+                self.param_num += 1
         self.num_losses: int = num_losses
         self.alpha: float = alpha
         self.beta: float = beta
@@ -66,11 +70,12 @@ def __call__(self, losses: List[paddle.Tensor], step: int = 0) -> "Relobralo":
 
         if self.step == 0:
             self.loss = losses_stacked.sum()
-            paddle.assign(losses_stacked.detach(), self.losses_init)
+            with paddle.no_grad():
+                paddle.assign(losses_stacked.detach(), self.losses_init)
         else:
             with paddle.no_grad():
                 # 1. update lambda_hist
-                rho = paddle.bernoulli(paddle.tensor(self.beta))
+                rho = paddle.bernoulli(paddle.to_tensor(self.beta))
                 lmbda_hist = rho * self.lmbda + (1 - rho) * self._compute_bal(
                     losses_stacked, self.losses_init
                 )
@@ -83,11 +88,12 @@ def __call__(self, losses: List[paddle.Tensor], step: int = 0) -> "Relobralo":
                     self.lmbda,
                 )
 
-                # 3. compute reweighted total loss with lambda
-                self.loss = losses_stacked * self.lmbda
+            # 3. compute reweighted total loss with lambda
+            self.loss = (losses_stacked * self.lmbda).sum()
 
         # update losses_prev at the end of each step
-        paddle.assign(losses_stacked.detach(), self.losses_prev)
+        with paddle.no_grad():
+            paddle.assign(losses_stacked.detach(), self.losses_prev)
         return self
 
     def backward(self) -> None:

From 6f2f841631a32fd41cf7164b5ef82d8742aedc09 Mon Sep 17 00:00:00 2001
From: HydrogenSulfate <490868991@qq.com>
Date: Tue, 12 Dec 2023 02:49:31 +0000
Subject: [PATCH 3/5] add final ReloBRalo code

---
 docs/zh/api/loss.md         |  3 ++-
 ppsci/loss/mtl/relobralo.py | 33 +++++++++++++++++++++++++--------
 2 files changed, 27 insertions(+), 9 deletions(-)

diff --git a/docs/zh/api/loss.md b/docs/zh/api/loss.md
index ccdabceee..a2b59e24b 100644
--- a/docs/zh/api/loss.md
+++ b/docs/zh/api/loss.md
@@ -23,8 +23,9 @@
     handler: python
     options:
       members:
+        - AGDA
         - LossAggregator
         - PCGrad
-        - AGDA
+        - Relobralo
       show_root_heading: true
       heading_level: 3
diff --git a/ppsci/loss/mtl/relobralo.py b/ppsci/loss/mtl/relobralo.py
index d52b92234..4bbc73e2d 100644
--- a/ppsci/loss/mtl/relobralo.py
+++ b/ppsci/loss/mtl/relobralo.py
@@ -21,15 +21,37 @@
 
 
 class Relobralo(nn.Layer):
-    """Base class of loss aggregator mainly for multitask learning.
+    r"""
+    **Re**lative **Lo**ss **B**alancing with **Ra**ndom **Lo**okback
+
+    [Multi-Objective Loss Balancing for Physics-Informed Deep Learning](https://arxiv.org/abs/2110.09813)
 
     Args:
-        model (nn.Layer): Training model.
+        num_losses (int): Number of losses.
+        alpha (float, optional): Ability for remembering past in paper. Defaults to 0.95.
+        beta (float, optional): Parameter for generating $\rho$ from bernoulli distribution,
+            and $E[\rho](=\beta)$ should be close to 1. Defaults to 0.99.
+        tau (float, optional): Temperature factor. Equivalent to softmax when $\tau$=1.0,
+            equivalent to argmax when $\tau$=0. Defaults to 1.0.
+        eps (float, optional): $\epsilon$ to avoid divided by 0 in losses. Defaults to 1e-8.
+
+    Examples:
+        >>> import paddle
+        >>> from ppsci.loss import mtl
+        >>> model = paddle.nn.Linear(3, 4)
+        >>> loss_aggregator = mtl.Relobralo(2)
+        >>> for i in range(5):
+        ...     x1 = paddle.randn([8, 3])
+        ...     x2 = paddle.randn([8, 3])
+        ...     y1 = model(x1)
+        ...     y2 = model(x2)
+        ...     loss1 = paddle.sum(y1)
+        ...     loss2 = paddle.sum((y2 - 2) ** 2)
+        ...     loss_aggregator([loss1, loss2]).backward()
     """
 
     def __init__(
         self,
-        model: nn.Layer,
         num_losses: int,
         alpha: float = 0.95,
         beta: float = 0.99,
@@ -37,12 +59,7 @@ def __init__(
         eps: float = 1e-8,
     ) -> None:
         super().__init__()
-        self.model = model
         self.step = 0
-        self.param_num = 0
-        for param in self.model.parameters():
-            if not param.stop_gradient:
-                self.param_num += 1
         self.num_losses: int = num_losses
         self.alpha: float = alpha
         self.beta: float = beta

From acd2c9fdead4d65f329ed97e5904cb42ed159369 Mon Sep 17 00:00:00 2001
From: HydrogenSulfate <490868991@qq.com>
Date: Tue, 12 Dec 2023 09:20:58 +0000
Subject: [PATCH 4/5] update document

---
 docs/index.md | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/docs/index.md b/docs/index.md
index e825c16e3..8577bef00 100644
--- a/docs/index.md
+++ b/docs/index.md
@@ -26,8 +26,9 @@
 | 微分方程 | [洛伦兹方程](./zh/examples/lorenz.md) | 数据驱动 | Transformer-Physx | 监督学习 | [Data](https://github.com/zabaras/transformer-physx) | [Paper](https://arxiv.org/abs/2010.03957) |
 | 微分方程 | [若斯叻方程](./zh/examples/rossler.md) | 数据驱动 | Transformer-Physx | 监督学习 | [Data](https://github.com/zabaras/transformer-physx) | [Paper](https://arxiv.org/abs/2010.03957) |
 | 算子学习 | [DeepONet](./zh/examples/deeponet.md) | 数据驱动 | MLP | 监督学习 | [Data](https://deepxde.readthedocs.io/en/latest/demos/operator/antiderivative_unaligned.html) | [Paper](https://export.arxiv.org/pdf/1910.03193.pdf) |
-| 微分方程 | 梯度增强的物理知识融合PDE求解<sup>coming soon</sup> | 机理驱动 | gPINN | 半监督学习 | - |  [Paper](https://www.sciencedirect.com/science/article/abs/pii/S0045782522001438?via%3Dihub) |
+| 微分方程 | [梯度增强的物理知识融合 PDE 求解](https://github.com/PaddlePaddle/PaddleScience/blob/develop/examples/gpinn/poisson_1d.py) | 机理驱动 | gPINN | 无监督学习 | - |  [Paper](https://doi.org/10.1016/j.cma.2022.114823) |
 | 积分方程 | [沃尔泰拉积分方程](./zh/examples/volterra_ide.md) | 机理驱动 | MLP | 无监督学习 | - | [Project](https://github.com/lululxvi/deepxde/blob/master/examples/pinn_forward/Volterra_IDE.py) |
+| 微分方程 | [分数阶微分方程](https://github.com/PaddlePaddle/PaddleScience/blob/develop/examples/fpde/fractional_poisson_2d.py) | 机理驱动 | MLP | 无监督学习 | - | - |
 
 <br>
 <p align="center"><b>技术科学(AI for Technology)</b></p>

From f741cd6910383325ced2d420d0e87b8d7d640e27 Mon Sep 17 00:00:00 2001
From: HydrogenSulfate <490868991@qq.com>
Date: Tue, 12 Dec 2023 09:24:17 +0000
Subject: [PATCH 5/5] update css

---
 docs/stylesheets/extra.css | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/docs/stylesheets/extra.css b/docs/stylesheets/extra.css
index ebc042903..97d34cc25 100644
--- a/docs/stylesheets/extra.css
+++ b/docs/stylesheets/extra.css
@@ -7,7 +7,7 @@
 
 .md-grid {
   /* readable page width */
-  max-width: 70%;
+  max-width: 1440px;
 }
 
 .md-header__topic > .md-ellipsis {