def main():
args = parse_args()
env = Protein_Folding_Environment(ref_pdb=args.ref_pdb)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
net = Net(args, device)
print(f'Using Device: {device}')
if args.parameters is not None:
net.load_state_dict(torch.load(args.parameters))
net.to(device)
agent = PPO(model=net, env=env, args=args, device=device)
agent.train()
agent.done()
policy = CategoricalPolicy(model, recurrent, action_size)
else:
raise NotImplementedError
policy.to(device)
#############
## STORAGE ##
#############
print('INITIALIZAING STORAGE...')
hidden_state_dim = model.output_dim
storage = Storage(observation_shape, hidden_state_dim, n_steps, n_envs,
device)
###########
## AGENT ##
###########
print('INTIALIZING AGENT...')
algo = hyperparameters.get('algo', 'ppo')
if algo == 'ppo':
from agents.ppo import PPO as AGENT
else:
raise NotImplementedError
agent = AGENT(env, policy, logger, storage, device, num_checkpoints,
**hyperparameters)
##############
## TRAINING ##
##############
print('START TRAINING...')
agent.train(num_timesteps)