SAC：柔性 Actor-Critic

import torch
import torch.nn.functional as F

def compute_entropy(policy_logits: torch.Tensor) -> torch.Tensor:
    """计算策略分布的熵值"""
    policy_dist = torch.distributions.Categorical(logits=policy_logits)
    entropy = policy_dist.entropy()  # H = -sum(p * log p)
    return entropy.mean()

def max_ent_reward(rewards: torch.Tensor, 
                   policy_logits: torch.Tensor, 
                   alpha: float = 0.2) -> torch.Tensor:
    """最大熵回报：r + alpha * H"""
    entropy = compute_entropy(policy_logits)
    return rewards.mean() + alpha * entropy

class MaxEntropyValueIteration:
    """最大熵值迭代算法"""
    def __init__(self, n_states: int, n_actions: int, 
                 gamma: float = 0.99, alpha: float = 0.1):
        self.V = torch.zeros(n_states)
        self.gamma = gamma
        self.alpha = alpha
        
    def soft_value(self, q_values: torch.Tensor) -> torch.Tensor:
        """软值函数：log-sum-exp 形式"""
        # V_soft(s) = alpha * log(sum(exp(Q(s,a)/alpha)))
        return self.alpha * torch.logsumexp(q_values / self.alpha, dim=-1)
    
    def update(self, rewards: torch.Tensor, 
               transitions: torch.Tensor) -> None:
        """执行一次软值迭代"""
        q_values = rewards + self.gamma * transitions @ self.V
        self.V = self.soft_value(q_values)

import torch
import torch.nn as nn
import torch.nn.functional as F

class SoftQNetwork(nn.Module):
    """SAC 的软 Q 网络"""
    def __init__(self, state_dim: int, action_dim: int, 
                 hidden_dim: int = 256):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(state_dim + action_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, 1)
        )
    
    def forward(self, state: torch.Tensor, 
                action: torch.Tensor) -> torch.Tensor:
        x = torch.cat([state, action], dim=-1)
        return self.net(x)

class GaussianPolicy(nn.Module):
    """SAC 的高斯策略网络"""
    def __init__(self, state_dim: int, action_dim: int,
                 hidden_dim: int = 256, log_std_min: float = -20,
                 log_std_max: float = 2):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(state_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
        )
        self.mean_head = nn.Linear(hidden_dim, action_dim)
        self.log_std_head = nn.Linear(hidden_dim, action_dim)
        self.log_std_min = log_std_min
        self.log_std_max = log_std_max
    
    def forward(self, state: torch.Tensor):
        features = self.net(state)
        mean = self.mean_head(features)
        log_std = self.log_std_head(features)
        log_std = torch.clamp(log_std, self.log_std_min, self.log_std_max)
        return mean, log_std

def soft_q_target(rewards: torch.Tensor, 
                      next_states: torch.Tensor, 
                      dones: torch.Tensor,
                      next_actions: torch.Tensor, 
                      next_log_probs: torch.Tensor,
                      target_q1: torch.Tensor, 
                      target_q2: torch.Tensor,
                      gamma: float = 0.99, 
                      alpha: float = 0.2) -> torch.Tensor:
    """计算 SAC 软 Q 目标值"""
    # 取两个 Q 网络的较小值（缓解高估）
    min_q_next = torch.min(target_q1, target_q2)
    # 软价值：V = min(Q1, Q2) - alpha * log pi
    next_v = min_q_next - alpha * next_log_probs
    # Bellman 备份
    targets = rewards + gamma * (1 - dones) * next_v
    return targets

def sac_q_loss(q_values: torch.Tensor, 
               targets: torch.Tensor) -> torch.Tensor:
    """SAC Q 网络损失：MSE"""
    return F.mse_loss(q_values, targets.detach())

def soft_policy_eval(q1: SoftQNetwork, q2: SoftQNetwork,
                         policy: GaussianPolicy, 
                         states: torch.Tensor, 
                         actions: torch.Tensor,
                         rewards: torch.Tensor, 
                         next_states: torch.Tensor,
                         dones: torch.Tensor,
                         alpha: float, gamma: float = 0.99,
                         tau: float = 0.005) -> dict:
    """软策略评估：更新 Q 网络"""
    with torch.no_grad():
        # 采样下一动作
        next_mean, next_log_std = policy(next_states)
        next_std = next_log_std.exp()
        next_dist = torch.distributions.Normal(next_mean, next_std)
        next_action_raw = next_dist.rsample()
        next_action = torch.tanh(next_action_raw)
        
        # 计算熵修正
        next_log_prob = next_dist.log_prob(next_action_raw)
        next_log_prob = next_log_prob.sum(dim=-1, keepdim=True)
        # tanh 修正
        next_log_prob -= torch.log(1 - next_action**2 + 1e-6).sum(dim=-1, keepdim=True)
        
        next_q1 = q1(next_states, next_action)
        next_q2 = q2(next_states, next_action)
        next_v = torch.min(next_q1, next_q2) - alpha * next_log_prob
        
        q_targets = rewards + gamma * (1 - dones) * next_v
    
    q1_pred = q1(states, actions)
    q2_pred = q2(states, actions)
    q1_loss = F.mse_loss(q1_pred, q_targets.detach())
    q2_loss = F.mse_loss(q2_pred, q_targets.detach())
    return {"q1_loss": q1_loss, "q2_loss": q2_loss}

