brozek/rltorch
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity

Made sure everything went to their appropriate devices

Browse source
This commit is contained in:
Brandon Rozek 2019-02-03 00:45:14 -05:00
parent 39643f04e1
commit 9cd3625fd3
5 changed files with 14 additions and 9 deletions
Download patch file Download diff file Expand all files Collapse all files

6
examples/acrobot.py
View file

@ -90,11 +90,11 @@ logwriter = rltorch.log.LogWriter(logger, SummaryWriter())
# Setting up the networks
device = torch.device("cuda:0" if torch.cuda.is_available() and not config['disable_cuda'] else "cpu")
net = rn.Network(Value(state_size, action_size),
torch.optim.Adam, config, logger = logger, name = "DQN")
target_net = rn.TargetNetwork(net)
torch.optim.Adam, config, device = device, logger = logger, name = "DQN")
target_net = rn.TargetNetwork(net, device = device)
# Actor takes a net and uses it to produce actions from given states
actor = ArgMaxSelector(net, action_size)
actor = ArgMaxSelector(net, action_size, device = device)
# Memory stores experiences for later training
memory = M.ReplayMemory(capacity = config['memory_size'])

6
examples/pong.py
View file

@ -107,11 +107,11 @@ logwriter = rltorch.log.LogWriter(logger, SummaryWriter())
# Setting up the networks
device = torch.device("cuda:0" if torch.cuda.is_available() and not config['disable_cuda'] else "cpu")
net = rn.Network(Value(state_size, action_size),
torch.optim.Adam, config, logger = logger, name = "DQN")
target_net = rn.TargetNetwork(net)
torch.optim.Adam, config, device = device, logger = logger, name = "DQN")
target_net = rn.TargetNetwork(net, device = device)
# Actor takes a network and uses it to produce actions from given states
actor = ArgMaxSelector(net, action_size)
actor = ArgMaxSelector(net, action_size, device = device)
# Memory stores experiences for later training
memory = M.ReplayMemory(capacity = config['memory_size'])