Homework

3D grid environment to a goal location

Code

import numpy as np
import matplotlib.pyplot as plt
from IPython.display import HTML
import env
environment = env.SingleAgentThreeDGridEnv()

learning_rate = 0.1
discount_factor = 0.99
exploration_rate = 1.0
max_exploration_rate = 1.0
min_exploration_rate = 0.1
exploration_decay_rate = 0.01
total_episodes = 10000
def get_state_index(state):
    x, y, z = state
    return x * environment.shape[1] * environment.shape[2] + y * environment.shape[2] + z


# 创建环境实例
environment = env.SingleAgentThreeDGridEnv()

# 获取动作空间的大小
action_space_size = len(environment.move)  # 因为move列表中的每个元素都对应一个动作

# 获取状态空间的大小
state_space_size = np.prod(environment.shape)  # 假设每个维度的每个值都是一个可能的状态

q_table = np.zeros((state_space_size, action_space_size))

# 初始化环境
environment = env.SingleAgentThreeDGridEnv()

# 定义状态映射函数
def get_state_index(state):
    # 查找代理位置的索引
    return np.where(state.flatten() == 1)[0][0]

# 定义Q表的大小
state_space_size = np.prod(environment.shape)
action_space_size = len(environment.move)
q_table = np.zeros((state_space_size, action_space_size))

# Q学习参数
learning_rate = 0.1
discount_factor = 0.99
exploration_rate = 1.0
max_exploration_rate = 1.0
min_exploration_rate = 0.1
exploration_decay_rate = 0.995
total_episodes = 10000  # 根据需要调整

for episode in range(total_episodes):
    state = environment.reset()
    state_index = get_state_index(state)
    done = False
    total_reward = 0

    while not done:

      # 探索或利用
        if np.random.uniform(0, 1) < exploration_rate:
            action = np.random.choice(action_space_size)  # 随机选择动作
        else:
            action = np.argmax(q_table[state_index, :])  # 利用Q表选择最佳动作

        # 执行动作并观察结果
        new_state, reward, done = environment.step(action)
        new_state_index = get_state_index(new_state)

        # 更新Q表
        q_table[state_index, action] = (1 - learning_rate) * q_table[state_index, action] + \
            learning_rate * (reward + discount_factor * np.max(q_table[new_state_index, :]))

        # 更新状态
        state_index = new_state_index
        total_reward += reward

        # 更新探索率
        exploration_rate = min_exploration_rate + \
        (max_exploration_rate - min_exploration_rate) * np.exp(-exploration_decay_rate * episode)
        # 打印每个回合的总奖励
        if episode % 100 == 0:
          print(f"Episode {episode}, Total reward: {total_reward}")

# 打印训练后的Q表
print("Trained Q-table:")
print(q_table)

# 定义从状态到Q表索引的映射函数
def get_state_index(state):
    return np.where(state.flatten() == 1)[0][0]

# 生成最优策略路径
optimal_policy = []
state = environment.reset()

while True:  # 使用while循环直到任务完成
    state_index = get_state_index(state)
    action = np.argmax(q_table[state_index, :])
    optimal_policy.append(action)
    new_state, reward, done = environment.step(action)

    state = new_state
    if done:
        if reward == environment.goal_reward:
            print("Reached the goal!")
        else:
            print("Hit an obstacle or took too many steps.")
        break

# 使用环境的animate_solution函数生成动画
environment.reset()
ani = environment.animate_solution(optimal_policy)
ani.save('/Users/Chang/Desktop/animation.mp4')  # 保存为视频文件

env.py

import numpy as np
from matplotlib import pyplot as plt, animation
from mpl_toolkits.mplot3d import Axes3D
from matplotlib.animation import FuncAnimation

move = [(0, 0, 0), (0, 0, 1), (0, 0, -1), (0, -1, 0), (0, 1, 0), (1, 0, 0), (-1, 0, 0)]

