How to extend SACPID to multi-constraint situation?

[TomatoZ2]

Thank you for your contribution, it's really helpful for a newcomer!

How to extend SACPID to multi-constraint situation? Although I've noticed similar discussions[#258], I still haven't figured out how to do it. I think more cost_critics should be introduced to estimate the cost, but I don't know how to deal with these λs.

[Gaiejj]

Thanks for your recognition, and it is a good question. Suppose you have multiple critics, then the formulation of CMDPs with multiple costs is:
$$\text{max } J^R$$ $$\text{while }J^{C_i}< d_i$$
where $J^R$ is the reward objective and $J^{C_i}$ is the cost objective, $d_i$ is the cost_limit and $i$ denotes the index of each single cost.
Since you use a Lagrangian method, and you have multiple Lagrangian Multipliers $\lambda$ s, suppose you are in the $k_{th}$ learning iteration, the reward advantage and cost advantages are $A^R$ and $A^{C_i}$ , a naive method can be:
$$\text{max} [A^R-\lambda^k_i A^{C_i}]$$ $$\lambda^{k+1}_i = \text{max}(\lambda^k_i+\eta_i*\text{PID}(J^{C_i}-d_i), 0)$$
Specifically, in off-policy algorithm SACPID, we have $A^R=Q^R$ and $A^{C_i}=Q^{C_i}$, where $Q$ denotes the Critic estimation. Besides, in OmniSafe implementation, the $J^{C_i}$ is always the episodic costs (without discounted) obtained from Logger.

[TomatoZ2]

Thanks for your answer! I want to clarify some details.
According to this paper Responsive Safety in Reinforcement Learning by PID Lagrangian Methods and my recent study of the Lagrangian relaxation method, I think the entire solution process should be as follows.
For the following multi-constraints optimization problem, there are k constraints.
$$\min _{\mathbf{x}} f(\mathbf{x}) \quad \text { s.t. } g_i(\mathbf{x})=0 \quad i=1,2,3..k$$
The extended optimazition objective should be like this
$$\mathcal{L}(\mathbf{x},\mathbf{ \lambda})=f(\mathbf{x})+\mathbf{\lambda} \mathbf{G}(\mathbf{x})$$ $$\mathbf{\lambda}=[\lambda_1,\lambda_2,...,\lambda_k]$$ $$\mathbf{G}(\mathbf{x})=[g_1(\mathbf{x}),g_2(\mathbf{x}),...,g_k(\mathbf{x})]^\intercal$$
And iterative solution process
$$-\nabla_x\mathcal{L}(\mathbf{x},\mathbf{\lambda})=-\nabla_xf(x)-\mathbf{\lambda}\nabla_x\mathbf{G}(\mathbf{x})$$ $$\nabla_i \mathcal{L}(\mathbf{x}, \lambda_i)=g_i(\mathbf{x})$$
So the extened optimazition objective should be $\mathcal{L}(\mathbf{x},\mathbf{ \lambda})$ not $\max \left[A^R-\lambda_i^k A^{C_i}\right]$. I don't know why the optimization objective you provided me contains only one $\lambda$, but I think the update of the final value function should be based on $-\nabla_x\mathcal{L}(\mathbf{x},\mathbf{\lambda})$ which contains all $\lambda$.
If there are any errors in my opinions, please point them out in time, or could you please explain why your optimization objective is $\max \left[A^R-\lambda_i^k A^{C_i}\right]$.

[Gaiejj]

My bad. There is something easily misunderstood in my expression. I agree with your formulation, and I think the objective function should be
$$\text{max}[A^R-\sum^k_i{\lambda_i^k A^{C_i}}]$$
This naive method sums all cost advantage functions combined with each $\lambda_i$. As Lagrangian methods are sensitive to hyperparameters, it takes effort to trade off multiple costs and reward. The "learning rate" $\eta_i$ (corresponding to $K_i^p$, $K_i^i$ and $K_i^d$ in PID-Lagrangian methods) and the initial value of each $\lambda_i$ should be carefully designed.

[TomatoZ2]

It's not a a problem right now

