1{
2  "op": "replace",
3  "filename": "cleanrl/cleanrl/custom_onpolicy_continuous.py",
4  "content": "class Agent(nn.Module):\n    \"\"\"Squashed-Gaussian PPO agent with orthogonal init and state-dependent std.\"\"\"\n\n    def __init__(self, obs_dim, action_dim):\n        super().__init__()\n        h = 256\n        self.critic = nn.Sequential(\n            layer_init(nn.Linear(obs_dim, h)), nn.Tanh(),\n            layer_init(nn.Linear(h, h)), nn.Tanh(),\n            layer_init(nn.Linear(h, 1), std=1.0),\n        )\n        self.actor = nn.Sequential(\n            layer_init(nn.Linear(obs_dim, h)), nn.Tanh(),\n            layer_init(nn.Linear(h, h)), nn.Tanh(),\n        )\n        self.actor_mean = layer_init(nn.Linear(h, action_dim), std=0.01)\n        self.actor_logstd = layer_init(nn.Linear(h, action_dim), std=0.01, bias_const=-0.5)\n\n    def get_value(self, obs):\n        return self.critic(obs)\n\n    def get_action_and_value(self, obs, action=None):\n        hidden = self.actor(obs)\n        mean = self.actor_mean(hidden)\n        logstd = self.actor_logstd(hidden).clamp(-5.0, 1.0)\n        std = logstd.exp()\n        probs = Normal(mean, std)\n        if action is None:\n            raw_action = probs.rsample()\n            action = torch.tanh(raw_action)\n        else:\n            action = action.clamp(-0.999999, 0.999999)\n            raw_action = 0.5 * (torch.log1p(action) - torch.log1p(-action))\n        squash_correction = 2.0 * (np.log(2.0) - raw_action - F.softplus(-2.0 * raw_action))\n        log_prob = probs.log_prob(raw_action) - squash_correction\n        return action, log_prob.sum(1), probs.entropy().sum(1), self.critic(obs)\n\n\ndef compute_losses(agent, mb_obs, mb_actions, mb_logprobs, mb_advantages, mb_returns, mb_values, args):\n    \"\"\"PPO clipped surrogate objective + clipped value loss for squashed Gaussian policy.\"\"\"\n    _, newlogprob, entropy, newvalue = agent.get_action_and_value(mb_obs, mb_actions)\n    logratio = newlogprob - mb_logprobs\n    ratio = logratio.exp()\n\n    with torch.no_grad():\n        approx_kl = ((ratio - 1) - logratio).mean()\n        clipfrac = ((ratio - 1.0).abs() > args.clip_coef).float().mean().item()\n\n    pg_loss1 = -mb_advantages * ratio\n    pg_loss2 = -mb_advantages * torch.clamp(ratio, 1 - args.clip_coef, 1 + args.clip_coef)\n    pg_loss = torch.max(pg_loss1, pg_loss2).mean()\n\n    newvalue = newvalue.view(-1)\n    if args.clip_vloss:\n        v_loss_unclipped = (newvalue - mb_returns) ** 2\n        v_clipped = mb_values + torch.clamp(newvalue - mb_values, -args.clip_coef, args.clip_coef)\n        v_loss_clipped = (v_clipped - mb_returns) ** 2\n        v_loss = 0.5 * torch.max(v_loss_unclipped, v_loss_clipped).mean()\n    else:\n        v_loss = 0.5 * ((newvalue - mb_returns) ** 2).mean()\n\n    entropy_loss = entropy.mean()\n    loss = pg_loss - args.ent_coef * entropy_loss + args.vf_coef * v_loss\n    return loss, pg_loss, v_loss, entropy_loss, approx_kl, clipfrac\n",
5  "after_line": 142,
6  "start_line": 142,
7  "end_line": 215
8}

1{
2  "is_final": false,
3  "use_last": false
4}

1{
2  "is_final": true,
3  "use_last": true
4}

