# === 数值类型 ===
reward = 1.5            # float: 奖励值通常是浮点数
episode = 100            # int: 回合数是整数
done = True              # bool: 回合是否结束

# === 字符串 ===
env_name = "CartPole-v1"  # str: 环境名称
print(f"环境: {env_name}, 回合: {episode}")  # f-string 格式化

# === 列表 (List) - 有序可变序列 ===
rewards = [0.5, 1.0, -0.3, 2.0]   # 存储每步的奖励
rewards.append(1.5)                   # 添加新奖励
total = sum(rewards)                   # 计算总奖励

# === 字典 (Dict) - 键值对映射 ===
config = {
    "learning_rate": 3e-4,    # 学习率
    "gamma": 0.99,             # 折扣因子
    "batch_size": 64,          # 批大小
    "hidden_dim": 256,         # 隐藏层维度
}

# === 元组 (Tuple) - 有序不可变序列 ===
state_shape = (4,)          # 状态空间维度
action_range = (-1.0, 1.0)  # 动作取值范围

# === 类型转换 ===
steps_str = "1000"
steps_int = int(steps_str)           # 字符串 → 整数
ratio = float(steps_int) / 500      # 整数 → 浮点数

# === 条件判断 ===
reward = -0.5

if reward > 0:
    print("正奖励，智能体做了好的决策")
elif reward == 0:
    print("零奖励，中性结果")
else:
    print("负奖励，智能体需要改进策略")

# === for 循环 - 训练回合 ===
num_episodes = 5
for ep in range(num_episodes):
    episode_reward = 0
    # 模拟一个回合内的多个步骤
    for step in range(100):
        r = 1.0  # 假设每步奖励为 1
        episode_reward += r
    print(f"回合 {ep}: 总奖励 = {episode_reward}")

# === while 循环 + break ===
step_count = 0
max_steps = 200
while True:
    step_count += 1
    done = step_count >= max_steps
    if done:
        print(f"回合在第 {step_count} 步结束")
        break

# === continue - 跳过无效数据 ===
rewards_list = [1.0, None, 0.5, None, 2.0]
valid_rewards = []
for r in rewards_list:
    if r is None:
        continue  # 跳过空值
    valid_rewards.append(r)

# === 基本函数 ===
def compute_return(rewards, gamma=0.99):
    """计算折扣回报 G_t = r_t + gamma * r_{t+1} + gamma^2 * r_{t+2} + ..."""
    g = 0
    returns = []
    for r in reversed(rewards):
        g = r + gamma * g
        returns.insert(0, g)
    return returns

rewards = [1.0, 0.0, 1.0, 0.5]
G = compute_return(rewards, gamma=0.99)
print(f"折扣回报: {G}")

# === *args 和 **kwargs ===
def create_env(env_id, *wrappers, **kwargs):
    """灵活创建环境，支持任意多个包装器和配置参数"""
    print(f"创建环境: {env_id}")
    for w in wrappers:
        print(f"  应用包装器: {w}")
    for k, v in kwargs.items():
        print(f"  配置: {k} = {v}")

create_env("Hopper-v4", "NormalizeObservation", "ClipAction",
           max_episode_steps=1000, render_mode="human")

# === lambda 表达式 ===
# 常用于简短的奖励变换或回调函数
clip_reward = lambda r: max(-1.0, min(1.0, r))
print(clip_reward(5.0))   # 输出 1.0
print(clip_reward(-3.0))  # 输出 -1.0

# === 列表推导式 ===
# 快速生成数据，在 RL 中用于批量处理非常常见
squared = [x ** 2 for x in range(10)]
positive_rewards = [r for r in rewards if r > 0]

# === 生成器 - 节省内存 ===
def episode_generator(num_episodes):
    """逐个生成回合数据，无需一次性加载所有数据到内存"""
    for i in range(num_episodes):
        # 模拟生成一个回合的奖励序列
        ep_rewards = [1.0] * 50
        yield i, ep_rewards

for ep_id, ep_r in episode_generator(3):
    print(f"回合 {ep_id}: 奖励步数 = {len(ep_r)}")

import random

class ReplayBuffer:
    """经验回放缓冲区 - 存储并采样训练数据"""

    def __init__(self, capacity):
        self.capacity = capacity   # 最大容量
        self.buffer = []            # 存储列表
        self.position = 0          # 当前写入位置

    def push(self, state, action, reward, next_state, done):
        """存储一条经验"""
        experience = (state, action, reward, next_state, done)
        if len(self.buffer) < self.capacity:
            self.buffer.append(experience)
        else:
            self.buffer[self.position] = experience
        self.position = (self.position + 1) % self.capacity

    def sample(self, batch_size):
        """随机采样一批经验"""
        return random.sample(self.buffer, batch_size)

    def __len__(self):
        return len(self.buffer)

