jxx123's picture
Toy cartpole simulator app
c785ba7
import streamlit as st
import gymnasium
from gymnasium.wrappers.record_video import RecordVideo
from PIL import Image
import time
steps = st.number_input("Step Number", min_value=1, value=100)
sim_clicked = st.button("Simulate")
if sim_clicked:
env = gymnasium.make("MountainCarContinuous-v0", render_mode="rgb_array")
env = RecordVideo(env, video_folder="./videos")
obs, info = env.reset()
with st.empty():
for i in range(steps):
action = env.action_space.sample()
obs, reward, terminated, truncated, info = env.step(action)
rgb_array = env.render()
image = Image.fromarray(rgb_array)
st.image(image, caption=f"Step {i}")
time.sleep(0.02)
if terminated or truncated:
break
env.close()
st.video("./videos/rl-video-episode-0.mp4")