1{
2  "op": "replace",
3  "filename": "omnisafe/omnisafe/algorithms/on_policy/naive_lagrange/custom_lag.py",
4  "start_line": 48,
5  "end_line": 70,
6  "content": "    def _init(self) -> None:\n        super()._init()\n        self._cost_limit: float = self._cfgs.lagrange_cfgs.cost_limit\n        self._lagrangian_multiplier: float = 0.0\n        # EMA for cost smoothing\n        self._cost_ema: float = 0.0\n        self._cost_ema_beta: float = 0.85\n        # PID integral state with anti-windup\n        self._pid_integral: float = 0.0\n        # Previous cost for derivative\n        self._prev_cost_ema: float = 0.0\n        # Whether we've seen any cost yet\n        self._initialized: bool = False\n\n    def _init_log(self) -> None:\n        super()._init_log()\n        self._logger.register_key('Metrics/LagrangeMultiplier', min_and_max=True)\n\n    def _update(self) -> None:\n        Jc = self._logger.get_stats('Metrics/EpCost')[0]\n        assert not np.isnan(Jc), 'cost is nan'\n\n        # Update EMA\n        if not self._initialized:\n            self._cost_ema = Jc\n            self._prev_cost_ema = Jc\n            self._initialized = True\n        else:\n            beta = self._cost_ema_beta\n            self._cost_ema = beta * self._cost_ema + (1 - beta) * Jc\n\n        # Compute violation signal (smoothed)\n        violation = self._cost_ema - self._cost_limit\n        \n        # --- Adaptive PID multiplier update ---\n        # Proportional: adaptive gain that increases with violation severity\n        kp_base = 0.05\n        kp = kp_base + 0.15 * max(0.0, violation)\n        p_term = kp * max(0.0, violation)\n        \n        # Integral: accumulate only when violating (anti-windup)\n        ki = 0.02\n        if violation > 0:\n            self._pid_integral += ki * violation\n            self._pid_integral = min(self._pid_integral, 5.0)  # anti-windup cap\n        else:\n            # Decay integral when safe\n            self._pid_integral *= 0.98\n        \n        # Derivative: on cost trend (not on violation)\n        kd = 0.03\n        cost_derivative = self._cost_ema - self._prev_cost_ema\n        self._prev_cost_ema = self._cost_ema\n        d_term = kd * max(0.0, cost_derivative)  # only penalize increasing cost\n        \n        # Combine PID terms\n        pid_output = p_term + self._pid_integral + d_term\n        \n        # Update multiplier with bounded step\n        if violation > 0:\n            self._lagrangian_multiplier += pid_output\n        else:\n            # Decay when safely under limit\n            self._lagrangian_multiplier *= 0.95\n        \n        # Ensure non-negative and bound\n        self._lagrangian_multiplier = max(0.0, min(self._lagrangian_multiplier, 5.0))\n\n        super()._update()\n        self._logger.store({'Metrics/LagrangeMultiplier': self._lagrangian_multiplier})\n\n    def _compute_adv_surrogate(self, adv_r: torch.Tensor, adv_c: torch.Tensor) -> torch.Tensor:\n        \"\"\"Combine reward and cost advantages using adaptive barrier formulation.\n        \n        Key innovations:\n        1. Bounded penalty via tanh (prevents reward signal collapse)\n        2. Adaptive normalization using reward advantage magnitude\n        3. Pure reward optimization when multiplier is 0 (like naive baseline)\n        \"\"\"\n        lam = self._lagrangian_multiplier\n        \n        if lam <= 1e-6:\n            # When safe: optimize reward only\n            return adv_r\n        \n        # Soft-bounded penalty (prevents over-penalization)\n        # tanh bounds the effective penalty to [0, ~0.9] even for large lambda\n        soft_lam = float(np.tanh(lam))\n        \n        # Normalize cost advantage relative to reward advantage magnitude\n        # This prevents cost from dominating when reward signal is small\n        adv_r_abs = adv_r.abs().mean()\n        adv_c_scale = adv_c.abs().mean() + 1e-8\n        normalized_adv_c = adv_c * (adv_r_abs / adv_c_scale)\n        \n        # Combined advantage with adaptive penalty\n        # When lam is small: mostly reward, slight cost awareness\n        # When lam is large: balanced reward/cost with bounded penalty\n        return adv_r - soft_lam * normalized_adv_c"
7}