# === 继承 - 创建特殊版本的缓冲区 ===
class PrioritizedReplayBuffer(ReplayBuffer):
    """优先经验回放 - 按 TD 误差优先级采样"""

    def __init__(self, capacity, alpha=0.6):
        super().__init__(capacity)  # 调用父类初始化
        self.alpha = alpha           # 优先级指数
        self.priorities = []          # 优先级列表

    def push(self, state, action, reward, next_state, done, td_error=1.0):
        super().push(state, action, reward, next_state, done)
        priority = abs(td_error) ** self.alpha
        if len(self.priorities) < self.capacity:
            self.priorities.append(priority)
        else:
            self.priorities[self.position - 1] = priority

# 使用示例
buf = ReplayBuffer(capacity=10000)
buf.push([1.0, 0.5], 0, 1.0, [1.1, 0.4], False)
print(f"缓冲区大小: {len(buf)}")

import numpy as np

# === 创建数组 ===
state = np.array([0.5, -0.3, 1.2, 0.0])       # 从列表创建
zeros = np.zeros((3, 4))                       # 3x4 零矩阵
ones = np.ones(5)                               # 全 1 向量
random_actions = np.random.randn(10, 2)         # 10 个 2 维随机动作
indices = np.arange(0, 100, 10)                # [0, 10, 20, ..., 90]

# === 索引与切片 ===
rewards = np.array([1.0, 0.5, -0.3, 2.0, 0.8, -1.0])
print(rewards[0])        # 第一个: 1.0
print(rewards[-1])       # 最后一个: -1.0
print(rewards[1:4])      # 切片: [0.5, -0.3, 2.0]
print(rewards[rewards > 0])  # 布尔索引: [1.0, 0.5, 2.0, 0.8]

# === 广播 (Broadcasting) ===
# 标量与数组运算：对每个元素应用
normalized = (rewards - rewards.mean()) / (rewards.std() + 1e-8)
print(f"归一化后: {normalized}")

# === 矩阵运算 ===
# 模拟一个简单的线性策略: action = W @ state + b
W = np.random.randn(2, 4)    # 权重矩阵 (2 个动作, 4 维状态)
b = np.zeros(2)               # 偏置
action = W @ state + b         # 矩阵乘法 + 广播加法
print(f"动作: {action}")

# === 常用统计 ===
print(f"平均奖励: {rewards.mean()}")
print(f"最大奖励: {rewards.max()}")
print(f"标准差:   {rewards.std()}")
print(f"形状:     {random_actions.shape}")  # (10, 2)

import numpy as np                 # 数值计算
import torch                          # PyTorch 深度学习框架
import torch.nn as nn               # 神经网络模块
import torch.optim as optim          # 优化器
from torch.distributions import Normal  # 概率分布（用于随机策略）
import gymnasium as gym              # RL 环境接口
from collections import deque         # 高效的双端队列
import matplotlib.pyplot as plt      # 绘图

import torch
import numpy as np

# === 创建张量 ===
# 从 Python 列表创建
state = torch.tensor([0.5, -0.3, 1.2, 0.0])
print(f"状态张量: {state}")
print(f"形状: {state.shape}, 数据类型: {state.dtype}")

# 常用创建函数
zeros = torch.zeros(3, 4)           # 3x4 零矩阵
ones = torch.ones(2, 3)              # 2x3 全 1 矩阵
rand = torch.rand(5, 2)              # 5x2 均匀分布 [0,1)
randn = torch.randn(64, 4)           # 64x4 标准正态分布（常用于初始化）
arange = torch.arange(0, 10)         # [0, 1, 2, ..., 9]
eye = torch.eye(3)                    # 3x3 单位矩阵

# === 指定数据类型 ===
int_tensor = torch.tensor([1, 0, 1], dtype=torch.long)      # 离散动作索引
float_tensor = torch.tensor([0.5], dtype=torch.float32)     # 连续值
bool_tensor = torch.tensor([True, False])                    # 掩码

# === 设备管理 (CPU / GPU) ===
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"当前设备: {device}")
gpu_tensor = torch.randn(3, 4).to(device)    # 将张量移动到指定设备

# === NumPy 互转 ===
np_array = np.array([1.0, 2.0, 3.0])
from_np = torch.from_numpy(np_array)         # NumPy → Tensor（共享内存）
back_to_np = from_np.numpy()                 # Tensor → NumPy

