render example episode

create_movie.py

@@ -0,0 +1,42 @@
+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()