Allow overwritting action by environment

Master.lua

@@ -77,7 +77,7 @@ function Master:train()
   -- Catch CTRL-C to save
   self:catchSigInt()
 
-  local reward, state, terminal = 0, self.env:start(), false
+  local reward, state, terminal, actionTaken = 0, self.env:start(), false, false
 
   -- Set environment and agent to training mode
   self.env:training()
@@ -97,7 +97,11 @@ function Master:train()
     local action = self.agent:observe(reward, state, terminal) -- As results received, learn in training mode
     if not terminal then
       -- Act on environment (to cause transition)
-      reward, state, terminal = self.env:step(action)
+      reward, state, terminal, actionTaken = self.env:step(action)
+      -- Update experience memory with actual action
+      if actionTaken and actionTaken ~= action then
+        self.agent.memory.actions[self.agent.memory.index] = actionTaken
+      end
       -- Track score
       episodeScore = episodeScore + reward
     else

README.md

@@ -54,6 +54,8 @@ You can use a custom environment (as the path to a Lua file/`rlenvs`-namespaced
 
 If the environment has separate behaviour during training and testing it should also implement `training` and `evaluate` methods - otherwise these will be added as empty methods during runtime. The environment can also implement a `getDisplay` method (with a mandatory `getDisplaySpec` method for determining screen size) which will be used for displaying the screen/computing saliency maps, where `getDisplay` must return a RGB (3D) tensor; this can also be utilised even if the state is not an image (although saliency can only be computed for states that are images). This **must** be implemented to have a visual display/computing saliency maps. The `-zoom` factor can be used to increase the size of small displays.
 
+Custom environments can also control the action selection process, specifying the actual action taken when it differs from that selected by the network. This allows the agent to learn from hand-crafted behaviours, human experts or pre-planned sequences. To achieve this environments can optionally return `actionTaken` from the `step` method. i.e. `return reward, state, terminal[, actionTaken]`.
+
 You can also use a custom model (body) with `-modelBody`, which replaces the usual DQN convolutional layers with a custom Torch neural network (as the path to a Lua file/`models`-namespaced environment). The class must include a `createBody` method which returns the custom neural network. The model will receive a stack of the previous states (as determined by `-histLen`), and must reshape them manually if needed. The DQN "heads" will then be constructed as normal, with `-hiddenSize` used to change the size of the fully connected layer if needed.
 
 For an example on a GridWorld environment, run `./run.sh demo-grid` - the demo also works with `qlua` and experience replay agents. The custom environment and network can be found in the [examples](https://github.com/Kaixhin/Atari/tree/master/examples) folder.

async/A3CAgent.lua

@@ -40,6 +40,8 @@ function A3CAgent:learn(steps, from)
   log.info('A3CAgent starting | steps=%d', steps)
   local reward, terminal, state = self:start()
 
+  local actionTaken
+
   self.states:resize(self.batchSize, table.unpack(state:size():totable()))
 
   self.tic = torch.tic()
@@ -55,7 +57,10 @@ function A3CAgent:learn(steps, from)
 
       self.actions[self.batchIdx] = action
 
-      reward, terminal, state = self:takeAction(action)
+      reward, terminal, state, actionTaken = self:takeAction(action)
+      if actionTaken and actionTaken + self.actionOffset ~= action then
+        action = actionTaken + self.actionOffset
+      end
       self.rewards[self.batchIdx] = reward
 
       self:progress(steps)
@@ -98,7 +103,7 @@ function A3CAgent:accumulateGradients(terminal, state)
     local gradEntropy = torch.log(probability) + 1
     -- Add to target to improve exploration (prevent convergence to suboptimal deterministic policy)
     self.policyTarget:add(self.beta, gradEntropy)
-    
+
     self.policyNet_:backward(self.states[i], self.targets)
   end
 end

async/AsyncAgent.lua

@@ -78,7 +78,7 @@ end
 
 
 function AsyncAgent:takeAction(action)
-  local reward, rawObservation, terminal = self.env:step(action - self.actionOffset)
+  local reward, rawObservation, terminal, actionTaken = self.env:step(action - self.actionOffset)
   if self.rewardClip > 0 then
     reward = math.max(reward, -self.rewardClip)
     reward = math.min(reward, self.rewardClip)
@@ -91,7 +91,7 @@ function AsyncAgent:takeAction(action)
     self.stateBuffer:push(observation)
   end
 
-  return reward, terminal, self.stateBuffer:readAll()
+  return reward, terminal, self.stateBuffer:readAll(), actionTaken
 end
 
 

async/NStepQAgent.lua

@@ -32,6 +32,8 @@ function NStepQAgent:learn(steps, from)
   log.info('NStepQAgent starting | steps=%d | ε=%.2f -> %.2f', steps, self.epsilon, self.epsilonEnd)
   local reward, terminal, state = self:start()
 
+  local actionTaken
+
   self.states:resize(self.batchSize, table.unpack(state:size():totable()))
   self.tic = torch.tic()
   repeat
@@ -44,7 +46,10 @@ function NStepQAgent:learn(steps, from)
       local action = self:eGreedy(state, self.policyNet_)
       self.actions[self.batchIdx] = action
 
-      reward, terminal, state = self:takeAction(action)
+      reward, terminal, state, actionTaken = self:takeAction(action)
+      if actionTaken and actionTaken + self.actionOffset ~= action then
+        action = actionTaken + self.actionOffset
+      end
       self.rewards[self.batchIdx] = reward
 
       self:progress(steps)

async/OneStepQAgent.lua

@@ -24,13 +24,16 @@ function OneStepQAgent:learn(steps, from)
   log.info('%s starting | steps=%d | ε=%.2f -> %.2f', self.agentName, steps, self.epsilon, self.epsilonEnd)
   local reward, terminal, state = self:start()
 
-  local action, state_
+  local action, state_, actionTaken
 
   self.tic = torch.tic()
   for step1=1,steps do
     if not terminal then
       action = self:eGreedy(state, self.policyNet)
-      reward, terminal, state_ = self:takeAction(action)
+      reward, terminal, state_, actionTaken = self:takeAction(action)
+      if actionTaken and actionTaken + self.actionOffset ~= action then
+        action = actionTaken + self.actionOffset
+      end
     else
       reward, terminal, state_ = self:start()
     end