create_movie.py

import cv2
import numpy as np

from apple.envs.discrete_apple import get_apple_env
from apple.evaluation.render_episode import render_episode
from apple.models.categorical_policy import CategoricalPolicy

env_kwargs = dict(
    start_x=0,
    goal_x=8,
    c=0.5,
    time_limit=100,
    bias_in_state=True,
    position_in_state=False,
    apple_in_state=True,
)

env = get_apple_env("full", render_mode="rgb_array", **env_kwargs)

model = CategoricalPolicy(env.observation_space.shape[0], 1, 0, 10)
data = list(render_episode(env, model))

video_frames = np.stack([d["pixel_array"] for d in data], axis=0)


# Define the output file name and location
output_file = "apple_retrieval.mp4"

# Define the video codec and parameters
fourcc = cv2.VideoWriter_fourcc(*"XVID")
fps = 1.0

# Create a VideoWriter object
writer = cv2.VideoWriter(output_file, fourcc, fps, (video_frames.shape[2], video_frames.shape[1]))

# Write the video frames to the output file
for frame in video_frames:
    frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
    writer.write(frame)

# Release the VideoWriter object
writer.release()