import torch

# === reshape 和 view ===
# 假设我们有一批展平的状态向量，需要恢复为 (batch, channels, height, width)
flat = torch.randn(8, 84 * 84 * 3)        # 8 张展平的 84x84 RGB 图像
images = flat.view(8, 3, 84, 84)            # 恢复为 4D 张量
print(f"恢复后形状: {images.shape}")        # torch.Size([8, 3, 84, 84])

# === squeeze 和 unsqueeze ===
# squeeze: 移除大小为 1 的维度
action = torch.tensor([[0.5]])               # 形状 (1, 1)
print(action.squeeze().shape)                  # 形状 () - 变成标量

# unsqueeze: 添加大小为 1 的维度
state = torch.randn(4)                        # 形状 (4,)
batched = state.unsqueeze(0)                   # 形状 (1, 4) - 添加 batch 维度
print(f"添加 batch 维度: {batched.shape}")    # torch.Size([1, 4])

# === cat 和 stack ===
# cat: 沿现有维度拼接
states_a = torch.randn(32, 4)    # 第一批 32 个状态
states_b = torch.randn(16, 4)    # 第二批 16 个状态
combined = torch.cat([states_a, states_b], dim=0)
print(f"拼接后: {combined.shape}")   # torch.Size([48, 4])

# stack: 沿新维度堆叠
episode_returns = [torch.tensor(100.0), torch.tensor(150.0), torch.tensor(120.0)]
stacked = torch.stack(episode_returns)
print(f"堆叠后: {stacked}")             # tensor([100., 150., 120.])

# === 高级索引 ===
q_values = torch.randn(4, 3)       # 4 个状态，3 个动作的 Q 值
actions = torch.tensor([0, 2, 1, 0])  # 每个状态选择的动作索引
# 取出每个状态对应动作的 Q 值
selected_q = q_values.gather(1, actions.unsqueeze(1)).squeeze(1)
print(f"选中的 Q 值: {selected_q}")

import torch

# === 基本梯度计算 ===
# requires_grad=True 告诉 PyTorch 追踪这个张量的所有运算
w = torch.tensor([2.0, 3.0], requires_grad=True)
x = torch.tensor([1.0, 1.0])

# 前向传播: y = w · x
y = (w * x).sum()          # y = 2*1 + 3*1 = 5

# 反向传播: 计算 dy/dw
y.backward()
print(f"梯度 dy/dw: {w.grad}")  # tensor([1., 1.]) — 因为 dy/dw_i = x_i

# === 计算图与 detach ===
a = torch.tensor([1.0], requires_grad=True)
b = a * 2          # b 在计算图中
c = b.detach()     # c 脱离计算图，梯度不再回传到 a
print(f"b 需要梯度: {b.requires_grad}")   # True
print(f"c 需要梯度: {c.requires_grad}")   # False

# === torch.no_grad() — 推理时关闭梯度追踪 ===
# 在 RL 中，收集经验时不需要计算梯度，这样可以节省内存和加速
model_params = torch.randn(4, 2, requires_grad=True)
state = torch.randn(4)

with torch.no_grad():
    # 这个块里的运算不会建立计算图
    action = state @ model_params
    print(f"推理动作: {action}")
    print(f"需要梯度: {action.requires_grad}")  # False

import torch
import torch.nn as nn

# === 方式一：使用 Sequential 快速搭建 ===
simple_policy = nn.Sequential(
    nn.Linear(4, 128),      # 输入层: 4 维状态 → 128 维隐藏层
    nn.ReLU(),               # 激活函数
    nn.Linear(128, 128),    # 隐藏层
    nn.ReLU(),
    nn.Linear(128, 2),      # 输出层: 2 维动作
    nn.Tanh(),               # 将动作限制在 [-1, 1]
)

# 测试前向传播
dummy_state = torch.randn(1, 4)    # batch_size=1, state_dim=4
action = simple_policy(dummy_state)
print(f"策略输出动作: {action}")

# === 方式二：自定义 Module（更灵活，推荐用于 RL）===
class PolicyNetwork(nn.Module):
    """高斯策略网络 - 输出动作的均值和标准差"""

    def __init__(self, state_dim, action_dim, hidden_dim=256):
        super().__init__()

        # 共享特征提取层
        self.shared = 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 = nn.Parameter(torch.zeros(action_dim))

    def forward(self, state):
        """前向传播: 状态 → 动作分布参数"""
        features = self.shared(state)
        mean = self.mean_head(features)
        std = self.log_std.exp()         # 确保标准差为正数
        return mean, std