1{
2  "op": "replace",
3  "filename": "cleanrl/cleanrl/custom_onpolicy_continuous.py",
4  "content": "    def get_action_and_value(self, obs, action=None):\n        if not hasattr(self, \"_init_done\"):\n            self._init_done = True\n            layer_init(self.critic[0])\n            layer_init(self.critic[2])\n            layer_init(self.critic[4], std=1.0)\n            layer_init(self.actor_mean[0])\n            layer_init(self.actor_mean[2])\n            layer_init(self.actor_mean[4], std=0.01)\n            nn.init.constant_(self.actor_logstd, -0.5)\n\n        action_mean = torch.clamp(self.actor_mean(obs), -5.0, 5.0)\n        action_logstd = torch.clamp(self.actor_logstd, -4.0, 1.0).expand_as(action_mean)\n        action_std = torch.exp(action_logstd)\n        probs = Normal(action_mean, action_std)\n        if action is None:\n            action = probs.sample()\n        return action, probs.log_prob(action).sum(1), probs.entropy().sum(1), self.critic(obs)\n\n\ndef compute_losses(agent, mb_obs, mb_actions, mb_logprobs, mb_advantages, mb_returns, mb_values, args):\n    \"\"\"Engineered PPO loss with clipped policy/value updates and stabilizers.\"\"\"\n    if args.target_kl is None:\n        args.target_kl = 0.03\n    args.clip_vloss = True\n\n    clip_coef = args.clip_coef\n    _, newlogprob, entropy, newvalue = agent.get_action_and_value(mb_obs, mb_actions)\n    logratio = torch.clamp(newlogprob - mb_logprobs, -20.0, 20.0)\n    ratio = logratio.exp()\n\n    with torch.no_grad():\n        approx_kl = ((ratio - 1) - logratio).mean()\n        clipfrac = ((ratio - 1.0).abs() > clip_coef).float().mean().item()\n\n    mb_advantages = torch.clamp(mb_advantages, -5.0, 5.0)\n    pg_loss1 = -mb_advantages * ratio\n    pg_loss2 = -mb_advantages * torch.clamp(ratio, 1 - clip_coef, 1 + clip_coef)\n    pg_loss = torch.max(pg_loss1, pg_loss2).mean()\n\n    newvalue = newvalue.view(-1)\n    v_loss_unclipped = (newvalue - mb_returns) ** 2\n    v_clipped = mb_values + torch.clamp(newvalue - mb_values, -clip_coef, clip_coef)\n    v_loss_clipped = (v_clipped - mb_returns) ** 2\n    v_loss = 0.5 * torch.max(v_loss_unclipped, v_loss_clipped).mean()\n\n    entropy_loss = entropy.mean()\n    loss = pg_loss - args.ent_coef * entropy_loss + v_loss * args.vf_coef\n\n    return loss, pg_loss, v_loss, entropy_loss, approx_kl, clipfrac",
5  "after_line": 0,
6  "start_line": 175,
7  "end_line": 219
8}

