Modify envs to be compatible with twrl

experiment.lua

@@ -1,11 +1,8 @@
-local image = require 'image'
 require 'rlenvs'
 local Catch = require('rlenvs.Catch')
--- Detect QT for image display
-local qt = pcall(require, 'qt')
 
 -- Initialise and start environment
-local env = Catch({level = 2})
+local env = Catch({level = 2, render = true})
 local getActionSpace = env:getActionSpace()
 local observation = env:start()
 
@@ -14,7 +11,7 @@ local episodes, totalReward = 0, 0
 local nEpisodes = 1000
 
 -- Display
-local window = qt and image.display({image=observation, zoom=10})
+env:render()
 
 for i = 1, nEpisodes do
   while not terminal do
@@ -24,9 +21,7 @@ for i = 1, nEpisodes do
     totalReward = totalReward + reward
 
     -- Display
-    if qt then
-        image.display({image=observation, zoom=10, win=window})
-    end
+    env:render()
   end
 
   episodes = episodes + 1

rlenvs/Env.lua

@@ -19,6 +19,10 @@ function Env:_init(opts)
     else
         self.maxSteps = 1000
     end
+    if opts.render then
+        require 'image' self.qt = pcall(require, 'qt')
+        if not self.qt then print('Was not able to load qt to render, are you using qlua to run the script?') end
+    end
     self.currentStep = 1
 end
 
@@ -34,7 +38,15 @@ end
 
 function Env:start()
     self.currentStep = 1
-    return self:_start()
+    local obs = self:_start()
+    return obs
+end
+
+function Env:render()
+    if self.qt and self.getDisplay then
+        self.window = self.window == nil and image.display({ image = self:getDisplay(), zoom = 10 }) or self.window
+        image.display({ image = self:getDisplay(), zoom = 10, win = self.window })
+    end
 end
 
 return Env