【Hierarchical RL】Options Framework（选项框架）

"""《Options Framework实现项目》
    时间：2024.10
    环境：CartPole
    作者：不去幼儿园
"""
import gym
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
import time

# 环境设置
env = gym.make('CartPole-v1')
state_dim = env.observation_space.shape[0]
action_dim = env.action_space.n

# 高层策略网络
class ManagerNet(nn.Module):
    def __init__(self, state_dim, option_dim):
        super(ManagerNet, self).__init__()
        self.fc = nn.Sequential(
            nn.Linear(state_dim, 128),
            nn.ReLU(),
            nn.Linear(128, option_dim)
        )

    def forward(self, state):
        return self.fc(state)

# 选项策略网络
class OptionNet(nn.Module):
    def __init__(self, state_dim, action_dim):
        super(OptionNet, self).__init__()
        self.fc = nn.Sequential(
            nn.Linear(state_dim, 128),
            nn.ReLU(),
            nn.Linear(128, action_dim)
        )

    def forward(self, state):
        return self.fc(state)

# 选项类定义
class Option:
    def __init__(self, policy_net, termination_net):
        self.policy_net = policy_net
        self.termination_net = termination_net

    def select_action(self, state):
        state = torch.FloatTensor(state).unsqueeze(0)
        action_probs = self.policy_net(state)
        action = torch.argmax(action_probs).item()
        return action

    def should_terminate(self, state):
        state = torch.FloatTensor(state).unsqueeze(0)
        termination_prob = self.termination_net(state)
        return termination_prob.item() > 0.5

# 初始化管理器和选项
option_dim = 2  # 两个选项，分别对应左右
options = [Option(OptionNet(state_dim, action_dim), OptionNet(state_dim, 1)) for _ in range(option_dim)]
manager_net = ManagerNet(state_dim, option_dim)
manager_optimizer = optim.Adam(manager_net.parameters(), lr=0.001)

# 训练
def train_options_framework(env, num_episodes=1000):
    for episode in range(num_episodes):
        state, _ = env.reset()
        total_reward = 0
        option = np.random.choice(len(options))  # 随机选择一个选项

        for t in range(200):  # 每个episode最多执行200步
            # 选择高层选项
            state_tensor = torch.FloatTensor(state).unsqueeze(0)
            option_probs = manager_net(state_tensor)
            option = torch.argmax(option_probs).item()

            # 执行选项的策略
            action = options[option].select_action(state)
            next_state, reward, done, _, _ = env.step(action)
            total_reward += reward

            # 检查选项终止条件
            if options[option].should_terminate(next_state) or done:
                option = np.random.choice(len(options))  # 重新选择选项

            state = next_state

            if done:
                break

        print(f'Episode {episode}: Total Reward: {total_reward}')


# 开始训练
print("Starting training...")
train_options_framework(env, num_episodes=100)

# 测试函数（带动画显示）
def test_options_framework_with_render(env, num_episodes=5):
    for episode in range(num_episodes):
        state, _ = env.reset()
        total_reward = 0
        option = np.random.choice(len(options))  # 随机选择一个选项

        for t in range(200):  # 每个episode最多执行200步
            env.render()  # 显示动画

            # 选择高层选项
            state_tensor = torch.FloatTensor(state).unsqueeze(0)
            option_probs = manager_net(state_tensor)
            option = torch.argmax(option_probs).item()

            # 执行选项的策略
            action = options[option].select_action(state)
            next_state, reward, done, _, _ = env.step(action)
            total_reward += reward

            # 检查选项终止条件
            if options[option].should_terminate(next_state) or done:
                option = np.random.choice(len(options))  # 重新选择选项

            state = next_state

            if done:
                break

        print(f'Test Episode {episode}: Total Reward: {total_reward}')
        time.sleep(0.01)  # 控制动画播放速度

    env.close()  # 关闭环境


# 测试阶段（显示动画）
print("Starting testing with render...")
env = gym.make('CartPole-v1', render_mode="human")
test_options_framework_with_render(env, num_episodes=5)