Spaces:
Running
on
CPU Upgrade
Running
on
CPU Upgrade
import gym | |
from gym.wrappers import RecordVideo | |
from matplotlib import pyplot as plt | |
from interpretable.interpretable import InterpretablePolicyExtractor | |
from interpretable.utils import generate_dataset_from_expert, rollouts | |
if __name__ == "__main__": | |
env_name = "CartPole-v1" | |
dataset_path = generate_dataset_from_expert("ppo", env_name, force=True) | |
ipe = InterpretablePolicyExtractor(env_name) | |
results = ipe.train_from_dataset(dataset_path) | |
ipe.policy.prune() | |
ipe.policy.plot(mask=True) | |
plt.savefig("kan-policy.png") | |
env = gym.make(env_name, render_mode="rgb_array") | |
env = RecordVideo(env, video_folder="videos", episode_trigger=lambda x: True, name_prefix=f"kan-{env_name}") | |
ipe.policy.auto_symbolic() | |
ipe.policy.plot(mask=True) | |
plt.savefig("sym-policy.png") | |
print(ipe.policy.symbolic_formula()) | |
rollouts(env, ipe.forward, 2) | |