|
|
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() |
|
|
|
|
|
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) |
|
|
|
|
|
parser.add_argument("--save_dir", default="results", type=str) |
|
|
|
|
|
parser.add_argument("--bias", default="force", type=str) |
|
|
|
|
|
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) |
|
|
|
|
|
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) |
|
|
|
|
|
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() |