def soft_policy_improvement(q1: SoftQNetwork, 
                              q2: SoftQNetwork,
                              policy: GaussianPolicy, 
                              policy_optimizer: torch.optim.Optimizer,
                              states: torch.Tensor,
                              alpha: float, n_samples: int = 10) -> dict:
    """软策略改进：更新策略网络"""
    mean, log_std = policy(states)
    std = log_std.exp()
    dist = torch.distributions.Normal(mean, std)
    
    # 重参数化采样
    z = dist.rsample()
    action = torch.tanh(z)
    
    log_prob = dist.log_prob(z)
    log_prob = log_prob.sum(dim=-1, keepdim=True)
    log_prob -= torch.log(1 - action**2 + 1e-6).sum(dim=-1, keepdim=True)
    
    # 策略目标：max E[Q - alpha * log pi]
    q1_val = q1(states, action)
    q2_val = q2(states, action)
    q_val = torch.min(q1_val, q2_val)
    
    policy_loss = (alpha * log_prob - q_val).mean()
    
    policy_optimizer.zero_grad()
    policy_loss.backward()
    policy_optimizer.step()
    
    return {"policy_loss": policy_loss.item()}

class AutoTemperature:
    """SAC 自动温度调节器"""
    def __init__(self, initial_alpha: float = 0.2, 
                 target_entropy: float = None,
                 lr: float = 3e-4, action_dim: int = None):
        self.log_alpha = torch.tensor(
            [torch.log(torch.tensor(initial_alpha))],
            requires_grad=True
        )
        self.alpha_optimizer = torch.optim.Adam(
            [self.log_alpha], lr=lr
        )
        if target_entropy is None and action_dim is not None:
            # 默认目标熵 = -动作维度
            self.target_entropy = -action_dim
        else:
            self.target_entropy = target_entropy
    
    @property
    def alpha(self) -> torch.Tensor:
        return torch.exp(self.log_alpha.detach())
    
    def update(self, log_probs: torch.Tensor) -> float:
        """更新温度系数 alpha"""
        # 目标: alpha * (log pi + H_target) = 0
        alpha_loss = -(self.log_alpha * 
                      (log_probs.detach() + self.target_entropy)).mean()
        
        self.alpha_optimizer.zero_grad()
        alpha_loss.backward()
        self.alpha_optimizer.step()
        
        return alpha_loss.item()

def train_with_auto_alpha(replay_buffer, q1, q2, 
                           policy, alpha_controller,
                           q1_optimizer, q2_optimizer, 
                           policy_optimizer, 
                           batch_size: int = 256,
                           gamma: float = 0.99,
                           tau: float = 0.005,
                           steps: int = 100000) -> list:
    """使用自动温度调节的 SAC 训练循环"""
    losses = []
    
    for step in range(steps):
        batch = replay_buffer.sample(batch_size)
        states, actions, rewards, next_states, dones = batch
        
        # 1) 更新 Q 网络
        q_losses = soft_policy_eval(q1, q2, policy, 
                                     states, actions, rewards, 
                                     next_states, dones,
                                     alpha=alpha_controller.alpha,
                                     gamma=gamma, tau=tau)
        q1_optimizer.zero_grad()
        q_losses["q1_loss"].backward()
        q1_optimizer.step()
        q2_optimizer.zero_grad()
        q_losses["q2_loss"].backward()
        q2_optimizer.step()
        
        # 2) 更新策略
        policy_loss = soft_policy_improvement(
            q1, q2, policy, policy_optimizer,
            states, alpha=alpha_controller.alpha
        )
        
        # 3) 更新 alpha
        alpha_loss = alpha_controller.update(policy_loss["log_prob"])
        
        # 4) 软更新目标网络
        soft_target_update(q1, q1_target, tau)
        soft_target_update(q2, q2_target, tau)
        
        losses.append({
            "q1_loss": q_losses["q1_loss"].item(),
            "policy_loss": policy_loss["policy_loss"],
            "alpha": alpha_controller.alpha.item()
        })
    
    return losses

