File size: 4,438 Bytes
7efee70 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 |
import os
import sys
import argparse
import torch
import wandb
from entangled.mds import MDs
from entangled.utils.logging import Logger
from entangled.entangledsbm import EntangledSBM
def main():
parser = argparse.ArgumentParser()
# System Config
parser.add_argument("--date", type=str)
parser.add_argument("--seed", default=2, type=int)
parser.add_argument("--device", default="cuda:1", type=str)
parser.add_argument("--molecule", default="aldp", type=str)
parser.add_argument('--wandb', action='store_true', default=False)
parser.add_argument('--run_name', default=None, type=str)
# Logger Config
parser.add_argument("--save_dir", default="results", type=str)
# Policy Config
parser.add_argument("--bias", default="force", type=str)
# Sampling Config
parser.add_argument("--start_state", default="c5", type=str)
parser.add_argument("--end_state", default="c7ax", type=str)
parser.add_argument("--num_steps", default=1000, type=int)
parser.add_argument("--timestep", default=1, type=float)
parser.add_argument("--sigma", default=0.1, type=float)
parser.add_argument("--num_samples", default=16, type=int)
parser.add_argument("--temperature", default=300, type=float)
parser.add_argument("--friction", default=0.001, type=float)
parser.add_argument("--rbf", action='store_true', default=False)
parser.add_argument("--use_delta_to_target", action='store_true', default=False)
# Training Config
parser.add_argument("--start_temperature", default=600, type=float)
parser.add_argument("--end_temperature", default=300, type=float)
parser.add_argument("--num_rollouts", default=1000, type=int)
parser.add_argument("--trains_per_rollout", default=1000, type=int)
parser.add_argument("--log_z_lr", default=1e-3, type=float)
parser.add_argument("--policy_lr", default=1e-4, type=float)
parser.add_argument("--batch_size", default=16, type=int)
parser.add_argument("--buffer_size", default=1000, type=int)
parser.add_argument("--max_grad_norm", default=1, type=int)
parser.add_argument("--control_variate", default="global", type=str)
parser.add_argument("--self_normalize", action='store_true', default=False)
parser.add_argument("--importance_sample", action='store_true', default=False)
# path objective
parser.add_argument("--objective", default="ce", type=str)
parser.add_argument("--curriculum_rollouts", default=50, type=int)
parser.add_argument("--sigma_min", default=0.25, type=float)
parser.add_argument("--sigma_max", default=2.0, type=float)
parser.add_argument("--resample_every", default=0, type=int)
parser.add_argument("--resample_jitter", default=1e-3, type=float)
parser.add_argument("--branch_beta", default=5.0, type=float)
parser.add_argument("--bias_scale", default=1.0, type=float)
parser.add_argument("--control_cost", action='store_true', default=False)
parser.add_argument("--early_termination", action='store_true', default=False)
parser.add_argument("--success_distance_threshold", default=1.0, type=float)
parser.add_argument("--adaptive_bias", action='store_true', default=False)
parser.add_argument("--min_bias_scale", default=0.5, type=float)
parser.add_argument("--max_bias_scale", default=3.0, type=float)
parser.add_argument("--target_dist_threshold", default=2.0, type=float)
parser.add_argument("--vel_conditioned", action='store_true', default=False)
args = parser.parse_args()
args.training = True
args.save_dir = args.save_dir
positions_dir = f"{args.save_dir}/positions"
if not os.path.exists(positions_dir):
os.makedirs(positions_dir)
if args.wandb:
wandb.init(project="entangled-tps", config=args, name=args.run_name)
torch.manual_seed(args.seed)
mds = MDs(args)
logger = Logger(args, mds)
target_pos = mds.target_position.squeeze(0)
model = EntangledSBM(args, mds)
temperatures = torch.linspace(
args.start_temperature, args.end_temperature, args.num_rollouts
)
for rollout in range(args.num_rollouts):
model.sample(args, mds, temperatures[rollout])
loss, positions = model.train(args, mds)
model.increment_rollout()
logger(loss, rollout, model.bias_net)
logger.plot()
if __name__ == "__main__":
main() |