base gym env

app.py

@@ -0,0 +1,276 @@
+import streamlit as st
+import gymnasium as gym
+import numpy as np
+from PIL import Image
+import time
+
+# Initialize session state variables if they don't exist
+if "env" not in st.session_state:
+    st.session_state.env = gym.make("LunarLander-v2", render_mode="rgb_array")
+    st.session_state.env.reset()
+    st.session_state.frame = st.session_state.env.render()
+if "paused" not in st.session_state:
+    st.session_state.paused = False
+
+
+# Function to reset the environment
+def reset_environment():
+    st.session_state.env.reset()
+
+
+# Function to toggle pause state
+def toggle_pause():
+    st.session_state.paused = not st.session_state.paused
+
+
+# Create the Streamlit app
+st.title("Gymnasium Environment Viewer")
+
+# Add control buttons in a horizontal layout
+col1, col2 = st.columns(2)
+with col1:
+    st.button("Reset Environment", on_click=reset_environment)
+with col2:
+    if st.session_state.paused:
+        st.button("Resume", on_click=toggle_pause)
+    else:
+        st.button("Pause", on_click=toggle_pause)
+
+# Create a placeholder for the image
+image_placeholder = st.empty()
+
+# Create a container for environment info
+sidebar_container = st.sidebar.container()
+
+# Main simulation loop using rerun
+if not st.session_state.paused:
+    # Take a random action
+    action = st.session_state.env.action_space.sample()
+    observation, reward, terminated, truncated, info = (
+        st.session_state.env.step(action)
+    )
+
+    # Render the environment
+    st.session_state.frame = st.session_state.env.render()
+
+    # Reset if the episode is done
+    if terminated or truncated:
+        st.session_state.env.reset()
+# Display the frame
+if st.session_state.paused:
+    image_placeholder.image(
+        st.session_state.frame,
+        caption="Environment Visualization (Paused)",
+        use_column_width=True,
+    )
+else:
+    image_placeholder.image(
+        st.session_state.frame,
+        caption="Environment Visualization",
+        use_column_width=True,
+    )
+
+# Display some information about the environment
+with sidebar_container:
+    st.header("Environment Info")
+    st.write(f"Action Space: {st.session_state.env.action_space}")
+    st.write(f"Observation Space: {st.session_state.env.observation_space}")
+
+# Add auto-refresh logic
+if not st.session_state.paused:
+    time.sleep(0.1)  # Add a small delay to control refresh rate
+    st.rerun()
+
+# fig, ax = plt.subplots()
+# ax.imshow(env.render())
+# st.pyplot(fig)
+# st.image(env.render())
+
+
+# import gymnasium as gym
+# import streamlit as st
+# import numpy as np
+# from udrl.policies import SklearnPolicy
+# from udrl.agent import UpsideDownAgent, AgentHyper
+# from pathlib import Path
+
+# # import json
+
+
+# def normalize_value(value, is_bounded, low=None, high=None):
+#     return (value - low) / (high - low)
+
+
+# def visualize_environment(
+#     state,
+#     env,
+#     # paused,
+#     feature_importances,
+#     epoch,
+#     max_epoch=200,
+# ):
+
+#     st.image(env.render())
+#     st.image(e)
+#     # Render the Gym environment
+#     # env_render = env.render()
+
+#     # # Display the rendered image using Streamlit
+#     # st.image(env_render, caption=f"Epoch {epoch}", use_column_width=True)
+
+#     # Display feature importances using Streamlit metrics
+#     # cols = st.columns(len(feature_importances))
+#     # for i, col in enumerate(cols):
+#     #     col.metric(
+#     #         label=f"Importance {i}", value=f"{feature_importances[i]:.2f}"
+#     #     )
+
+#     # Create buttons using Streamlit
+#     # reset_button = st.button("Reset")
+#     # pause_play_button = st.button("Pause" if not paused else "Play")
+#     # next_button = st.button("Next")
+#     # save_button = st.button("Save")
+
+#     # return reset_button, pause_play_button, next_button, save_button
+
+
+# def run_visualization(
+#     env_name,
+#     agent,
+#     init_desired_return,
+#     init_desired_horizon,
+#     max_epoch,
+#     base_path,
+# ):
+#     # base_path = (
+#     #     Path(base_path) / env_name / agent.policy.estimator.__str__()[:-2]
+#     # )
+#     # base_path.mkdir(parents=True, exist_ok=True)
+#     desired_return = init_desired_return
+#     desired_horizon = init_desired_horizon
+
+#     # Initialize the Gym environment
+#     env = gym.make(env_name, render_mode="rgb_array")
+#     state, _ = env.reset()
+
+#     epoch = 0
+#     # save_index = 0
+
+#     # paused = False
+#     # step = False
+
+#     # # Use Streamlit session state to manage paused state
+#     # if "paused" not in st.session_state:
+#     #     st.session_state.paused = False
+
+#     while True:
+#         # Render and display the environment
+#         env_render = env.render()
+#         # if not st.session_state.pausedor step:
+#         command = np.array(
+#             [
+#                 desired_return * agent.conf.return_scale,
+#                 desired_horizon * agent.conf.horizon_scale,
+#             ]
+#         )
+#         command = np.expand_dims(command, axis=0)
+#         state = np.expand_dims(state, axis=0)
+
+#         action = agent.policy(state, command, True)
+
+#         ext_state = np.concatenate((state, command), axis=1)
+
+#         state, reward, done, truncated, info = env.step(action)
+
+#         feature_importances = {idx: [] for idx in range(ext_state.shape[1])}
+
+#         for t in agent.policy.estimator.estimators_:
+#             branch = np.array(t.decision_path(ext_state).todense(), dtype=bool)
+#             imp = t.tree_.impurity[branch[0]]
+
+#             for f, i in zip(
+#                 t.tree_.feature[branch[0]][:-1], imp[:-1] - imp[1:]
+#             ):
+#                 feature_importances.setdefault(f, []).append(i)
+
+#         # Line 8 Algorithm 2
+#         desired_return -= reward
+#         # Line 9 Algorithm 2
+#         desired_horizon = max(desired_horizon - 1, 1)
+
+#         summed_importances = [
+#             sum(feature_importances.get(k, [0.001]))
+#             for k in range(len(feature_importances.keys()))
+#         ]
+
+#         epoch += 1
+#         visualize_environment(
+#             state,
+#             env,
+#             # st.session_state.paused,  # Use session state
+#             summed_importances,
+#             epoch,
+#             max_epoch,
+#         )
+#         # reset_button, pause_play_button, next_button, save_button = (
+
+#         # )
+
+#         if done or truncated:
+#             state, _ = env.reset()
+#             desired_horizon = init_desired_horizon
+#             desired_return = init_desired_return
+#             epoch = 0
+
+#         # step = False
+
+#         # Handle button clicks
+#         # if reset_button:
+#         #     state, _ = env.reset()
+#         #     desired_horizon = init_desired_horizon
+#         #     desired_return = init_desired_return
+#         #     epoch = 0
+#         # elif pause_play_button:
+#         #     st.session_state.paused = (
+#         #         not st.session_state.paused
+#         #     )  # Toggle paused state
+#         # elif next_button and st.session_state.paused:
+#         #     step = True
+#         # elif save_button:
+#         #     # Save image and info using Streamlit
+#         #     st.image(
+#         #         env_render, caption=f"Epoch {epoch}", use_column_width=True
+#         #     )
+#         #     st.write(
+#         #         {
+#         #             "state": {i: str(val) for i, val in enumerate(state)},
+#         #             "feature": {
+#         #                 i: str(val) for i, val in enumerate(summed_importances)
+#         #             },
+#         #             "action": str(action),
+#         #             "reward": str(reward),
+#         #             "desired_return": str(desired_return + reward),
+#         #             "desired_horizon": str(desired_horizon + 1),
+#         #         }
+#         #     )
+
+#     env.close()
+
+
+# env = "Acrobot-v1"
+# desired_return = -79
+# desired_horizon = 82
+# max_epoch = 500
+
+
+# policy = SklearnPolicy.load("policy")
+# hyper = AgentHyper(
+#     env,
+#     warm_up=0,
+# )
+
+# agent = UpsideDownAgent(hyper, policy)
+
+# run_visualization(
+#     env, agent, desired_return, desired_horizon, max_epoch, "data/viz_examples"
+# )