# 实例化并测试
policy = PolicyNetwork(state_dim=4, action_dim=2)
state = torch.randn(8, 4)   # 一批 8 个状态
mean, std = policy(state)
print(f"均值形状: {mean.shape}")    # torch.Size([8, 2])
print(f"标准差: {std}")

# 查看网络参数数量
total_params = sum(p.numel() for p in policy.parameters())
print(f"总参数量: {total_params}")

import torch
import torch.nn as nn
import torch.optim as optim

# === 常用损失函数 ===
mse_loss = nn.MSELoss()    # 均方误差 — 用于价值函数回归
ce_loss = nn.CrossEntropyLoss()  # 交叉熵 — 用于离散动作分类

# MSE 示例: 预测 Q 值 vs 目标 Q 值
predicted_q = torch.tensor([1.5, 2.3, 0.8])
target_q = torch.tensor([1.0, 2.0, 1.0])
loss = mse_loss(predicted_q, target_q)
print(f"MSE 损失: {loss.item():.4f}")

# === 自定义损失 — 策略梯度损失 ===
def policy_gradient_loss(log_probs, advantages):
    """策略梯度损失: L = -E[log π(a|s) · A(s,a)]"""
    return -(log_probs * advantages).mean()

# === 优化器与训练步骤 ===
# 构建一个简单的 Q 网络
q_network = nn.Sequential(
    nn.Linear(4, 64),
    nn.ReLU(),
    nn.Linear(64, 2),  # 2 个动作的 Q 值
)

# Adam 优化器 — RL 中最常用的优化器
optimizer = optim.Adam(q_network.parameters(), lr=3e-4)

# === 标准训练三步曲 ===
# 1. 清空梯度
optimizer.zero_grad()

# 2. 前向 + 计算损失 + 反向传播
states = torch.randn(32, 4)
q_pred = q_network(states)[:, 0]  # 取动作 0 的 Q 值
q_target = torch.randn(32)        # 目标 Q 值（实际中用贝尔曼方程计算）
loss = mse_loss(q_pred, q_target)
loss.backward()                    # 计算梯度

# 3. 更新参数
optimizer.step()
print(f"损失: {loss.item():.4f}")

import torch
import torch.nn as nn
import torch.optim as optim

# ── 1. 定义网络 ──
class SimpleMLP(nn.Module):
    """简单多层感知机，可作为策略网络或价值网络的基础结构"""
    def __init__(self, input_dim, output_dim, hidden_dim=128):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, output_dim),
        )

    def forward(self, x):
        return self.net(x)

# ── 2. 准备数据（模拟"状态→Q值"映射）──
torch.manual_seed(42)
num_samples = 1000
state_dim = 4
action_dim = 2

# 生成假数据：模拟 Q 值标签
X = torch.randn(num_samples, state_dim)
# 假设真实 Q 值是状态特征的某种线性组合 + 噪声
W_true = torch.randn(state_dim, action_dim)
Y = X @ W_true + 0.1 * torch.randn(num_samples, action_dim)

# ── 3. 初始化模型和优化器 ──
model = SimpleMLP(input_dim=state_dim, output_dim=action_dim)
optimizer = optim.Adam(model.parameters(), lr=1e-3)
loss_fn = nn.MSELoss()

# ── 4. 训练循环 ──
num_epochs = 200
batch_size = 64

for epoch in range(num_epochs):
    # 随机采样一个 mini-batch
    indices = torch.randint(0, num_samples, (batch_size,))
    batch_x = X[indices]
    batch_y = Y[indices]

    # 前向传播
    predictions = model(batch_x)

    # 计算损失
    loss = loss_fn(predictions, batch_y)

    # 反向传播与参数更新
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    # 每 50 个 epoch 打印一次
    if (epoch + 1) % 50 == 0:
        print(f"Epoch {epoch+1}/{num_epochs}, 损失: {loss.item():.4f}")

# ── 5. 推理（模拟在环境中使用策略）──
model.eval()  # 切换到评估模式
with torch.no_grad():
    test_state = torch.randn(1, state_dim)
    q_values = model(test_state)
    best_action = q_values.argmax(dim=1)
    print(f"测试状态的 Q 值: {q_values}")
    print(f"选择动作: {best_action.item()}")

import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np

