# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# !pip install 'stable-baselines3[extra]'

# ## 1. Import dependecies

import os
import gym
from stable_baselines3 import PPO
from stable_baselines3.common.vec_env import DummyVecEnv
from stable_baselines3.common.evaluation import evaluate_policy 

# ### space type
#
# Box : n dimensional tensor, range of values ( e.g Box(0,1,shape = (3,3))
#
# Discrete : set of items ( e.g Discrete(3))
#
# Tuple : tuple of other spaces e.g Box or Discrete 
#         ( e.g. Tuple((Discrete(2), Box(0, 100, shape = (1, ))))
#    
# MultiBinary : One-hot encoded binary value ( e.g. MultiDiscrete([5, 2, 2])

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

# +
print(env.action_space)
print(env.action_space.sample())

print(env.observation_space)
print(env.observation_space.sample())
# -

episodes = 5 
for epi in range(1, episodes + 1):
    state = env.reset()
    done = False
    score = 0
    
    while not done:
        # env.render()
        action = env.action_space.sample()
        n_state, reward, done, info = env.step(action)
        score += reward
    print('Episodes:{}, Score:{}'.format(epi, score))
env.close

# ## 2. Training

# make your Directory first
log_path = os.path.join('Training', 'logs')
log_path

env = gym.make('CartPole-v0')
env = DummyVecEnv([lambda: env])
model = PPO('MlpPolicy', env, verbose = 1, tensorboard_log = log_path)

model.learn(total_timesteps = 20000)

# ## 3. save and reload model

PPO_Path = os.path.join('Training', 'Saved Models', 'PPO_Model_Cartpole')
PPO_Path

model.save(PPO_Path)
model

model = PPO.load(PPO_Path, env = env)

model.learn(total_timesteps = 100)

# ## 4. Evaluation

evaluate_policy(model, env, n_eval_episodes = 10, render = False)

env.close()

# # 5. Test Model

obs = env.reset()
model.predict(obs)    # 두개가 나오므로 하나는 고려하지 않는다.

#
#  model.predict??
#  
#        Get the model's action(s) from an observation
#
#        :param observation: the input observation
#        :param state: The last states (can be None, used in recurrent policies)
#        :param mask: The last masks (can be None, used in recurrent policies)
#        :param deterministic: Whether or not to return deterministic actions.
#        :return: the model's action and the next state
#

action, _ = model.predict(obs)
action

episodes = 5 
for epi in range(1, episodes + 1):
    obs = env.reset()                 # state = env.reset() 에서 변경
    done = False
    score = 0
    
    while not done:
        # env.render()
        action, _ = model.predict(obs)      # env.action_space.sample() 에서 변경
                                      # random action이 아닌 모델이 action을 선택한다
        obs, reward, done, info = env.step(action)
        score += reward
    print('Episodes:{}, Score:{}'.format(epi, score))
env.close

obs = env.reset()

obs

model.predict(obs)        # array([1]) : 최적의 선택

env.action_space.sample()   # 1 : 랜덤하게 선택

action, _ = model.predict(obs)
state, reward, done, info = env.step(action)
print(state)
print(reward)
print(done)
print(info)

# # 6. Viewing Logs in Tensorboard

training_log_path = os.path.join(log_path, 'PPO_2')
training_log_path

# !tensorboard --logdir={training_log_path}

# # 7. Adding a callboack to the training stage

# +
# import Callback helpers

from stable_baselines3.common.callbacks import EvalCallback, StopTrainingOnRewardThreshold
# EvalCallback : call model back and evaluate to find best model
# StopTrainingOnRewardThreshold : stop training when get reward treshold
# -

save_path = os.path.join('Training', 'Saved Models')

# +
stop_callback = StopTrainingOnRewardThreshold(reward_threshold = 200, verbose = 1)

eval_callback = EvalCallback(env,
                            callback_on_new_best=stop_callback,
                            eval_freq=10000,
                            best_model_save_path=save_path,
                            verbose=1)
# -

model = PPO('MlpPolicy', env, verbose = 1, tensorboard_log = log_path)

model.learn(total_timesteps = 20000, callback = eval_callback)

# # 8. Changing Policies : specipy new neural network

net_arch = [dict(pi=[128,128,128,128], vf=[128,128,128,128])]
# pi : policy (custom actor) , and get 4 layers which has 128 units(nodes)
# vf : value functino (custom critic) , and get 4 layers which has 128 units(nodes)

model = PPO('MlpPolicy', env, verbose = 1, tensorboard_log = log_path, 
            policy_kwargs={'net_arch':net_arch})

model.learn(total_timesteps = 20000, callback = eval_callback)

# # 10. Using an Alternative Algorithm

from stable_baselines3 import DQN

model = DQN('MlpPolicy', env, verbose = 1, tensorboard_log = log_path)

model.learn(total_timesteps = 20000)

StarcoderdataJupyterScriptsDedupFiltered