File size: 3,010 Bytes
08e2c92 | 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 | #!/usr/bin/env python3
"""Record landing videos of the 14-parameter LunarLander policy (headless).
Writes `lunarlander-14params.mp4` (a few episodes) and `lunarlander-14params.gif`
(one episode, downscaled) next to this file.
"""
import os
os.environ.setdefault("SDL_VIDEODRIVER", "dummy")
os.environ.setdefault("SDL_AUDIODRIVER", "dummy")
import numpy as np
import gymnasium as gym
import imageio.v2 as imageio
from model import TinyLunarPolicy
FPS = 50
N_MP4_EPISODES = 3
HERE = os.path.dirname(os.path.abspath(__file__))
SEEDS = [40122, 42940, 51305, 55316, 60944, 61197, 93978, 96327, 103680, 157789]
def record(policy, n_episodes, min_return=250.0):
env = gym.make("LunarLander-v3", render_mode="rgb_array")
clips = []
for seed in SEEDS:
for ep in range(100):
o, _ = env.reset(seed=seed * 1000 + ep)
frames, tot, last, term, trunc = [], 0.0, 0.0, False, False
frames.append(env.render())
while not (term or trunc):
o, r, term, trunc, _ = env.step(policy.act(o))
tot += r
last = r
frames.append(env.render())
if term and not trunc and last == 100.0 and tot >= min_return:
clips.append((seed, ep, tot, frames))
print(" клип: seed %d ep %d возврат %.1f кадров %d" % (seed, ep, tot, len(frames)),
flush=True)
break # один клип на сид -> разные старты в видео
if len(clips) >= n_episodes:
env.close()
return clips
env.close()
return clips
def main():
policy = TinyLunarPolicy.from_safetensors(os.path.join(HERE, "model.safetensors"))
print("запись эпизодов...", flush=True)
clips = record(policy, N_MP4_EPISODES)
if not clips:
raise SystemExit("не нашёл подходящих эпизодов")
# mp4: склеиваем клипы с короткой белой паузой между ними
white = np.full_like(clips[0][3][0], 255)
frames = []
for i, (s, e, r, fr) in enumerate(clips):
frames += fr
frames += [white] * 15
mp4 = os.path.join(HERE, "lunarlander-14params.mp4")
with imageio.get_writer(mp4, fps=FPS, codec="libx264", quality=8,
pixelformat="yuv420p", macro_block_size=8) as w:
for f in frames:
w.append_data(np.asarray(f, np.uint8))
print("saved:", mp4, "| кадров", len(frames), "| ~%.1f c" % (len(frames) / FPS))
# gif: один клип, уменьшенный вдвое и 25 fps
fr = clips[0][3]
small = [np.asarray(f, np.uint8)[::2, ::2] for f in fr[::2]]
gif = os.path.join(HERE, "lunarlander-14params.gif")
imageio.mimsave(gif, small, fps=25)
print("saved:", gif, "| кадров", len(small), "| размер", small[0].shape)
if __name__ == "__main__":
main()
|