1{
2  "op": "replace",
3  "filename": "cleanrl/cleanrl/custom_onpolicy_continuous.py",
4  "content": "    def get_action_and_value(self, obs, action=None):\n        \"\"\"Bounded Gaussian policy with boundary-aware exploration.\n\n        We sample in an unconstrained space and squash with tanh so that the\n        policy matches the Box action bounds more faithfully than an\n        unsquashed Gaussian followed by environment-side clipping.\n        Additionally, the exploration scale is reduced near saturated means,\n        which improves stability for precise-control regimes.\n        \"\"\"\n        eps = 1e-6\n        action_mean = self.actor_mean(obs)\n\n        # Mean-dependent exploration: keep broad exploration near the origin\n        # but shrink it smoothly as the policy pushes toward action bounds.\n        squash_mean = torch.tanh(action_mean)\n        mean_gate = 1.0 - squash_mean.pow(2)\n        action_logstd = self.actor_logstd.expand_as(action_mean) + 0.5 * torch.log(mean_gate + eps)\n        action_logstd = torch.clamp(action_logstd, -5.0, 2.0)\n        action_std = torch.exp(action_logstd)\n        probs = Normal(action_mean, action_std)\n\n        if action is None:\n            pre_tanh_action = probs.rsample()\n            action = torch.tanh(pre_tanh_action)\n        else:\n            action = torch.clamp(action, -1.0 + eps, 1.0 - eps)\n            pre_tanh_action = 0.5 * (torch.log1p(action) - torch.log1p(-action))\n\n        logprob = probs.log_prob(pre_tanh_action) - torch.log(1.0 - action.pow(2) + eps)\n        logprob = logprob.sum(1)\n\n        # Monte-Carlo entropy estimate for the squashed policy.\n        entropy = probs.entropy().sum(1) + torch.log(1.0 - action.pow(2) + eps).sum(1)\n        return action, logprob, entropy, self.critic(obs)\n\n\ndef compute_losses(agent, mb_obs, mb_actions, mb_logprobs, mb_advantages, mb_returns, mb_values, args):\n    \"\"\"Hybrid bounded PPO with positive-advantage regression and expectile critic.\n\n    Main ingredients:\n      1) tanh-squashed bounded policy (implemented in get_action_and_value)\n      2) asymmetric PPO clipping: more room to reinforce good actions and a\n         tighter correction on bad ones\n      3) positive-advantage log-ratio bonus: an on-policy AWR-style term that\n         selectively boosts high-advantage samples without abandoning PPO's\n         trust-region behavior\n      4) clipped expectile value regression for a slightly optimistic critic\n    \"\"\"\n    _awr_coef = 0.5\n    _awr_beta = 0.75\n    _kl_coef = 0.02\n    _expectile = 0.65\n\n    _, newlogprob, entropy, newvalue = agent.get_action_and_value(mb_obs, mb_actions)\n    logratio = newlogprob - mb_logprobs\n    ratio = logratio.exp()\n\n    with torch.no_grad():\n        approx_kl = ((ratio - 1.0) - logratio).mean()\n        clipfrac = ((ratio - 1.0).abs() > args.clip_coef).float().mean().item()\n\n    adv = torch.clamp(mb_advantages, -4.0, 4.0)\n    shaped_adv = torch.where(\n        adv >= 0,\n        adv * (1.0 + 0.25 * adv),\n        0.75 * adv,\n    )\n\n    pos_clip_ratio = torch.clamp(ratio, 1.0 - args.clip_coef, 1.0 + 1.5 * args.clip_coef)\n    neg_clip_ratio = torch.clamp(ratio, 1.0 - 0.75 * args.clip_coef, 1.0 + args.clip_coef)\n    clipped_ratio = torch.where(shaped_adv >= 0, pos_clip_ratio, neg_clip_ratio)\n    ppo_loss = -torch.min(ratio * shaped_adv, clipped_ratio * shaped_adv).mean()\n\n    with torch.no_grad():\n        pos_weights = torch.exp(torch.clamp(torch.relu(adv) / _awr_beta, max=2.0)) - 1.0\n        pos_weights = pos_weights / (pos_weights.mean() + 1e-8)\n    awr_loss = -(logratio * pos_weights).mean()\n\n    kl_loss = ((ratio - 1.0) - logratio).mean()\n    pg_loss = ppo_loss + _awr_coef * awr_loss + _kl_coef * kl_loss\n\n    newvalue = newvalue.view(-1)\n    v_error = mb_returns - newvalue\n    v_weight = torch.where(v_error > 0, _expectile, 1.0 - _expectile)\n    v_loss_unclipped = v_weight * v_error.pow(2)\n    if args.clip_vloss:\n        v_clipped = mb_values + torch.clamp(newvalue - mb_values, -args.clip_coef, args.clip_coef)\n        v_error_clipped = mb_returns - v_clipped\n        v_weight_clipped = torch.where(v_error_clipped > 0, _expectile, 1.0 - _expectile)\n        v_loss_clipped = v_weight_clipped * v_error_clipped.pow(2)\n        v_loss = 0.5 * torch.max(v_loss_unclipped, v_loss_clipped).mean()\n    else:\n        v_loss = 0.5 * v_loss_unclipped.mean()\n\n    entropy_loss = entropy.mean()\n    loss = pg_loss - args.ent_coef * entropy_loss + v_loss * args.vf_coef\n    return loss, pg_loss, v_loss, entropy_loss, approx_kl, clipfrac",
5  "after_line": 174,
6  "start_line": 175,
7  "end_line": 219
8}

1{
2  "is_final": false,
3  "use_last": false
4}

1{
2  "is_final": false,
3  "use_last": false
4}

1{
2  "is_final": true,
3  "use_last": true
4}

1{
2  "is_final": true,
3  "use_last": true
4}

