rl_ppo/visualize.py at master · seoneum/rl_ppo · 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
128
129
130
131
132
133

import os, sys, re, glob, argparse, inspect
import numpy as np, imageio, ray
from gymnasium.wrappers import TimeLimit
from ray.tune.registry import register_env
from ray.rllib.algorithms.algorithm import Algorithm
from envs.quad_env import QuadrupedEnv

def _has_display():
    if sys.platform.startswith("linux"):
        return os.environ.get("DISPLAY") is not None
    return True

def make_env(env_config):
    env = QuadrupedEnv(**env_config)
    return TimeLimit(env, max_episode_steps=env_config.get("time_limit", 1000))

def _find_ckpt(path):
    p = os.path.expanduser(path)
    if os.path.isdir(p):
        cands = glob.glob(os.path.join(p, "**", "checkpoint_*"), recursive=True)
        cands += glob.glob(os.path.join(p, "**", "checkpoint-*"), recursive=True)
        if not cands: raise FileNotFoundError("checkpoint not found in dir")
        def key(x):
            m = re.search(r"(checkpoint[_\-])(\d+)$", x)
            return int(m.group(2)) if m else -1
        return sorted(cands, key=key)[-1]
    return p

def parse_args():
    ap = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    ap.add_argument("--checkpoint", type=str, default=None)
    ap.add_argument("--exp-dir", type=str, default=None)
    ap.add_argument("--render", type=str, default="human", choices=["human", "rgb_array"])
    ap.add_argument("--output", type=str, default="eval.mp4")
    ap.add_argument("--fps", type=int, default=30)
    ap.add_argument("--steps", type=int, default=2000)
    ap.add_argument("--seed", type=int, default=0)

    # env overrides (중요한 것만)
    ap.add_argument("--xml-path", type=str, default=None)
    ap.add_argument("--time-limit", type=int, default=None)
    ap.add_argument("--min-episode-steps", type=int, default=None)
    ap.add_argument("--terminate-when-unhealthy", type=int, default=None, help="1/0")
    ap.add_argument("--health-grace", type=float, default=None)
    ap.add_argument("--shin-term", type=float, default=None)
    ap.add_argument("--back-pitch-deg", type=float, default=None)
    ap.add_argument("--back-pitch-time", type=float, default=None)
    ap.add_argument("--healthy-tilt-deg", type=float, default=None)

    # 제어/PD/액션모드
    ap.add_argument("--action-mode", type=str, default=None,
                      choices=["raw_torque","pd_residual","pd_target","pd_target_residual"])
    ap.add_argument("--pd-kp", type=float, default=None)
    ap.add_argument("--pd-kd", type=float, default=None)
    ap.add_argument("--pd-strength", type=float, default=None)
    ap.add_argument("--residual-torque-scale", type=float, default=None)
    ap.add_argument("--action-angle-scale", type=float, default=None)

    return ap.parse_args()

def main():
    args = parse_args()

    if not (args.render == "human" and _has_display()):
        os.environ["MUJOCO_GL"] = "egl"
    else:
        os.environ["MUJOCO_GL"] = "glfw"

    ray.init(ignore_reinit_error=True, include_dashboard=False, log_to_driver=True)
    register_env("Quadruped-v0", make_env)

    if args.checkpoint:
        ckpt = _find_ckpt(args.checkpoint)
    elif args.exp_dir:
        ckpt = _find_ckpt(args.exp_dir)
    else:
        raise ValueError("--checkpoint or --exp-dir required")

    algo = Algorithm.from_checkpoint(ckpt)
    env_cfg = dict(getattr(algo.config, "env_config", {}) or {})

    def set_if(k, v):
        if v is not None: env_cfg[k] = v

    # overrides
    set_if("xml_path", args.xml_path)
    set_if("time_limit", args.time_limit or args.steps)
    set_if("min_episode_steps", args.min_episode_steps)
    if args.terminate_when_unhealthy is not None:
        env_cfg["terminate_when_unhealthy"] = bool(int(args.terminate_when_unhealthy))
    set_if("health_terminate_grace_time", args.health_grace)
    set_if("shin_contact_terminate_time", args.shin_term)
    set_if("back_pitch_terminate_deg", args.back_pitch_deg)
    set_if("back_pitch_terminate_time", args.back_pitch_time)
    set_if("healthy_tilt_deg", args.healthy_tilt_deg)
    set_if("action_mode", args.action_mode)
    set_if("pd_kp", args.pd_kp)
    set_if("pd_kd", args.pd_kd)
    set_if("pd_strength", args.pd_strength)
    set_if("residual_torque_scale", args.residual_torque_scale)
    set_if("action_angle_scale", args.action_angle_scale)
    env_cfg["render_mode"] = args.render

    # build env
    env = make_env(env_cfg)
    obs, info = env.reset(seed=args.seed)
    frames, total = [], 0.0
    for t in range(int(env_cfg.get("time_limit", args.steps))):
        # policy step
        out = algo.compute_single_action(obs, explore=False)
        act = out[0] if isinstance(out, tuple) else out
        obs, rew, term, trunc, info = env.step(act)
        total += float(rew)
        if args.render == "human":
            env.render()
        else:
            frame = env.render()
            if frame is not None:
                frames.append(np.clip(frame, 0, 255).astype(np.uint8))
        if term or trunc:
            break

    if args.render == "rgb_array" and frames:
        imageio.mimsave(args.output, frames, fps=args.fps)
        print(f"saved: {args.output}")

    print(f"done. total reward: {total:.3f}")
    env.close()
    ray.shutdown()

if __name__ == "__main__":
    main()