test_code/run_v2.py at main · plutozc/test_code · GitHub

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
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
#!/usr/bin/env python3
"""CLI for corrected LfD V2 experiments."""

import argparse
import json
from dataclasses import replace
from pathlib import Path

from lfd_v2.experiments import (
    ExperimentConfig,
    run_limit_study,
    run_freeze_plan_study,
    run_main_sweep,
    run_ood_study,
    run_oracle_study,
    run_scaling_study,
    run_smoothness_study,
    write_comparison_report,
)


def _environment_tuple(value):
    return ("compatible", "rigid_body") if value == "both" else (value,)


def parse_cli(argv=None):
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "command",
        choices=("smoke", "main", "smoothness", "limits", "freeze-plan", "ood", "scaling", "oracle", "report", "full"),
    )
    parser.add_argument("--env", choices=("compatible", "rigid_body", "both"), default="both")
    parser.add_argument("--output", type=Path, default=None)
    parser.add_argument("--no-resume", action="store_true")
    parser.add_argument("--seed-start", type=int)
    parser.add_argument("--seed-end", type=int)
    args = parser.parse_args(argv)

    if args.command == "smoke":
        config = ExperimentConfig(
            environments=_environment_tuple(args.env),
            methods=("bc", "position_bc", "promp_cl"),
            budgets=(5, 25),
            seeds=(42, 43),
            eval_episodes=10,
            epochs=20,
            output_root=args.output or Path("results_v2/smoke"),
        )
    elif args.command == "scaling":
        config = ExperimentConfig(
            study="scaling",
            environments=_environment_tuple(args.env),
            methods=("position_bc", "promp_cl"),
            budgets=(50, 100, 200, 400, 800),
            seeds=(42, 43, 44, 45, 46),
            output_root=args.output or Path("results_v2"),
        )
    elif args.command == "smoothness":
        config = ExperimentConfig(
            environments=_environment_tuple(args.env),
            seeds=tuple(range(42, 62)),
            output_root=args.output or Path("results_v2"),
        )
    elif args.command == "freeze-plan":
        config = ExperimentConfig(
            study="freeze_plan",
            environments=_environment_tuple(args.env),
            methods=("promp_cl",),
            budgets=(50, 200),
            output_root=args.output or Path("results_v2"),
        )
    else:
        config = ExperimentConfig(
            environments=_environment_tuple(args.env),
            output_root=args.output or Path("results_v2"),
        )
    if (args.seed_start is None) != (args.seed_end is None):
        parser.error("--seed-start and --seed-end must be provided together")
    if args.seed_start is not None:
        if args.seed_end < args.seed_start:
            parser.error("--seed-end must be >= --seed-start")
        config = replace(config, seeds=tuple(range(args.seed_start, args.seed_end + 1)))
    return args, config


def main(argv=None):
    args, config = parse_cli(argv)
    resume = not args.no_resume
    if args.command in {"smoke", "main"}:
        result = run_main_sweep(config, resume=resume)
    elif args.command == "smoothness":
        result = run_smoothness_study(config, seeds=config.seeds, resume=resume)
    elif args.command == "limits":
        result = run_limit_study(config, resume=resume)
    elif args.command == "freeze-plan":
        result = run_freeze_plan_study(config, resume=resume)
    elif args.command == "ood":
        result = run_ood_study(config, resume=resume)
    elif args.command == "scaling":
        result = run_scaling_study(config, resume=resume)
    elif args.command == "oracle":
        result = run_oracle_study(config, resume=resume)
    elif args.command == "report":
        result = write_comparison_report(config)
    else:
        stages = {}
        stages["main"] = run_main_sweep(config, resume=resume)
        stages["oracle"] = run_oracle_study(config, resume=resume)
        stages["limits"] = run_limit_study(config, resume=resume)
        stages["ood"] = run_ood_study(config, resume=resume)
        stages["smoothness"] = run_smoothness_study(config, resume=resume)
        scaling = ExperimentConfig(
            study="scaling",
            environments=config.environments,
            methods=("position_bc", "promp_cl"),
            budgets=(50, 100, 200, 400, 800),
            seeds=(42, 43, 44, 45, 46),
            output_root=config.output_root,
        )
        stages["scaling"] = run_scaling_study(scaling, resume=resume)
        stages["report"] = write_comparison_report(config)
        result = stages
    print(json.dumps({"command": args.command, **result}, sort_keys=True))


if __name__ == "__main__":
    main()