make max bound SAC

reiniscimurs · reiniscimurs · commit 4dda8fa92658 · 2025-04-03T10:19:05.000+02:00
diff --git a/robot_nav/models/SAC/BSA1C.py b/robot_nav/models/SAC/BSA1C.py
@@ -0,0 +1,316 @@
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from statistics import mean
+import robot_nav.models.SAC.SAC_utils as utils
+from robot_nav.models.SAC.BSA1C_critic import QCritic as critic_model
+from robot_nav.models.SAC.SAC_actor import DiagGaussianActor as actor_model
+from torch.utils.tensorboard import SummaryWriter
+from robot_nav.utils import get_max_bound
+
+
+class BSA1C(object):
+    """SAC algorithm."""
+
+    def __init__(
+        self,
+        state_dim,
+        action_dim,
+        device,
+        max_action,
+        discount=0.99,
+        init_temperature=0.1,
+        alpha_lr=1e-4,
+        alpha_betas=(0.9, 0.999),
+        actor_lr=1e-4,
+        actor_betas=(0.9, 0.999),
+        actor_update_frequency=1,
+        critic_lr=1e-4,
+        critic_betas=(0.9, 0.999),
+        critic_tau=0.005,
+        critic_target_update_frequency=2,
+        learnable_temperature=True,
+        save_every=0,
+        load_model=False,
+        log_dist_and_hist=False,
+        save_directory=Path("robot_nav/models/SAC/checkpoint"),
+        model_name="BSAC",
+        load_directory=Path("robot_nav/models/SAC/checkpoint"),
+        bound_weight=0.25,
+    ):
+        super().__init__()
+
+        self.state_dim = state_dim
+        self.action_dim = action_dim
+        self.action_range = (-max_action, max_action)
+        self.device = torch.device(device)
+        self.discount = discount
+        self.critic_tau = critic_tau
+        self.actor_update_frequency = actor_update_frequency
+        self.critic_target_update_frequency = critic_target_update_frequency
+        self.learnable_temperature = learnable_temperature
+        self.save_every = save_every
+        self.model_name = model_name
+        self.save_directory = save_directory
+        self.log_dist_and_hist = log_dist_and_hist
+        self.bound_weight = bound_weight
+
+        self.train_metrics_dict = {
+            "train_critic/loss_av": [],
+            "train_actor/loss_av": [],
+            "train_actor/target_entropy_av": [],
+            "train_actor/entropy_av": [],
+            "train_alpha/loss_av": [],
+            "train_alpha/value_av": [],
+            "train/batch_reward_av": [],
+        }
+
+        self.critic = critic_model(
+            obs_dim=self.state_dim,
+            action_dim=action_dim,
+            hidden_dim=400,
+            hidden_depth=2,
+        ).to(self.device)
+        self.critic_target = critic_model(
+            obs_dim=self.state_dim,
+            action_dim=action_dim,
+            hidden_dim=400,
+            hidden_depth=2,
+        ).to(self.device)
+        self.critic_target.load_state_dict(self.critic.state_dict())
+
+        self.actor = actor_model(
+            obs_dim=self.state_dim,
+            action_dim=action_dim,
+            hidden_dim=400,
+            hidden_depth=2,
+            log_std_bounds=[-5, 2],
+        ).to(self.device)
+
+        if load_model:
+            self.load(filename=model_name, directory=load_directory)
+
+        self.log_alpha = torch.tensor(np.log(init_temperature)).to(self.device)
+        self.log_alpha.requires_grad = True
+        # set target entropy to -|A|
+        self.target_entropy = -action_dim
+
+        # optimizers
+        self.actor_optimizer = torch.optim.Adam(
+            self.actor.parameters(), lr=actor_lr, betas=actor_betas
+        )
+
+        self.critic_optimizer = torch.optim.Adam(
+            self.critic.parameters(), lr=critic_lr, betas=critic_betas
+        )
+
+        self.log_alpha_optimizer = torch.optim.Adam(
+            [self.log_alpha], lr=alpha_lr, betas=alpha_betas
+        )
+
+        self.critic_target.train()
+
+        self.actor.train(True)
+        self.critic.train(True)
+        self.step = 0
+        self.writer = SummaryWriter(comment=model_name)
+
+    def save(self, filename, directory):
+        Path(directory).mkdir(parents=True, exist_ok=True)
+        torch.save(self.actor.state_dict(), "%s/%s_actor.pth" % (directory, filename))
+        torch.save(self.critic.state_dict(), "%s/%s_critic.pth" % (directory, filename))
+        torch.save(
+            self.critic_target.state_dict(),
+            "%s/%s_critic_target.pth" % (directory, filename),
+        )
+
+    def load(self, filename, directory):
+        self.actor.load_state_dict(
+            torch.load("%s/%s_actor.pth" % (directory, filename))
+        )
+        self.critic.load_state_dict(
+            torch.load("%s/%s_critic.pth" % (directory, filename))
+        )
+        self.critic_target.load_state_dict(
+            torch.load("%s/%s_critic_target.pth" % (directory, filename))
+        )
+        print(f"Loaded weights from: {directory}")
+
+    def train(self, replay_buffer, iterations, batch_size):
+        for _ in range(iterations):
+            self.update(
+                replay_buffer=replay_buffer, step=self.step, batch_size=batch_size
+            )
+
+        for key, value in self.train_metrics_dict.items():
+            if len(value):
+                self.writer.add_scalar(key, mean(value), self.step)
+            self.train_metrics_dict[key] = []
+        self.step += 1
+
+        if self.save_every > 0 and self.step % self.save_every == 0:
+            self.save(filename=self.model_name, directory=self.save_directory)
+
+    @property
+    def alpha(self):
+        return self.log_alpha.exp()
+
+    def get_action(self, obs, add_noise):
+        if add_noise:
+            return (
+                self.act(obs) + np.random.normal(0, 0.2, size=self.action_dim)
+            ).clip(self.action_range[0], self.action_range[1])
+        else:
+            return self.act(obs)
+
+    def act(self, obs, sample=False):
+        obs = torch.FloatTensor(obs).to(self.device)
+        obs = obs.unsqueeze(0)
+        dist = self.actor(obs)
+        action = dist.sample() if sample else dist.mean
+        action = action.clamp(*self.action_range)
+        assert action.ndim == 2 and action.shape[0] == 1
+        return utils.to_np(action[0])
+
+    def update_critic(self, obs, action, reward, next_obs, done, step):
+        dist = self.actor(next_obs)
+        next_action = dist.rsample()
+        log_prob = dist.log_prob(next_action).sum(-1, keepdim=True)
+        target_q = self.critic_target(next_obs, next_action)
+        target_V = target_q - self.alpha.detach() * log_prob
+        target_Q = reward + ((1 - done) * self.discount * target_V)
+        target_Q = target_Q.detach()
+
+        # get current Q estimates
+        current_Q = self.critic(obs, action)
+
+        max_bound = get_max_bound(
+            next_obs, self.discount, 0.5, 1, 0.3, 10, 100, reward, done, self.device
+        )
+
+        max_excess_Q = F.relu(current_Q - max_bound)
+        max_bound_loss = (max_excess_Q**2).mean()
+        max_bound_loss = self.bound_weight * max_bound_loss
+        critic_loss = (
+            F.mse_loss(current_Q, target_Q)
+            + max_bound_loss
+        )
+        self.train_metrics_dict["train_critic/loss_av"].append(critic_loss.item())
+        self.writer.add_scalar("train_critic/loss", critic_loss, step)
+        self.writer.add_scalar("train_critic/max_bound_loss", max_bound_loss, step)
+
+        # Optimize the critic
+        self.critic_optimizer.zero_grad()
+        critic_loss.backward()
+        self.critic_optimizer.step()
+        if self.log_dist_and_hist:
+            self.critic.log(self.writer, step)
+
+    def update_actor_and_alpha(self, obs, step):
+        dist = self.actor(obs)
+        action = dist.rsample()
+        log_prob = dist.log_prob(action).sum(-1, keepdim=True)
+        actor_Q = self.critic(obs, action)
+
+        # actor_Q = torch.min(actor_Q, max_bound)
+        actor_loss = (self.alpha.detach() * log_prob - actor_Q).mean()
+        self.train_metrics_dict["train_actor/loss_av"].append(actor_loss.item())
+        self.train_metrics_dict["train_actor/target_entropy_av"].append(
+            self.target_entropy
+        )
+        self.train_metrics_dict["train_actor/entropy_av"].append(
+            -log_prob.mean().item()
+        )
+        self.writer.add_scalar("train_actor/loss", actor_loss, step)
+        self.writer.add_scalar("train_actor/target_entropy", self.target_entropy, step)
+        self.writer.add_scalar("train_actor/entropy", -log_prob.mean(), step)
+
+        # optimize the actor
+        self.actor_optimizer.zero_grad()
+        actor_loss.backward()
+        self.actor_optimizer.step()
+        if self.log_dist_and_hist:
+            self.actor.log(self.writer, step)
+
+        if self.learnable_temperature:
+            self.log_alpha_optimizer.zero_grad()
+            alpha_loss = (
+                self.alpha * (-log_prob - self.target_entropy).detach()
+            ).mean()
+            self.train_metrics_dict["train_alpha/loss_av"].append(alpha_loss.item())
+            self.train_metrics_dict["train_alpha/value_av"].append(self.alpha.item())
+            self.writer.add_scalar("train_alpha/loss", alpha_loss, step)
+            self.writer.add_scalar("train_alpha/value", self.alpha, step)
+            alpha_loss.backward()
+            self.log_alpha_optimizer.step()
+
+    def update(
+        self,
+        replay_buffer,
+        step,
+        batch_size,
+        max_lin_vel=0.5,
+        max_ang_vel=1,
+        goal_reward=100,
+        distance_norm=10,
+        time_step=0.3,
+    ):
+        (
+            batch_states,
+            batch_actions,
+            batch_rewards,
+            batch_dones,
+            batch_next_states,
+        ) = replay_buffer.sample_batch(batch_size)
+
+        state = torch.Tensor(batch_states).to(self.device)
+        next_state = torch.Tensor(batch_next_states).to(self.device)
+        action = torch.Tensor(batch_actions).to(self.device)
+        reward = torch.Tensor(batch_rewards).to(self.device)
+        done = torch.Tensor(batch_dones).to(self.device)
+
+        self.train_metrics_dict["train/batch_reward_av"].append(
+            batch_rewards.mean().item()
+        )
+        self.writer.add_scalar("train/batch_reward", batch_rewards.mean(), step)
+
+        self.update_critic(state, action, reward, next_state, done, step)
+
+        if step % self.actor_update_frequency == 0:
+            self.update_actor_and_alpha(state, step)
+
+        if step % self.critic_target_update_frequency == 0:
+            utils.soft_update_params(self.critic, self.critic_target, self.critic_tau)
+
+    def prepare_state(self, latest_scan, distance, cos, sin, collision, goal, action):
+        # update the returned data from ROS into a form used for learning in the current model
+        latest_scan = np.array(latest_scan)
+
+        inf_mask = np.isinf(latest_scan)
+        latest_scan[inf_mask] = 7.0
+
+        max_bins = self.state_dim - 5
+        bin_size = int(np.ceil(len(latest_scan) / max_bins))
+
+        # Initialize the list to store the minimum values of each bin
+        min_values = []
+
+        # Loop through the data and create bins
+        for i in range(0, len(latest_scan), bin_size):
+            # Get the current bin
+            bin = latest_scan[i : i + min(bin_size, len(latest_scan) - i)]
+            # Find the minimum value in the current bin and append it to the min_values list
+            min_values.append(min(bin) / 7)
+
+        # Normalize to [0, 1] range
+        distance /= 10
+        lin_vel = action[0] * 2
+        ang_vel = (action[1] + 1) / 2
+        state = min_values + [distance, cos, sin] + [lin_vel, ang_vel]
+
+        assert len(state) == self.state_dim
+        terminal = 1 if collision or goal else 0
+
+        return state, terminal
diff --git a/robot_nav/models/SAC/BSA1C_critic.py b/robot_nav/models/SAC/BSA1C_critic.py
@@ -0,0 +1,30 @@
+import torch
+from torch import nn
+
+import robot_nav.models.SAC.SAC_utils as utils
+
+
+class QCritic(nn.Module):
+    """Critic network, employes double Q-learning."""
+
+    def __init__(self, obs_dim, action_dim, hidden_dim, hidden_depth):
+        super().__init__()
+
+        self.Q1 = utils.mlp(obs_dim + action_dim, hidden_dim, 1, hidden_depth)
+
+        self.outputs = dict()
+        self.apply(utils.weight_init)
+
+    def forward(self, obs, action):
+        assert obs.size(0) == action.size(0)
+
+        obs_action = torch.cat([obs, action], dim=-1)
+        q1 = self.Q1(obs_action)
+
+        self.outputs["q1"] = q1
+
+        return q1
+
+    def log(self, writer, step):
+        for k, v in self.outputs.items():
+            writer.add_histogram(f"train_critic/{k}_hist", v, step)
diff --git a/robot_nav/test_random.py b/robot_nav/test_random.py
@@ -7,6 +7,7 @@
 from robot_nav.models.BPG.BTD3 import BTD3
 from robot_nav.models.CNNTD3.CNNTD3 import CNNTD3
 from robot_nav.models.SAC.BSAC import BSAC