class ActorCritic(nn.Module):
    """Actor-Critic 网络：策略和价值共享底层特征"""
    def __init__(self, obs_dim, act_dim, hidden=256):
        super().__init__()
        # 共享特征提取层
        self.shared = nn.Sequential(
            nn.Linear(obs_dim, hidden), nn.Tanh(),
            nn.Linear(hidden, hidden), nn.Tanh(),
        )
        # Actor 头：输出动作均值
        self.actor_mean = nn.Linear(hidden, act_dim)
        # 可学习的对数标准差
        self.actor_log_std = nn.Parameter(torch.zeros(act_dim))
        # Critic 头：输出状态价值
        self.critic = nn.Linear(hidden, 1)

    def forward(self, obs):
        features = self.shared(obs)
        return self.actor_mean(features), self.critic(features)

    def get_action(self, obs):
        """采样动作并返回 log_prob 和 value"""
        mean, value = self.forward(obs)
        std = self.actor_log_std.exp()
        dist = torch.distributions.Normal(mean, std)
        action = dist.sample()
        log_prob = dist.log_prob(action).sum(-1)
        return action, log_prob, value.squeeze(-1)

    def evaluate(self, obs, actions):
        """计算给定 obs-action 对的 log_prob, value, entropy"""
        mean, value = self.forward(obs)
        std = self.actor_log_std.exp()
        dist = torch.distributions.Normal(mean, std)
        log_prob = dist.log_prob(actions).sum(-1)
        entropy = dist.entropy().sum(-1)
        return log_prob, value.squeeze(-1), entropy


def compute_gae(rewards, values, dones, gamma=0.99, lam=0.95):
    """计算广义优势估计 (GAE)"""
    advantages = []
    gae = 0
    # 从后往前计算
    for t in reversed(range(len(rewards))):
        if t == len(rewards) - 1:
            next_value = 0  # 终止状态
        else:
            next_value = values[t + 1]
        delta = rewards[t] + gamma * next_value * (1 - dones[t]) - values[t]
        gae = delta + gamma * lam * (1 - dones[t]) * gae
        advantages.insert(0, gae)
    return torch.tensor(advantages, dtype=torch.float32)


def ppo_update(model, optimizer, obs, actions, old_log_probs,
               returns, advantages, clip_eps=0.2, epochs=10,
               batch_size=64, ent_coef=0.01, vf_coef=0.5):
    """PPO 策略更新"""
    dataset_size = obs.shape[0]
    for _ in range(epochs):
        # 随机打乱数据
        indices = np.random.permutation(dataset_size)
        for start in range(0, dataset_size, batch_size):
            end = start + batch_size
            idx = indices[start:end]

            # 取 mini-batch
            mb_obs = obs[idx]
            mb_actions = actions[idx]
            mb_old_log = old_log_probs[idx]
            mb_returns = returns[idx]
            mb_adv = advantages[idx]

            # 优势标准化（关键trick！）
            mb_adv = (mb_adv - mb_adv.mean()) / (mb_adv.std() + 1e-8)

            # 计算新策略下的 log_prob
            new_log_prob, values, entropy = model.evaluate(mb_obs, mb_actions)

            # 概率比 r_t(θ)
            ratio = (new_log_prob - mb_old_log).exp()

            # PPO-Clip 目标
            surr1 = ratio * mb_adv
            surr2 = torch.clamp(ratio, 1 - clip_eps, 1 + clip_eps) * mb_adv
            policy_loss = -torch.min(surr1, surr2).mean()

            # 价值函数损失
            value_loss = nn.functional.mse_loss(values, mb_returns)

            # 总损失 = 策略损失 + 价值损失 - 熵奖励
            loss = policy_loss + vf_coef * value_loss - ent_coef * entropy.mean()

            optimizer.zero_grad()
            loss.backward()
            nn.utils.clip_grad_norm_(model.parameters(), 0.5)  # 梯度裁剪
            optimizer.step()

import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
from collections import deque
import random

# ---- 经验回放缓冲区 ----
class ReplayBuffer:
    def __init__(self, capacity=1_000_000):
        self.buffer = deque(maxlen=capacity)

    def add(self, s, a, r, s_next, done):
        self.buffer.append((s, a, r, s_next, done))

    def sample(self, batch_size=256):
        batch = random.sample(self.buffer, batch_size)
        s, a, r, s2, d = zip(*batch)
        return (torch.FloatTensor(np.array(s)),
                torch.FloatTensor(np.array(a)),
                torch.FloatTensor(r).unsqueeze(1),
                torch.FloatTensor(np.array(s2)),
                torch.FloatTensor(d).unsqueeze(1))

# ---- Q 网络 ----
class QNetwork(nn.Module):
    def __init__(self, obs_dim, act_dim, hidden=256):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(obs_dim + act_dim, hidden), nn.ReLU(),
            nn.Linear(hidden, hidden), nn.ReLU(),
            nn.Linear(hidden, 1))

    def forward(self, obs, action):
        return self.net(torch.cat([obs, action], dim=-1))

