num_inputs = 4
elif observation_type.startswith('partial_'):
observation_size = int(observation_type.split('_')[-1])
observation_width = 2 * observation_size + 1
num_inputs = 3 * (observation_width**2)
else:
raise ValueError(
'Feedforward agent only compatible with partial and position observations.'
)
model = agents.FeedforwardAgent(num_actions=4,
num_inputs=num_inputs,
num_layers=2,
hidden_units=64).to(args.device)
elif agent_type == 'random':
model = agents.RandomAgent(num_actions=4, device=args.device)
else:
raise ValueError('Unrecognised agent')
if reload:
model.load_state_dict(torch.load(agent_path))
if args.train:
model.train()
else:
torch.no_grad()
model.eval()
print(model)
if args.agent != 'random' and args.train:
optimizer = optim.Adam(model.parameters(), lr=args.lr)
if agent_type == 'conv':
models.append(
agents.ConvAgent(
num_actions=num_actions, num_initial_convs=2, in_channels=in_channels, conv_channels=32,
num_residual_convs=2, num_feedforward=1, feedforward_dim=64,
num_heads=num_heads).to(device=args.device, dtype=dtype)
)
elif agent_type == 'gru':
models.append(
agents.GRUAgent(
num_actions=num_actions, num_initial_convs=2, in_channels=in_channels, conv_channels=32,
num_residual_convs=2, num_feedforward=1, feedforward_dim=64,
num_heads=num_heads).to(device=args.device, dtype=dtype)
)
elif agent_type == 'random':
models.append(agents.RandomAgent(num_actions=num_actions, device=args.device))
else:
raise ValueError('Unrecognised agent type.')
# Load state dict if reload
if reload:
models[i].load_state_dict(
torch.load(args.agent[i])
)
if args.train:
for m in models:
m.train()
else:
torch.no_grad()
for m in models: