From fcc33d7bdcac834238ed361149892d060941d117 Mon Sep 17 00:00:00 2001
From: Dominik Roth <dominik.roth.dev@gmail.com>
Date: Sun, 19 Jun 2022 15:50:54 +0200
Subject: [PATCH] Work on Columbus integration

---
 test.py | 11 ++++++++++-
 1 file changed, 10 insertions(+), 1 deletion(-)

diff --git a/test.py b/test.py
index dbe7a53..f01ffec 100644
--- a/test.py
+++ b/test.py
@@ -1,4 +1,5 @@
 import gym
+from gym.envs.registration import register
 import numpy as np
 import time
 
@@ -6,9 +7,17 @@ from stable_baselines3 import SAC, PPO, A2C
 from stable_baselines3.common.evaluation import evaluate_policy
 
 from sb3_trl.trl_pg import TRL_PG
+from subtrees.columbus import env
+
+register(
+    id='Columbus_Test3.1-v0',
+    entry_point=env.ColumbusEnv,
+    max_episode_steps=1000,
+)
 
 def main():
-    env = gym.make("LunarLander-v2")
+    #env = gym.make("LunarLander-v2")
+    env = gym.make("Columbus_test3.1-v0")
     ppo = PPO(
         "MlpPolicy",
         env,