# ---- Squashed Gaussian 策略 ----
class GaussianPolicy(nn.Module):
    def __init__(self, obs_dim, act_dim, hidden=256):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(obs_dim, hidden), nn.ReLU(),
            nn.Linear(hidden, hidden), nn.ReLU())
        self.mean_head = nn.Linear(hidden, act_dim)
        self.log_std_head = nn.Linear(hidden, act_dim)

    def forward(self, obs):
        feat = self.net(obs)
        mean = self.mean_head(feat)
        log_std = self.log_std_head(feat).clamp(-20, 2)
        return mean, log_std

    def sample(self, obs):
        """重参数化采样 + tanh 压缩"""
        mean, log_std = self.forward(obs)
        std = log_std.exp()
        dist = torch.distributions.Normal(mean, std)
        # 重参数化: u = μ + σ·ε
        u = dist.rsample()
        # tanh 压缩到 [-1, 1]
        action = torch.tanh(u)
        # 计算 log_prob（含雅可比修正）
        log_prob = dist.log_prob(u) - torch.log(1 - action.pow(2) + 1e-6)
        log_prob = log_prob.sum(-1, keepdim=True)
        return action, log_prob

# ---- SAC Agent ----
class SACAgent:
    def __init__(self, obs_dim, act_dim, lr=3e-4, gamma=0.99, tau=0.005):
        self.gamma = gamma
        self.tau = tau

        # 网络
        self.policy = GaussianPolicy(obs_dim, act_dim)
        self.q1 = QNetwork(obs_dim, act_dim)
        self.q2 = QNetwork(obs_dim, act_dim)
        self.q1_target = QNetwork(obs_dim, act_dim)
        self.q2_target = QNetwork(obs_dim, act_dim)
        # 初始化 target = 原网络
        self.q1_target.load_state_dict(self.q1.state_dict())
        self.q2_target.load_state_dict(self.q2.state_dict())

        # 自动温度调节
        self.target_entropy = -act_dim
        self.log_alpha = torch.zeros(1, requires_grad=True)

        # 优化器
        self.policy_opt = optim.Adam(self.policy.parameters(), lr=lr)
        self.q1_opt = optim.Adam(self.q1.parameters(), lr=lr)
        self.q2_opt = optim.Adam(self.q2.parameters(), lr=lr)
        self.alpha_opt = optim.Adam([self.log_alpha], lr=lr)

    @property
    def alpha(self):
        return self.log_alpha.exp().detach()

    def select_action(self, obs):
        obs_t = torch.FloatTensor(obs).unsqueeze(0)
        with torch.no_grad():
            action, _ = self.policy.sample(obs_t)
        return action.squeeze(0).numpy()

    def update(self, batch):
        s, a, r, s2, done = batch

        # ---- 更新 Q 网络 ----
        with torch.no_grad():
            a2, log_prob2 = self.policy.sample(s2)
            q1_target = self.q1_target(s2, a2)
            q2_target = self.q2_target(s2, a2)
            q_target = torch.min(q1_target, q2_target) - self.alpha * log_prob2
            y = r + self.gamma * (1 - done) * q_target

        q1_loss = nn.functional.mse_loss(self.q1(s, a), y)
        q2_loss = nn.functional.mse_loss(self.q2(s, a), y)
        self.q1_opt.zero_grad(); q1_loss.backward(); self.q1_opt.step()
        self.q2_opt.zero_grad(); q2_loss.backward(); self.q2_opt.step()

        # ---- 更新策略 ----
        a_new, log_prob = self.policy.sample(s)
        q_new = torch.min(self.q1(s, a_new), self.q2(s, a_new))
        policy_loss = (self.alpha * log_prob - q_new).mean()
        self.policy_opt.zero_grad(); policy_loss.backward(); self.policy_opt.step()

        # ---- 更新温度 α ----
        alpha_loss = -(self.log_alpha * (log_prob.detach() + self.target_entropy)).mean()
        self.alpha_opt.zero_grad(); alpha_loss.backward(); self.alpha_opt.step()

        # ---- 软更新 target 网络 ----
        for p, pt in zip(self.q1.parameters(), self.q1_target.parameters()):
            pt.data.copy_(self.tau * p.data + (1 - self.tau) * pt.data)
        for p, pt in zip(self.q2.parameters(), self.q2_target.parameters()):
            pt.data.copy_(self.tau * p.data + (1 - self.tau) * pt.data)

