Function main

mimic.py:22–113

(argv)

def main(argv):
    with open(FLAGS.config, 'r') as f:
        args = AttrDict(yaml.safe_load(f))

    logdir = "logs/exp"
    pretrain = None
    for k, v in args.items():
        if k == 'motion_dir':
            logdir += f"_{v.split('/')[-1].split('.')[0]}"
        elif k == 'pretrain':
            from brax.io import model
            pretrain = model.load_params(v)
            logdir += f"_resume"
        else:
            logdir += f"_{v}"

    model_fn = functools.partial(
        apg_networks.make_apg_networks,
        hidden_layer_sizes=(512, 512, 512) if not args.get('large', False) else (1024,) * 5
    )

    local_device_count = jax.local_device_count()
    train_set = MotionDataset(motion_dir=args.motion_dir, seq_len=args.ep_len, subset=args.get('subset', None), resample=args.get('resample', False))
    train_sampler = RandomSampler(train_set, replacement=True, num_samples=2**31)
    train_loader = NumpyLoader(
        train_set,
        sampler=train_sampler,
        batch_size=args.num_envs * local_device_count,
        num_workers=0
    )

    test_set = MotionDataset(motion_dir=args.motion_dir, seq_len=args.ep_len_eval, subset=args.get('subset', None))
    test_sampler = RandomSampler(test_set, replacement=True, num_samples=2**31)
    test_loader = NumpyLoader(
        test_set,
        batch_size=args.num_eval_envs,
        sampler=test_sampler,
        num_workers=0
    )
    train_env = envs.get_environment(
        env_name="humanoid_mimic_train",
        system_config=args.system_config,
        early_termination=args.early_termination,
        demo_replay_mode=args.demo_replay_mode,
        err_threshold=args.threshold,
        replay_rate=args.replay_rate,
        reward_scaling=args.reward_scaling,
        vel_weight=args.vel_weight,
        rot_weight=args.get('rot_weight', 0.25),
        ang_weight=args.get('ang_weight', 0.01),
        foot_weight=args.get('foot_weight', 0.0),
        local=args.get('local', 'none'),
    )

    eval_env = envs.get_environment(
        env_name="humanoid_mimic",
        system_config=args.system_config,
        vel_weight=args.vel_weight,