pyproject.toml

@@ -26,6 +26,7 @@ python = "3.10.14"
 scikit-learn = "^1.5.2"
 matplotlib = "^3.9.2"
 gymnasium = {extras = ["box2d"], version = "^0.29.1"}
+numpy = "1.24.4"
 scikit-image = "^0.24.0"
 tqdm = "^4.66.5"
 torch = "^2.4.1"

udrl/__main__.py

@@ -1,6 +1,6 @@
-from .agent import UpsideDownAgent, AgentHyper
-from .policies import SklearnPolicy, NeuralPolicy
-from .catch import CatchAdaptor
+from udrl.agent import UpsideDownAgent, AgentHyper
+from udrl.policies import SklearnPolicy, NeuralPolicy
+from udrl.catch import CatchAdaptor
 from dataclasses import dataclass, asdict
 import gymnasium as gym
 from tqdm import trange

udrl/agent.py

@@ -2,9 +2,9 @@ from dataclasses import dataclass
 import gymnasium as gym
 import numpy as np
 
-from .catch import CatchAdaptor
-from .policies import ABCPolicy
-from .buffer import ReplayBuffer
+from udrl.catch import CatchAdaptor
+from udrl.policies import ABCPolicy
+from udrl.buffer import ReplayBuffer
 
 
 @dataclass

udrl/inference.py

@@ -1,7 +1,7 @@
 import matplotlib.pyplot as plt
 import numpy as np
-from .policies import SklearnPolicy, NeuralPolicy
-from .agent import UpsideDownAgent, AgentHyper
+from udrl.policies import SklearnPolicy, NeuralPolicy
+from udrl.agent import UpsideDownAgent, AgentHyper
 from pathlib import Path
 from collections import Counter
 from tqdm import trange

udrl/plot.py

@@ -1,5 +1,5 @@
-from .policies import SklearnPolicy
-from .agent import UpsideDownAgent, AgentHyper
+from udrl.policies import SklearnPolicy
+from udrl.agent import UpsideDownAgent, AgentHyper
 from pathlib import Path
 import matplotlib.pyplot as plt
 import numpy as np

udrl/viz.py

@@ -1,8 +1,8 @@
 import gymnasium as gym
 import pygame
 import numpy as np
-from .policies import SklearnPolicy
-from .agent import UpsideDownAgent, AgentHyper
+from udrl.policies import SklearnPolicy
+from udrl.agent import UpsideDownAgent, AgentHyper
 from pathlib import Path
 import json