import numpy as np

def rotation_z(theta):
    """绕 Z 轴旋转的旋转矩阵"""
    c, s = np.cos(theta), np.sin(theta)
    return np.array([
        [c, -s, 0],
        [s,  c, 0],
        [0,  0, 1]
    ])

def homogeneous_transform(R, p):
    """构造 4×4 齐次变换矩阵"""
    T = np.eye(4)
    T[:3, :3] = R
    T[:3, 3] = p
    return T

# 示例：绕 Z 轴旋转 90 度，平移 (1, 0, 0.5)
R = rotation_z(np.pi / 2)
p = np.array([1.0, 0.0, 0.5])
T = homogeneous_transform(R, p)
print("齐次变换矩阵 T：")
print(np.round(T, 3))

import numpy as np

def forward_kinematics_2link(theta1, theta2, L1=1.0, L2=0.8):
    """二连杆机械臂正运动学
    参数:
        theta1: 关节1角度 (rad)
        theta2: 关节2角度 (rad)
        L1: 连杆1长度
        L2: 连杆2长度
    返回:
        (x, y) 末端执行器位置
    """
    # 关节1位置（肘部）
    x1 = L1 * np.cos(theta1)
    y1 = L1 * np.sin(theta1)

    # 末端执行器位置
    x2 = x1 + L2 * np.cos(theta1 + theta2)
    y2 = y1 + L2 * np.sin(theta1 + theta2)

    return x2, y2

# 测试不同关节角度
angles = [
    (0, 0),           # 完全伸展
    (np.pi/4, 0),     # 第一关节旋转45度
    (np.pi/4, -np.pi/2), # 肘部弯曲90度
]

for t1, t2 in angles:
    x, y = forward_kinematics_2link(t1, t2)
    print(f"θ1={np.degrees(t1):.0f}°, θ2={np.degrees(t2):.0f}° → x={x:.3f}, y={y:.3f}")

import numpy as np

class PDController:
    """简单的 PD 控制器"""

    def __init__(self, kp, kd, num_joints):
        self.kp = np.array(kp)  # 比例增益
        self.kd = np.array(kd)  # 微分增益
        self.num_joints = num_joints

    def compute_torque(self, q_desired, q_current, dq_desired, dq_current):
        """计算 PD 控制力矩
        参数:
            q_desired: 目标关节角度
            q_current: 当前关节角度
            dq_desired: 目标关节速度（通常为0）
            dq_current: 当前关节速度
        返回:
            tau: 关节力矩
        """
        pos_error = q_desired - q_current
        vel_error = dq_desired - dq_current
        tau = self.kp * pos_error + self.kd * vel_error
        return tau

# 示例：6自由度机械臂的 PD 控制器
controller = PDController(
    kp=[100, 100, 80, 50, 50, 30],  # 各关节刚度
    kd=[10, 10, 8, 5, 5, 3],         # 各关节阻尼
    num_joints=6
)

# 当前状态
q_current = np.array([0.5, -0.3, 0.8, 0.1, -0.2, 0.0])
dq_current = np.array([0.01, -0.02, 0.0, 0.03, 0.0, -0.01])

# 目标位置
q_desired = np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0])
dq_desired = np.zeros(6)

tau = controller.compute_torque(q_desired, q_current, dq_desired, dq_current)
print(f"控制力矩: {tau}")

<?xml version="1.0"?>
<robot name="two_link_arm">
  <!-- 基座（固定在世界中） -->
  <link name="base_link">
    <visual>
      <geometry><cylinder radius="0.05" length="0.1"/></geometry>
    </visual>
    <inertial>
      <mass value="5.0"/>
      <inertia ixx="0.01" ixy="0" ixz="0" iyy="0.01" iyz="0" izz="0.01"/>
    </inertial>
  </link>

  <!-- 连杆1：长0.5m -->
  <link name="link1">
    <visual>
      <origin xyz="0 0 0.25"/>
      <geometry><cylinder radius="0.03" length="0.5"/></geometry>
    </visual>
    <inertial>
      <mass value="2.0"/>
      <origin xyz="0 0 0.25"/>
      <inertia ixx="0.042" ixy="0" ixz="0" iyy="0.042" iyz="0" izz="0.001"/>
    </inertial>
  </link>

  <!-- 肩关节：绕 Z 轴旋转 -->
  <joint name="shoulder" type="revolute">
    <parent link="base_link"/>
    <child link="link1"/>
    <axis xyz="0 0 1"/>
    <limit lower="-3.14" upper="3.14" effort="50" velocity="3.0"/>
  </joint>

  <!-- 连杆2 和 肘关节 省略，结构类似 -->