I have a problem with programming.
How to expand the data stored in the logger into an array, such as cost, Epcost and Lagrange multiplier. My program has been uploaded, and its requirement are the same as ominisafe. Just run train-policy.py directly, no more peremeters needed. I wrote the new algorithm (multi-constraint SACPID) in new-algo.py, and the relevant parameter settings were written in the ominisafe source code (configs in the omnisafe file). I think the modification of the algorithm should be correct, but I don't know how to modify the logger. The cost returned by safe gymnasium is an array. The cost, Epcost and Lagrange multiplier that required for algorithm update, need to be stored as an array, but logger seems to only support int or float, which makes the algorithm unable to update.

I would be very grateful if you could help me modify the program.

safe_rlap.zip

[TomatoZ2]

Hey, I modified the logger, and the algorithm seems to be experimentally correct, but features like tensorboard don't work. . .
I'm not a professional programmer and my code may be inefficient and have potential logic errors. I would be very happy if you could help me check if there are any errors in the modified code and improve it. And how to make tansorboard work properly?
new_algo.txt

[Gaiejj]

Thank you very much for providing your implementation reference for supporting multiple costs in OmniSafe's Logger for the open-source community. Congratulations on the positive results from your experiments. However, due to significant differences between your code and OmniSafe's structure, we are having difficulty pinpointing the issues quickly. Could you please try incorporating your changes to the Logger directly inomnisafe/common/logger.py? This means modifying the relevant functions within the Logger without changing the interface used to call it. This is because OmniSafe's original Logger already supports TensorBoard. This effort will not only address your issue with using TensorBoard but also contribute to the open-source community. I am more than willing to review your pull request.

[TomatoZ2]

I'm trying to modify the logger to support the tensorboard. But I have a question when I read the code of SACPID. In function _loss_pi, $A^C$ is $Q^C$ not $Q^C-d$, which means you want the constraint to approach $0$ not $d$? Wouldn't this lead to an excessive penalty on the policy gradient? Could you please tell me the reason?

[Gaiejj]

I think there might be some misunderstandings here. In the paper Responsive Safety in Reinforcement Learning by PID Lagrangian Methods you referred to, the author uses the $J_c-d$ shown in the figure below to update the Lagrange multipliers,

and uses the following formula to update the policy $\pi$.

This is the same as the code implementation of OmniSafe.
As for your question about zero cost, imagine if the policy always meets the cost_limit constraint (though not a zero constraint), the Lagrange multiplier will be updated to 0 under PID control, and $A^C$ will no longer have an impact on the update of the policy. (since it is multiplied by a coefficient of 0).

[TomatoZ2]

Got it. But during the experiment, I discovered a phenomenon, which may be affected by the reward function. When the agent obtains more rewards, its output action will make the constraint much smaller than the threshold. When the constraint is already smaller than the threshold, the Lagrange multiplier remains positive. What is the role of the Lagrange multiplier at this time? Shouldn't the multiplier at this time be zero?

[Gaiejj]

Due to the I term of the PID Lagrange multiplier, which records the average of historical constraint violations, the stability of updates is achieved. This implementation helps prevent the Lagrange multiplier from oscillating violently due to individual epoch performances. Therefore, even though the algorithm complies with safety constraints in some epochs, the value of the Lagrange multiplier does not drop immediately.

[TomatoZ2]

I tuned the PID peremeters and the problem solved. Thank you.

[TomatoZ2]

Hi, I successfully modified logger to work with available tensorboard in multi-constraint case. But in fact, the settings in the logger exceeded my expectations, which prevented me from tesing some functions, such as the compatibility of vector with Delta and minmax.

To be honest, I am not a professional programmer but a student, so my initial understanding of the pull request in Github was only code sharing. After I learned about it, I was surprised that your team wanted me to contribute to the modification of the logger. But to be honest, I have never done any code development and I don’t know whether GitHub’s marge request directly merges the submitter’s code or something else. So I strongly recommend that you just refer to my modifications to the logger instead of merging the code directly (or at least you need to modify and test further).

The following are changes to the logger:

1. Added another label vector in addition to Delta and minmax, which means that the data you want to store is a vector rather than a number, such as multiple Lagrange multipliers.

def register_key(
            self,
            key: str,
            window_length: int | None = None,
            min_and_max: bool = False,
            delta: bool = False,
            vector: int | None = None,
    ) -> None:

Data labeled vector is available in function store for list, tensor and ndarry.

    def store(
            self,
            data: dict[str, Any] | None = None,
            /,
            **kwargs: Any | float | np.ndarray | torch.Tensor,
    ) -> None:
        if data is not None:
            kwargs.update(data)
        for key, val in kwargs.items():
            assert key in self._current_row, f'Key {key} has not been registered'
            if self._headers_vector[key]:
                if isinstance(val, list):
                    if len(val) == self._headers_vector[key]:
                        self._data[key].append(val)
                    else:
                        raise ValueError('Length error')
                elif isinstance(val, torch.Tensor):
                    t_size = val.size()
                    if len(t_size) == 1 and t_size[0] == self._headers_vector[key]:
                        self._data[key].append(val.tolist())
                    else:
                        raise ValueError('Length error')
                elif isinstance(val, np.ndarray):
                    t_size = val.shape
                    if len(t_size) == 1 and t_size[0] == self._headers_vector[key]:
                        self._data[key].append(val.tolist())
                    else:
                        raise ValueError('Length error')
                else:
                    raise ValueError(f'Unsupported type {type(val)}')
            elif isinstance(val, (int, float)):
                self._data[key].append(val)
            elif isinstance(val, torch.Tensor):
                self._data[key].append(val.mean().item())
            elif isinstance(val, np.ndarray):
                self._data[key].append(val.mean())
            else:
                raise ValueError(f'Unsupported type {type(val)}')

In the function get_state, the data marked vector will return the average value or other information according to the second dimension.

    def get_stats(
            self,
            key: str,
            min_and_max: bool = False,
    ) -> tuple[float | list, ...] | list:
        assert key in self._current_row, f'Key {key} has not been registered'
        vals = self._data[key]
        if isinstance(vals, deque):
            vals = list(vals)

        if self._headers_vector[key]:
            vals = np.array(vals)
            _, len_atrs = np.shape(vals)
            if min_and_max:
                means, stds, min_vals, max_vals = [], [], [], []
                for i in range(len_atrs):
                    mean, std, min_val, max_val = dist_statistics_scalar(torch.tensor(vals[:, i]).squeeze(),
                                                                         with_min_and_max=True)
                    means.append(mean.item())
                    stds.append(std.item())
                    min_vals.append(min_val.mean().item())
                    max_vals.append(max_val.mean().item())
                return means, stds, min_vals, max_vals  # not test
            else:
                means = []
                for i in range(len_atrs):
                    try:
                        mean, std = dist_statistics_scalar(torch.tensor(vals[:, i]).squeeze())  # 全是0时没法返回std
                    except:
                        mean = torch.mean(torch.tensor(vals[:, i]).squeeze())
                    means.append(mean.item())
                return means
        else:
            if min_and_max:
                mean, std, min_val, max_val = dist_statistics_scalar(
                    torch.tensor(vals),
                    with_min_and_max=True,
                )
                return mean.item(), min_val.mean().item(), max_val.mean().item(), std.item()
            else:
                mean, std = dist_statistics_scalar(torch.tensor(vals))
                return (mean.item(),)

2. In the console output section, process the data marked vector into the following output form.
   
3. In the tensorboard part, the data marked vector is automatically expanded into multiple data according to the length and stored.
   

For code reusability reasons, I strongly recommend changing the error below to a warning.

assert key not in self._current_row, f'Key {key} has been registered'

And I'm trying to summit the code in pull quest.

[Gaiejj]

Congratulations on successfully implementing the code modifications, and we are very pleased that our suggestions have been helpful to you. We encourage but don't strongly require you to merge your code implementation in the form of a Pull Request. Your contribution to OmniSafe is significant and inspires other researchers in the community, including our developers. If you feel you cannot complete the Pull Request independently, we will refer to your implementation for new feature additions and, on behalf of the OmniSafe open-source community, extend our sincere gratitude to you.

[TomatoZ2]

I think pull request and further testing are beyond my capabilities.

I hope my implementation of logger modifications and multi-constraint SACPID can help you add more features to omnisafe. Below is the code for my implementation. If you need the complete code for testing, please feel free to contact me. I look forward to future support for multi-constraint situations.
code.zip