+from robot_nav.models.SAC.BSA1C import BSA1C
 import statistics
 import numpy as np
 import tqdm
@@ -28,13 +29,13 @@ def main(args=None):
     max_steps = 300  # maximum number of steps in single episode
     test_scenarios = 1000
 
-    model = BSAC(
+    model = BSA1C(
         state_dim=state_dim,
         action_dim=action_dim,
         max_action=max_action,
         device=device,
         load_model=True,
-        model_name="BSACw025exp1",
+        model_name="BSA1Cw025exp1",
     )  # instantiate a model
 
     sim = SIM_ENV(
diff --git a/robot_nav/train.py b/robot_nav/train.py
@@ -5,6 +5,7 @@
 from robot_nav.models.BPG.BCNNPG import BCNNPG
 from robot_nav.models.SAC.SAC import SAC
 from robot_nav.models.SAC.BSAC import BSAC
+from robot_nav.models.SAC.BSA1C import BSA1C
 from robot_nav.models.HCM.hardcoded_model import HCM
 from robot_nav.models.PPO.PPO import PPO
 from robot_nav.models.CNNTD3.CNNTD3 import CNNTD3
@@ -40,14 +41,14 @@ def main(args=None):
     )
     save_every = 5  # save the model every n training cycles
 
-    model = BSAC(
+    model = BSA1C(
         state_dim=state_dim,
         action_dim=action_dim,
         max_action=max_action,
         device=device,
         save_every=save_every,
         load_model=False,
-        model_name="BSACw025exp1",
+        model_name="BSA1Cw025exp1",
         # bound_weight=0.0,
     )  # instantiate a model
 
