Files
kair_algorithms_draft/scripts/examples/lunarlander_continuous_v2/ddpg.py
T
Kyunghwan Kim ecb42d30d2 Add per (#8)
* Add per and modify etc

* Replace pre-commit-config.yaml and add pre-commit hook in .git

* Modify .gitignore

* Modify .gitignore

* Modify buffer and code

* Modify replay buffer and per

* Modify .gitignore
2019-02-16 17:30:12 +09:00

109 lines
2.6 KiB
Python

# -*- coding: utf-8 -*-
"""Run module for DDPG on LunarLanderContinuous-v2.
- Author: Curt Park
- Contact: curt.park@medipixel.io
"""
import argparse
import gym
import torch
import torch.optim as optim
from algorithms.common.networks.mlp import MLP
from algorithms.common.noise import OUNoise
from algorithms.ddpg.agent import Agent
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# hyper parameters
hyper_params = {
"GAMMA": 0.99,
"TAU": 1e-3,
"BUFFER_SIZE": int(1e5),
"BATCH_SIZE": 128,
"LR_ACTOR": 1e-3,
"LR_CRITIC": 1e-3,
"OU_NOISE_THETA": 0.0,
"OU_NOISE_SIGMA": 0.0,
"WEIGHT_DECAY": 1e-6,
}
def run(env: gym.Env, args: argparse.Namespace, state_dim: int, action_dim: int):
"""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_critic = [256, 256]
# create actor
actor = MLP(
input_size=state_dim,
output_size=action_dim,
hidden_sizes=hidden_sizes_actor,
output_activation=torch.tanh,
).to(device)
actor_target = MLP(
input_size=state_dim,
output_size=action_dim,
hidden_sizes=hidden_sizes_actor,
output_activation=torch.tanh,
).to(device)
actor_target.load_state_dict(actor.state_dict())
# create critic
critic = MLP(
input_size=state_dim + action_dim,
output_size=1,
hidden_sizes=hidden_sizes_critic,
).to(device)
critic_target = MLP(
input_size=state_dim + action_dim,
output_size=1,
hidden_sizes=hidden_sizes_critic,
).to(device)
critic_target.load_state_dict(critic.state_dict())
# create optimizer
actor_optim = optim.Adam(
actor.parameters(),
lr=hyper_params["LR_ACTOR"],
weight_decay=hyper_params["WEIGHT_DECAY"],
)
critic_optim = optim.Adam(
critic.parameters(),
lr=hyper_params["LR_CRITIC"],
weight_decay=hyper_params["WEIGHT_DECAY"],
)
# noise
noise = OUNoise(
action_dim,
theta=hyper_params["OU_NOISE_THETA"],
sigma=hyper_params["OU_NOISE_SIGMA"],
)
# make tuples to create an agent
models = (actor, actor_target, critic, critic_target)
optims = (actor_optim, critic_optim)
# create an agent
agent = Agent(env, args, hyper_params, models, optims, noise)
# run
if args.test:
agent.test()
else:
agent.train()