1{
2  "op": "replace",
3  "filename": "cleanrl/cleanrl/custom_onpolicy_continuous.py",
4  "content": "    def get_action_and_value(self, obs, action=None):\n        # Action-bound-aware policy: tanh-squashed Gaussian.\n        # This removes the train/eval mismatch caused by environment-side clipping.\n        action_mean = self.actor_mean(obs)\n        action_logstd = self.actor_logstd.expand_as(action_mean).clamp(-5.0, 1.0)\n        action_std = torch.exp(action_logstd)\n        probs = Normal(action_mean, action_std)\n        eps = 1e-6\n\n        if action is None:\n            pre_tanh_action = probs.rsample()\n            action = torch.tanh(pre_tanh_action)\n        else:\n            action = action.clamp(-1.0 + eps, 1.0 - eps)\n            pre_tanh_action = 0.5 * (torch.log1p(action) - torch.log1p(-action))\n\n        logprob = probs.log_prob(pre_tanh_action) - torch.log(1.0 - action.pow(2) + eps)\n        logprob = logprob.sum(1)\n\n        # Exact entropy is intractable after squashing; base Gaussian entropy works\n        # well enough here and keeps the objective simple.\n        entropy = probs.entropy().sum(1)\n        return action, logprob, entropy, self.critic(obs)\n\n\ndef compute_losses(agent, mb_obs, mb_actions, mb_logprobs, mb_advantages, mb_returns, mb_values, args):\n    \"\"\"Asymmetric clipped PPO with positive-advantage regression and expectile value fitting.\n\n    Key ideas:\n    1) tanh-squashed Gaussian policy for action-bound awareness,\n    2) asymmetric PPO clipping: larger trust region for positive advantages,\n       tighter clipping for negative advantages,\n    3) small positive-advantage regression bonus (AWR-style) to reinforce strong actions,\n    4) clipped expectile critic to correct underestimation faster than overestimation.\n    \"\"\"\n    _pos_clip_mult = 1.5\n    _neg_clip_mult = 0.5\n    _awr_beta = 1.5\n    _awr_max_weight = 3.0\n    _awr_coef = 0.05\n    _trust_coef = 0.02\n    _expectile = 0.6\n\n    _, newlogprob, entropy, newvalue = agent.get_action_and_value(mb_obs, mb_actions)\n    logratio = newlogprob - mb_logprobs\n    ratio = logratio.exp()\n\n    with torch.no_grad():\n        approx_kl = ((ratio - 1) - logratio).mean()\n        clipfrac = ((ratio - 1.0).abs() > args.clip_coef).float().mean().item()\n\n    # Asymmetric clipped surrogate: reinforce good actions more aggressively,\n    # but damp noisy penalties on bad actions.\n    pos_clipped_ratio = torch.clamp(ratio, 1.0 - args.clip_coef, 1.0 + args.clip_coef * _pos_clip_mult)\n    neg_clipped_ratio = torch.clamp(ratio, 1.0 - args.clip_coef * _neg_clip_mult, 1.0 + args.clip_coef)\n    clipped_ratio = torch.where(mb_advantages >= 0, pos_clipped_ratio, neg_clipped_ratio)\n    surrogate = torch.min(ratio * mb_advantages, clipped_ratio * mb_advantages)\n    ppo_loss = -surrogate.mean()\n\n    # Positive-advantage regression bonus: pushes probability mass toward the\n    # best sampled actions without letting a few outliers dominate.\n    with torch.no_grad():\n        pos_adv = torch.relu(mb_advantages)\n        awr_weight = torch.exp(pos_adv / _awr_beta) - 1.0\n        awr_weight = torch.clamp(awr_weight, max=_awr_max_weight)\n    regression_loss = -(awr_weight * newlogprob).mean()\n\n    # Mild ratio regularizer because the regression bonus bypasses PPO clipping.\n    trust_loss = (logratio.pow(2)).mean()\n    pg_loss = ppo_loss + _awr_coef * regression_loss + _trust_coef * trust_loss\n\n    # Clipped expectile critic: adapt faster when returns exceed current values,\n    # while retaining PPO-style value clipping for stability.\n    newvalue = newvalue.view(-1)\n    value_error = mb_returns - newvalue\n    value_weight = torch.where(value_error > 0, _expectile, 1.0 - _expectile)\n    v_loss_unclipped = value_weight * value_error.pow(2)\n    if args.clip_vloss:\n        v_clipped = mb_values + torch.clamp(newvalue - mb_values, -args.clip_coef, args.clip_coef)\n        clipped_error = mb_returns - v_clipped\n        clipped_weight = torch.where(clipped_error > 0, _expectile, 1.0 - _expectile)\n        v_loss_clipped = clipped_weight * clipped_error.pow(2)\n        v_loss = 0.5 * torch.max(v_loss_unclipped, v_loss_clipped).mean()\n    else:\n        v_loss = 0.5 * v_loss_unclipped.mean()\n\n    entropy_loss = entropy.mean()\n    loss = pg_loss - args.ent_coef * entropy_loss + v_loss * args.vf_coef\n\n    return loss, pg_loss, v_loss, entropy_loss, approx_kl, clipfrac",
5  "after_line": 174,
6  "start_line": 175,
7  "end_line": 219
8}