Files
kair_algorithms_draft/scripts/examples/lunarlander_continuous_v2/sac.py
T
Kyunghwan Kim d2769dfa9d Convert code to python 2.7 (#35)
* Convert code format to python2.7 (SAC)

* Convert code format python2.7 (TD3, all fD)

* Remove no use import and black setting

* Change SAC param

* Change env name Reacher-v2 to v1

* Remove old version reacher training script

* Convert code format python2.7

* Modify .travis.yml

* Add install command python3.6 & black on Makefile

* Fix seperator to tab on Makefile

* Modify Makefile

* Fix little error

* Change td3 gamma parameter
2019-03-25 19:07:19 +09:00

112 lines
3.0 KiB
Python

# -*- coding: utf-8 -*-
"""Run module for SAC on LunarLanderContinuous-v2.
- Author: Curt Park
- Contact: curt.park@medipixel.io
"""
import numpy as np
import torch
import torch.optim as optim
from algorithms.common.networks.mlp import MLP, FlattenMLP, TanhGaussianDistParams
from algorithms.sac.agent import Agent
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# hyper parameters
hyper_params = {
"GAMMA": 0.99,
"TAU": 5e-3,
"W_ENTROPY": 1e-3,
"W_MEAN_REG": 1e-3,
"W_STD_REG": 1e-3,
"W_PRE_ACTIVATION_REG": 0.0,
"LR_ACTOR": 3e-4,
"LR_VF": 3e-4,
"LR_QF1": 3e-4,
"LR_QF2": 3e-4,
"LR_ENTROPY": 3e-4,
"DELAYED_UPDATE": 2,
"BUFFER_SIZE": int(1e6),
"BATCH_SIZE": 512,
"AUTO_ENTROPY_TUNING": True,
"WEIGHT_DECAY": 0.0,
"INITIAL_RANDOM_ACTION": 5000,
}
def run(env, args, state_dim, action_dim):
"""Run training or test.
Args:
env (gym.Env): openAI Gym environment with continuous action space
args (argparse.Namespace): arguments including training settings
state_dim (int): dimension of states
action_dim (int): dimension of actions
"""
hidden_sizes_actor = [256, 256]
hidden_sizes_vf = [256, 256]
hidden_sizes_qf = [256, 256]
# target entropy
target_entropy = -np.prod((action_dim,)).item() # heuristic
# create actor
actor = TanhGaussianDistParams(
input_size=state_dim, output_size=action_dim, hidden_sizes=hidden_sizes_actor
).to(device)
# create v_critic
vf = MLP(input_size=state_dim, output_size=1, hidden_sizes=hidden_sizes_vf).to(
device
)
vf_target = MLP(
input_size=state_dim, output_size=1, hidden_sizes=hidden_sizes_vf
).to(device)
vf_target.load_state_dict(vf.state_dict())
# create q_critic
qf_1 = FlattenMLP(
input_size=state_dim + action_dim, output_size=1, hidden_sizes=hidden_sizes_qf
).to(device)
qf_2 = FlattenMLP(
input_size=state_dim + action_dim, output_size=1, hidden_sizes=hidden_sizes_qf
).to(device)
# create optimizers
actor_optim = optim.Adam(
actor.parameters(),
lr=hyper_params["LR_ACTOR"],
weight_decay=hyper_params["WEIGHT_DECAY"],
)
vf_optim = optim.Adam(
vf.parameters(),
lr=hyper_params["LR_VF"],
weight_decay=hyper_params["WEIGHT_DECAY"],
)
qf_1_optim = optim.Adam(
qf_1.parameters(),
lr=hyper_params["LR_QF1"],
weight_decay=hyper_params["WEIGHT_DECAY"],
)
qf_2_optim = optim.Adam(
qf_2.parameters(),
lr=hyper_params["LR_QF2"],
weight_decay=hyper_params["WEIGHT_DECAY"],
)
# make tuples to create an agent
models = (actor, vf, vf_target, qf_1, qf_2)
optims = (actor_optim, vf_optim, qf_1_optim, qf_2_optim)
# create an agent
agent = Agent(env, args, hyper_params, models, optims, target_entropy)
# run
if args.test:
agent.test()
else:
agent.train()