def ppo_update(policy, value_net, states, actions,
               old_log_probs, advantages, returns,
               clip_ratio: float = 0.2, 
               value_coef: float = 0.5,
               entropy_coef: float = 0.01) -> dict:
    """PPO 裁剪策略更新"""
    # 新策略分布
    mean, log_std = policy(states)
    dist = torch.distributions.Normal(mean, log_std.exp())
    new_log_probs = dist.log_prob(actions).sum(dim=-1)
    
    # 重要性采样比率
    ratio = torch.exp(new_log_probs - old_log_probs)
    
    # 裁剪目标
    surr1 = ratio * advantages
    surr2 = torch.clamp(ratio, 1 - clip_ratio, 
                        1 + clip_ratio) * advantages
    policy_loss = -torch.min(surr1, surr2).mean()
    
    # 值函数损失
    values = value_net(states).squeeze()
    value_loss = F.mse_loss(values, returns)
    
    # 熵正则
    entropy = dist.entropy().mean()
    
    total_loss = policy_loss + value_coef * value_loss - entropy_coef * entropy
    return {"total_loss": total_loss.item(), 
            "policy_loss": policy_loss.item(),
            "value_loss": value_loss.item()}

def compare_sac_ppo(env_name: str = "HalfCheetah-v4",
                    sac_steps: int = 1_000_000,
                    ppo_episodes: int = 1000) -> dict:
    """SAC vs PPO 对比实验"""
    import gymnasium as gym
    
    env = gym.make(env_name)
    results = {}
    
    # SAC: off-policy, 样本效率高
    sac_buffer = ReplayBuffer(capacity=1_000_000)
    # ... SAC 训练 ...
    sac_reward_sac = train_sac(env, sac_buffer, sac_steps)
    results["SAC"] = {
        "type": "off-policy",
        "samples": sac_steps,
        "final_reward": sac_reward_sac,
        "sample_efficiency": "high"
    }
    
    # PPO: on-policy, 每次迭代需要新数据
    # ... PPO 训练 ...
    ppo_reward_ppo = train_ppo(env, ppo_episodes)
    results["PPO"] = {
        "type": "on-policy",
        "samples": ppo_episodes * env.spec.max_episode_steps,
        "final_reward": ppo_reward_ppo,
        "sample_efficiency": "low"
    }
    
    return results

import gymnasium as gym
import torch
import numpy as np

def evaluate_sac(policy: GaussianPolicy, 
                 env_name: str = "HalfCheetah-v4",
                 n_episodes: int = 10) -> dict:
    """评估 SAC 策略在连续控制环境中的性能"""
    env = gym.make(env_name)
    returns = []
    episode_lengths = []
    
    for ep in range(n_episodes):
        state, _ = env.reset()
        ep_reward = 0.0
        done = False
        
        while not done:
            state_t = torch.tensor(state, dtype=torch.float32).unsqueeze(0)
            with torch.no_grad():
                mean, log_std = policy(state_t)
                action = torch.tanh(mean).numpy().flatten()
            
            state, reward, terminated, truncated, _ = env.step(action)
            ep_reward += reward
            done = terminated or truncated
        
        returns.append(ep_reward)
        episode_lengths.append(env.unwrapped.time if hasattr(env.unwrapped, 'time') else 0)
    
    return {
        "mean_return": np.mean(returns),
        "std_return": np.std(returns),
        "max_return": np.max(returns),
        "mean_length": np.mean(episode_lengths)
    }

class SACAgent:
    """完整的 SAC 智能体（连续控制）"""
    def __init__(self, state_dim: int, action_dim: int,
                 action_low: np.ndarray, action_high: np.ndarray,
                 hidden_dim: int = 256, lr: float = 3e-4,
                 buffer_size: int = 1_000_000,
                 batch_size: int = 256,
                 gamma: float = 0.99, tau: float = 0.005):
        self.state_dim = state_dim
        self.action_dim = action_dim
        self.action_scale = torch.tensor(
            (action_high - action_low) / 2
        )
        self.action_bias = torch.tensor(
            (action_high + action_low) / 2
        )
        
        self.policy = GaussianPolicy(state_dim, action_dim, hidden_dim)
        self.q1 = SoftQNetwork(state_dim, action_dim, hidden_dim)
        self.q2 = SoftQNetwork(state_dim, action_dim, hidden_dim)
        self.q1_target = SoftQNetwork(state_dim, action_dim, hidden_dim)
        self.q2_target = SoftQNetwork(state_dim, action_dim, hidden_dim)
        
        self.copy_weights(self.q1, self.q1_target)
        self.copy_weights(self.q2, self.q2_target)
        
        self.replay_buffer = ReplayBuffer(buffer_size)
        self.batch_size = batch_size
        self.gamma = gamma
        self.tau = tau
        
        self.policy_optimizer = torch.optim.Adam(
            self.policy.parameters(), lr=lr)
        self.q1_optimizer = torch.optim.Adam(
            self.q1.parameters(), lr=lr)
        self.q2_optimizer = torch.optim.Adam(
            self.q2.parameters(), lr=lr)
        
        self.alpha_controller = AutoTemperature(
            initial_alpha=0.2, action_dim=action_dim)
    
    def copy_weights(self, src: nn.Module, dst: nn.Module):
        for src_p, dst_p in zip(src.parameters(), dst.parameters()):
            dst_p.data.copy_(src_p.data)