class SingleAgentThreeDGridEnv:
    def __init__(self, stay_penalty=-2, collision_penalty=-2, goal_reward=10, step_penalty=-1):
        self.move = [(0, 0, 0), (0, 0, 1), (0, 0, -1), (0, -1, 0), (0, 1, 0), (1, 0, 0), (-1, 0, 0)]
        self.shape = (10, 10, 10)
        self.start = (0, 0, 0)
        self.n_obstacles = 15
        self.stay_penalty = stay_penalty
        self.collision_penalty = collision_penalty
        self.goal_reward = goal_reward
        self.step_penalty = step_penalty


        self.destination= (9,9,9) #tuple(np.random.randint(0, s) for s in shape)

        self.obstacles = set()
        while len(self.obstacles) < self.n_obstacles:
            obstacle = tuple(np.random.randint(0, s) for s in self.shape)
            if obstacle != self.start and obstacle != self.destination:
                self.obstacles.add(obstacle)




        self.agent_position = self.start
        assert self.start not in self.obstacles, "Start position must not be an obstacle"
        assert self.destination not in self.obstacles, "Destination position must not be an obstacle"

    def reset(self):
        self.agent_position = self.start
        self.obstacles = set()
        while len(self.obstacles) < self.n_obstacles:
            obstacle = tuple(np.random.randint(0, s) for s in self.shape)
            if obstacle != self.start and obstacle != self.destination:
                self.obstacles.add(obstacle)
        return self._get_state()

    def step(self, action):

        new_position = tuple(np.add(self.agent_position, self.move[action]))

        if new_position != self.agent_position:
            if (0 <= new_position[0] < self.shape[0] and
                0 <= new_position[1] < self.shape[1] and
                0 <= new_position[2] < self.shape[2] and
                new_position not in self.obstacles):
                self.agent_position = new_position
                reward = self.goal_reward if new_position == self.destination else self.step_penalty
                done = new_position == self.destination
            else:
                reward = self.collision_penalty
                done = False
        else:
            reward = self.stay_penalty
            done = False

        return self._get_state(), reward, done

    def _get_state(self):
        state = np.zeros(self.shape)
        for obs in self.obstacles:
            state[obs] = -1
        state[self.destination] = -2
        state[self.agent_position] = 1
        return state

    def _get_reward(self, position):
        if position == self.destination:
            return 100
        return -1

    def render(self):
      self.ax.cla()  # Clear the axis to redraw

      grid = self._get_state()  # Get the grid state

      # Only fill spaces that are not empty
      filled = np.zeros(grid.shape, dtype=bool)
      filled[grid != 0] = True

      # Create a color mapping for the grid
      colors = np.empty(grid.shape, dtype=object)
      colors[grid == -1] = 'red'    # Obstacles in red
      colors[grid == 1] = 'green'   # Agent in green
      colors[grid == -2] = 'yellow' # Destination in yellow

      # Plot the grid using voxels
      self.ax.voxels(filled, facecolors=colors, edgecolor='k')

      # Setting the labels for clarity
      self.ax.set_xlabel('X axis')
      self.ax.set_ylabel('Y axis')
      self.ax.set_zlabel('Z axis')

    def animate_solution(self, policy):
      fig = plt.figure()
      ax = fig.add_subplot(111, projection='3d')
      self.ax = ax

      # Prepare the animation function
      def update(num):
          action = policy[num]
          self.step(action)
          self.render()

          return self.ax,

      # Create the animation
      ani = animation.FuncAnimation(fig, update, frames=len(policy), blit=False, interval=500, repeat=False)
      plt.close(fig)  # Prevent duplicate display
      return ani



class MultiAgentThreeDGridEnv:
    def __init__(self, shape=(10, 10, 10), starts=[(0, 0, 0), (0, 0, 9)], destinations=[(9, 9, 9), (9, 9, 0)], n_obstacles=15, stay_penalty=-15, collision_penalty=-10, goal_reward=100, step_penalty=-1):
        self.shape = shape
        self.starts = starts
        self.destinations = destinations
        self.n_obstacles = n_obstacles
        self.stay_penalty = stay_penalty
        self.collision_penalty = collision_penalty
        self.goal_reward = goal_reward
        self.step_penalty = step_penalty

        self.obstacles = self._generate_obstacles()
        self.agent_positions = list(starts)

        assert not (set(starts) & self.obstacles), "Start positions must not be in obstacles"
        assert not (set(destinations) & self.obstacles), "Destination positions must not be in obstacles"

    def _generate_obstacles(self):
        obstacles = set()
        while len(obstacles) < self.n_obstacles:
            obstacle = tuple(np.random.randint(0, s) for s in self.shape)
            if obstacle not in self.starts and obstacle not in self.destinations:
                obstacles.add(obstacle)
        return obstacles

    def reset(self):
        self.agent_positions = list(self.starts)
        self.obstacles = self._generate_obstacles()
        return [self._get_state(i) for i in range(len(self.starts))]

    def step(self, actions):
        new_positions = []
        rewards = []
        dones = []

        for i, action in enumerate(actions):
            move = [(0, 0, 0), (0, 0, 1), (0, 0, -1), (0, -1, 0), (0, 1, 0), (1, 0, 0), (-1, 0, 0)]
            new_position = tuple(np.add(self.agent_positions[i], move[action]))

            if new_position == self.agent_positions[i]:
                reward = self.stay_penalty
            elif (0 <= new_position[0] < self.shape[0] and
                  0 <= new_position[1] < self.shape[1] and
                  0 <= new_position[2] < self.shape[2] and
                  new_position not in self.obstacles and
                  new_positions.count(new_position) == 1):
                self.agent_positions[i] = new_position
                reward = self.goal_reward if new_position == self.destinations[i] else self.step_penalty
            else:
                reward = self.collision_penalty

            new_positions.append(new_position)
            rewards.append(reward)
            dones.append(new_position == self.destinations[i])

        states = [self._get_state(i) for i in range(len(self.starts))]
        return states, rewards, dones

    def _get_state(self, agent_index):
        state = np.zeros(self.shape)
        for obs in self.obstacles:
            state[obs] = -1
        state[self.destinations[agent_index]] = -2
        for i, pos in enumerate(self.agent_positions):
            if state[pos] == 0:
                state[pos] = i + 1
        return state

    def _get_reward(self, agent_index, position):
        if position == self.destinations[agent_index]:
            return self.goal_reward
        return self.step_penalty