1{
2  "filename": "omnisafe/omnisafe/algorithms/on_policy/naive_lagrange/custom_lag.py",
3  "start_line": 48,
4  "end_line": 146,
5  "op": "replace",
6  "content": "    def _init(self) -> None:\n        super()._init()\n        self._cost_limit: float = self._cfgs.lagrange_cfgs.cost_limit\n        self._lagrangian_multiplier: float = 0.0\n        # EMA for cost smoothing\n        self._cost_ema: float = 0.0\n        self._cost_ema_beta: float = 0.85\n        # PID state\n        self._pid_integral: float = 0.0\n        self._prev_cost_ema: float = 0.0\n        self._initialized: bool = False\n\n    def _init_log(self) -> None:\n        super()._init_log()\n        self._logger.register_key('Metrics/LagrangeMultiplier', min_and_max=True)\n\n    def _update(self) -> None:\n        Jc = self._logger.get_stats('Metrics/EpCost')[0]\n        assert not np.isnan(Jc), 'cost is nan'\n\n        # Update EMA of cost\n        if not self._initialized:\n            self._cost_ema = Jc\n            self._prev_cost_ema = Jc\n            self._initialized = True\n        else:\n            beta = self._cost_ema_beta\n            self._cost_ema = beta * self._cost_ema + (1 - beta) * Jc\n\n        # Compute smoothed violation\n        violation = self._cost_ema - self._cost_limit\n        \n        # --- Aggressive PID multiplier update ---\n        kp = 0.15\n        ki = 0.05\n        kd = 0.05\n        \n        p_term = kp * max(0.0, violation)\n        \n        if violation > 0:\n            self._pid_integral += ki * violation\n            self._pid_integral = min(self._pid_integral, 3.0)\n        else:\n            # Decay faster when under limit\n            self._pid_integral *= 0.92\n        \n        cost_derivative = self._cost_ema - self._prev_cost_ema\n        self._prev_cost_ema = self._cost_ema\n        d_term = kd * max(0.0, cost_derivative)\n        \n        pid_output = p_term + self._pid_integral + d_term\n        \n        # Update multiplier\n        if violation > 0:\n            self._lagrangian_multiplier += pid_output\n        else:\n            self._lagrangian_multiplier *= 0.9\n        \n        self._lagrangian_multiplier = max(0.0, min(self._lagrangian_multiplier, 5.0))\n\n        super()._update()\n        self._logger.store({'Metrics/LagrangeMultiplier': self._lagrangian_multiplier})\n\n    def _compute_adv_surrogate(self, adv_r: torch.Tensor, adv_c: torch.Tensor) -> torch.Tensor:\n        \"\"\"Combine reward and cost advantages using Lagrangian formulation.\n        \n        Uses standard: adv_r - lambda * adv_c, which directly penalizes \n        actions that increase cost relative to the advantage baseline.\n        \"\"\"\n        lam = self._lagrangian_multiplier\n        \n        if lam <= 1e-6:\n            return adv_r\n        \n        # Standard Lagrangian: penalize cost advantage\n        return adv_r - lam * adv_c"
7}

1{
2  "n": 2
3}

1{
2  "filename": "omnisafe/omnisafe/algorithms/on_policy/naive_lagrange/custom_lag.py",
3  "start_line": 48,
4  "end_line": 123,
5  "op": "replace",
6  "content": "    def _init(self) -> None:\n        super()._init()\n        self._cost_limit: float = self._cfgs.lagrange_cfgs.cost_limit\n        self._lagrangian_multiplier: float = 0.0\n        self._cost_ema: float = 0.0\n        self._prev_cost_ema: float = 0.0\n        self._pid_integral: float = 0.0\n        self._initialized: bool = False\n\n    def _init_log(self) -> None:\n        super()._init_log()\n        self._logger.register_key('Metrics/LagrangeMultiplier', min_and_max=True)\n\n    def _update(self) -> None:\n        Jc = self._logger.get_stats('Metrics/EpCost')[0]\n        assert not np.isnan(Jc), 'cost is nan'\n\n        if not self._initialized:\n            self._cost_ema = Jc\n            self._prev_cost_ema = Jc\n            self._initialized = True\n        else:\n            self._cost_ema = 0.85 * self._cost_ema + 0.15 * Jc\n\n        violation = self._cost_ema - self._cost_limit\n        \n        # Aggressive adaptive PID\n        kp = 0.1 + 0.5 * max(0.0, violation / self._cost_limit)\n        ki = 0.05\n        kd = 0.05\n\n        p_term = kp * max(0.0, violation)\n        \n        if violation > 0:\n            self._pid_integral += ki * violation\n            self._pid_integral = min(self._pid_integral, 10.0)\n        else:\n            self._pid_integral = max(0.0, self._pid_integral - 0.1)\n        \n        cost_deriv = self._cost_ema - self._prev_cost_ema\n        self._prev_cost_ema = self._cost_ema\n        d_term = kd * max(0.0, cost_deriv)\n        \n        new_lambda = p_term + self._pid_integral + d_term\n        self._lagrangian_multiplier = max(0.0, new_lambda)\n\n        super()._update()\n        self._logger.store({'Metrics/LagrangeMultiplier': self._lagrangian_multiplier})\n\n    def _compute_adv_surrogate(self, adv_r: torch.Tensor, adv_c: torch.Tensor) -> torch.Tensor:\n        \"\"\"Dual-objective advantage with normalized penalty.\"\"\"\n        lam = self._lagrangian_multiplier\n        if lam < 1e-8:\n            return adv_r\n        # Normalized Lagrangian: prevents reward collapse\n        return (adv_r - lam * adv_c) / (1.0 + lam)"
7}