import gymnasium as gym
import numpy as np
from collections import deque

class ReplayBuffer:
    """高效经验回放缓冲区"""
    def __init__(self, capacity: int = 1_000_000):
        self.buffer = deque(maxlen=capacity)
    
    def push(self, state, action, reward, 
             next_state, done):
        self.buffer.append((state, action, reward, 
                           next_state, done))
    
    def sample(self, batch_size: int) -> tuple:
        indices = np.random.choice(len(self.buffer), 
                                   batch_size, replace=False)
        batch = [self.buffer[i] for i in indices]
        states = np.array([b[0] for b in batch])
        actions = np.array([b[1] for b in batch])
        rewards = np.array([b[2] for b in batch])
        next_states = np.array([b[3] for b in batch])
        dones = np.array([b[4] for b in batch])
        return (torch.tensor(states, dtype=torch.float32),
                torch.tensor(actions, dtype=torch.float32),
                torch.tensor(rewards, dtype=torch.float32).unsqueeze(1),
                torch.tensor(next_states, dtype=torch.float32),
                torch.tensor(dones, dtype=torch.float32).unsqueeze(1))
    
    def __len__(self):
        return len(self.buffer)

def train_sac_pendulum(total_steps: int = 100_000,
                       batch_size: int = 256,
                       start_steps: int = 10_000) -> list:
    """在 Pendulum-v1 上训练 SAC"""
    env = gym.make("Pendulum-v1")
    state_dim = env.observation_space.shape[0]
    action_dim = env.action_space.shape[0]
    
    agent = SACAgent(state_dim, action_dim,
                     env.action_space.low,
                     env.action_space.high)
    
    state, _ = env.reset()
    episode_reward = 0
    training_logs = []
    
    for step in range(total_steps):
        if step < start_steps:
            action = env.action_space.sample()
        else:
            action = agent.select_action(state, deterministic=False)
        
        next_state, reward, terminated, truncated, _ = env.step(action)
        done = terminated or truncated
        agent.replay_buffer.push(state, action, reward, next_state, done)
        state = next_state
        episode_reward += reward
        
        if step >= start_steps and len(agent.replay_buffer) > batch_size:
            for _ in range(1):
                agent.update(batch_size)
        
        if done:
            state, _ = env.reset()
            training_logs.append({"step": step, "reward": episode_reward})
            episode_reward = 0
    
    env.close()
    return training_logs

def train_sac_mujoco(env_name: str = "HalfCheetah-v4",
                       total_steps: int = 1_000_000,
                       batch_size: int = 256,
                       start_steps: int = 10_000,
                       eval_interval: int = 10_000,
                       eval_episodes: int = 5) -> dict:
    """在 MuJoCo 环境上训练 SAC"""
    import gymnasium as gym
    
    env = gym.make(env_name)
    eval_env = gym.make(env_name)
    state_dim = env.observation_space.shape[0]
    action_dim = env.action_space.shape[0]
    
    agent = SACAgent(
        state_dim=state_dim,
        action_dim=action_dim,
        action_low=env.action_space.low,
        action_high=env.action_space.high,
        hidden_dim=256,
        lr=3e-4,
        gamma=0.99,
        tau=0.005,
        buffer_size=1_000_000,
        batch_size=batch_size
    )
    
    state, _ = env.reset()
    episode_reward = 0
    results = {"steps": [], "train_rewards": [], "eval_rewards": []}
    
    for step in range(total_steps):
        # 选择动作
        if step < start_steps:
            action = env.action_space.sample()
        else:
            action = agent.select_action(state, deterministic=False)
        
        next_state, reward, terminated, truncated, _ = env.step(action)
        done = terminated or truncated
        agent.replay_buffer.push(state, action, reward, next_state, done)
        state = next_state
        episode_reward += reward
        
        # 训练
        if step >= start_steps and len(agent.replay_buffer) >= batch_size:
            agent.update(batch_size)
        
        # episode 结束
        if done:
            state, _ = env.reset()
            results["train_rewards"].append(episode_reward)
            results["steps"].append(step)
            episode_reward = 0
        
        # 定期评估
        if (step + 1) % eval_interval == 0:
            eval_reward = evaluate_sac(
                agent.policy, env_name, eval_episodes
            )
            results["eval_rewards"].append(eval_reward)
            print(f"Step {step}: eval_return = {eval_reward['mean_return']:.1f}")
    
    env.close()
    eval_env.close()
    return results