import time

import gym

import numpy as np

env = gym.make('FrozenLake-v0')

render = False

running_reward = None

Q = np.zeros([env.observation_space.n, env.action_space.n])

lr = .85 # Qleaning的学习率。alpha, if use value function approximation, we can ignore it

lambd = .99 # 折扣率 decay factor

num_episodes = 10000 # 迭代次数，也就是开始10000次游戏

rList = [] # 用于记录每次迭代的总奖励，这样我们就可以知道智能体是否有进步了。rewards for each episode

for i in range(num_episodes):

## 重置环境初始状态

episode_time = time.time() #用于记录运行时间，我们可以通过比较运行时间判断算法效率。

s = env.reset() #重置初始状态。

rAll = 0 #用于记录这次游戏的总奖励，这里先初始化为0

## 开始Qlearning算法

for j in range(99):

if render: env.render() #判断是否渲染环境。

## [敲黑板]

## 从Q表格中，找到当前状态S最大Q值，并在Q值上加上噪音。

## 然后找到最大的Q+噪音的动作

a = np.argmax(Q[s, :] + np.random.randn(1, env.action_space.n) * (1. / (i + 1)))

## 与环境互动，把动作放到env.step()函数，并返回下一状态S1，奖励，done，info

s1, r, d, _ = env.step(a)

## 更新Q表格

Q[s, a] = Q[s, a] + lr * (r + lambd * np.max(Q[s1, :]) - Q[s, a])

rAll += r # rAll累加当前的收获。

s = s1 # 把下一状态赋值给s，准备开始下一步。

if d ==True: # 如果已经到达最终状态，就跳出for循环。(开始下一次迭代)

break

#每次的总收获都放到rlist。可以通过打印看看算法效率。

rList.append(rAll)

#每一次迭代获得的总收获rAll,会以0.01的份额加入到running_reward。(原代码这里rAll用了r，个人认为是rAll更合适)

running_reward = rAll if running_reward is None else running_reward * 0.99 + rAll * 0.01

print("Episode [%d/%d] sum reward: %f running reward: %f took: %.5fs " % \

(i, num_episodes, rAll, running_reward, time.time() - episode_time))

print("Final Q-Table Values:/n %s" % Q)