1{
2  "filename": "omnisafe/omnisafe/algorithms/on_policy/naive_lagrange/custom_lag.py",
3  "op": "replace",
4  "content": "# ===================================================================\n# EDITABLE: Custom imports\n# ===================================================================\nfrom collections import deque\n\n# ===================================================================\n# FIXED: Algorithm class definition\n# ===================================================================\n@registry.register\nclass CustomLag(PPO):\n    \"\"\"Custom Lagrangian-based safe RL algorithm.\n\n    Extends PPO with constraint handling for safe reinforcement learning.\n    The agent must design:\n      1. _init: Initialize constraint handler state (call super()._init() first)\n      2. _init_log: Register logging keys (call super()._init_log() first)\n      3. _update: Update lagrangian multiplier, then call super()._update()\n      4. _compute_adv_surrogate: Combine reward and cost advantages\n\n    Available config:\n        self._cfgs.lagrange_cfgs.cost_limit   (float, default 25.0)\n        self._cfgs.lagrange_cfgs.lambda_lr    (float, default 0.035)\n\n    Available logger:\n        self._logger.get_stats('Metrics/EpCost')[0]  -- current mean episode cost\n        self._logger.store({'key': value})             -- log a metric value\n    \"\"\"\n\n    # ===============================================================\n    # EDITABLE: Constraint handling mechanism\n    # ===============================================================\n    def _init(self) -> None:\n        super()._init()\n        self._cost_limit: float = self._cfgs.lagrange_cfgs.cost_limit\n        self._lagrangian_multiplier: float = 0.0\n\n    def _init_log(self) -> None:\n        super()._init_log()\n        self._logger.register_key('Metrics/LagrangeMultiplier', min_and_max=True)\n\n    def _update(self) -> None:\n        Jc = self._logger.get_stats('Metrics/EpCost')[0]\n        assert not np.isnan(Jc), 'cost is nan'\n        # Default: no multiplier update -- agent should design this\n        super()._update()\n        self._logger.store({'Metrics/LagrangeMultiplier': self._lagrangian_multiplier})\n\n    def _compute_adv_surrogate(self, adv_r: torch.Tensor, adv_c: torch.Tensor) -> torch.Tensor:\n        \"\"\"Combine reward and cost advantages.\n\n        Default: only use reward advantage (ignores safety constraints entirely).\n        Agent should incorporate self._lagrangian_multiplier to penalize cost.\n        \"\"\"\n        return adv_r",
5  "start_line": 17,
6  "end_line": 70
7}

1{
2  "content": "from collections import deque",
3  "filename": "omnisafe/omnisafe/algorithms/on_policy/naive_lagrange/custom_lag.py",
4  "op": "replace",
5  "start_line": 20,
6  "end_line": 20
7}

1{
2  "content": "    def _init(self) -> None:\n        super()._init()\n        self._cost_limit: float = self._cfgs.lagrange_cfgs.cost_limit\n        self._lagrangian_multiplier: float = 0.0\n        self._cost_ema: float = 0.0\n        self._prev_cost_ema: float = 0.0\n        self._pid_integral: float = 0.0\n        self._initialized: bool = False\n\n    def _init_log(self) -> None:\n        super()._init_log()\n        self._logger.register_key('Metrics/LagrangeMultiplier', min_and_max=True)\n\n    def _update(self) -> None:\n        Jc = self._logger.get_stats('Metrics/EpCost')[0]\n        assert not np.isnan(Jc), 'cost is nan'\n        \n        if not self._initialized:\n            self._cost_ema = Jc\n            self._prev_cost_ema = Jc\n            self._initialized = True\n        else:\n            self._cost_ema = 0.85 * self._cost_ema + 0.15 * Jc\n\n        violation = self._cost_ema - self._cost_limit\n        \n        kp = 0.1 + 0.5 * max(0.0, violation / self._cost_limit)\n        ki = 0.05\n        kd = 0.05\n\n        p_term = kp * max(0.0, violation)\n        \n        if violation > 0:\n            self._pid_integral += ki * violation\n            self._pid_integral = min(self._pid_integral, 10.0)\n        else:\n            self._pid_integral = max(0.0, self._pid_integral - 0.1)\n        \n        cost_deriv = self._cost_ema - self._prev_cost_ema\n        self._prev_cost_ema = self._cost_ema\n        d_term = kd * max(0.0, cost_deriv)\n        \n        new_lambda = p_term + self._pid_integral + d_term\n        self._lagrangian_multiplier = max(0.0, new_lambda)\n\n        super()._update()\n        self._logger.store({'Metrics/LagrangeMultiplier': self._lagrangian_multiplier})\n\n    def _compute_adv_surrogate(self, adv_r: torch.Tensor, adv_c: torch.Tensor) -> torch.Tensor:\n        lam = self._lagrangian_multiplier\n        if lam < 1e-8:\n            return adv_r\n        return (adv_r - lam * adv_c) / (1.0 + lam)",
3  "filename": "omnisafe/omnisafe/algorithms/on_policy/naive_lagrange/custom_lag.py",
4  "op": "replace",
5  "start_line": 48,
6  "end_line": 123
7}