</robot>

<mujoco model="two_link_arm">
  <option timestep="0.002" gravity="0 0 -9.81"/>

  <worldbody>
    <light pos="0 0 3"/>
    <geom type="plane" size="5 5 0.1"/>

    <!-- 基座固定在世界 -->
    <body name="base" pos="0 0 1">
      <geom type="cylinder" size="0.05 0.05" rgba="0.5 0.5 0.5 1"/>

      <!-- 连杆1: body 嵌套代替 parent/child -->
      <body name="link1">
        <joint name="shoulder" type="hinge" axis="0 0 1"
               range="-3.14 3.14"/>
        <geom type="capsule" fromto="0 0 0 0 0 0.5"
              size="0.03" mass="2.0"/>

        <body name="link2" pos="0 0 0.5">
          <joint name="elbow" type="hinge" axis="0 1 0"
                 range="-2.0 2.0"/>
          <geom type="capsule" fromto="0 0 0 0 0 0.4"
                size="0.025" mass="1.5"/>
        </body>
      </body>
    </body>
  </worldbody>

  <!-- 执行器（URDF 中没有！） -->
  <actuator>
    <motor joint="shoulder" ctrlrange="-50 50"/>
    <motor joint="elbow" ctrlrange="-30 30"/>
  </actuator>

  <!-- 传感器（URDF 中没有！） -->
  <sensor>
    <jointpos joint="shoulder"/>
    <jointvel joint="shoulder"/>
    <jointpos joint="elbow"/>
    <jointvel joint="elbow"/>
  </sensor>
</mujoco>

import mujoco
import numpy as np

# 加载模型
model = mujoco.MjModel.from_xml_path("robot.xml")
data = mujoco.MjData(model)

# 仿真循环
for step in range(1000):
    # 设置控制信号（关节力矩）
    data.ctrl[:] = np.array([1.0, -0.5])

    # 推进一步仿真
    mujoco.mj_step(model, data)

    # 读取状态
    qpos = data.qpos.copy()   # 关节位置
    qvel = data.qvel.copy()   # 关节速度
    sensor = data.sensordata  # 传感器数据

    print(f"Step {step}: qpos={qpos}, qvel={qvel}")

# 可视化（打开交互窗口）
import mujoco.viewer
mujoco.viewer.launch(model, data)

import gymnasium as gym

# 使用预置的 MuJoCo 环境
env = gym.make("Ant-v4", render_mode="human")
obs, info = env.reset()

for _ in range(1000):
    action = env.action_space.sample()  # 随机动作
    obs, reward, terminated, truncated, info = env.step(action)
    if terminated or truncated:
        obs, info = env.reset()
env.close()

import gymnasium as gym
from gymnasium import spaces
import mujoco
import numpy as np

class PendulumSwingUp(gym.Env):
    """自定义摆锤 Swing-up 环境"""

    def __init__(self):
        self.model = mujoco.MjModel.from_xml_path("pendulum.xml")
        self.data = mujoco.MjData(self.model)
        self.dt = self.model.opt.timestep

        # 定义观测空间：[cos(θ), sin(θ), dθ]
        self.observation_space = spaces.Box(
            low=-np.inf, high=np.inf, shape=(3,), dtype=np.float32)
        # 定义动作空间：力矩 [-10, 10]
        self.action_space = spaces.Box(
            low=-10.0, high=10.0, shape=(1,), dtype=np.float32)
        self.step_count = 0

    def _get_obs(self):
        theta = self.data.qpos[0]
        dtheta = self.data.qvel[0]
        return np.array([np.cos(theta), np.sin(theta), dtheta],
                        dtype=np.float32)

    def reset(self, seed=None, **kwargs):
        super().reset(seed=seed)
        mujoco.mj_resetData(self.model, self.data)
        # 随机初始角度（倒垂状态附近）
        self.data.qpos[0] = np.pi + np.random.uniform(-0.1, 0.1)
        self.step_count = 0
        return self._get_obs(), {}

    def step(self, action):
        self.data.ctrl[:] = action
        # 执行多个仿真步（降低控制频率）
        for _ in range(5):
            mujoco.mj_step(self.model, self.data)
        self.step_count += 1

        obs = self._get_obs()
        # 奖励：cos(θ) 在竖直向上时为 1
        theta = self.data.qpos[0]
        reward = np.cos(theta) - 0.01 * action[0]**2

        terminated = False
        truncated = self.step_count >= 200
        return obs, reward, terminated, truncated, {}