AbstractRecurrent:reinforce() handles multiple rewards

AbstractRecurrent.lua

@@ -155,9 +155,19 @@ function AbstractRecurrent:evaluate()
 end
 
 function AbstractRecurrent:reinforce(reward)
-   return self:includingSharedClones(function()
-      return parent.reinforce(self, reward)
-   end)
+   if torch.type(reward) == 'table' then
+      -- multiple rewards, one per time-step
+      local rewards = reward
+      for step, reward in ipairs(rewards) do
+         local sm = self:getStepModule(step)
+         sm:reinforce(reward)
+      end
+   else
+      -- one reward broadcast to all time-steps
+      return self:includingSharedClones(function()
+         return parent.reinforce(self, reward)
+      end)
+   end
 end
 
 -- used by Recursor() after calling stepClone.

Recurrent.lua

@@ -158,6 +158,26 @@ function Recurrent:includingSharedClones(f)
    return r
 end
 
+function Recurrent:reinforce(reward)
+   if torch.type(reward) == 'table' then
+      -- multiple rewards, one per time-step
+      local rewards = reward
+      for step, reward in ipairs(rewards) do
+         if step == 1 then
+            self.initialModule:reinforce(reward)
+         else
+            local sm = self:getStepModule(step)
+            sm:reinforce(reward)
+         end
+      end
+   else
+      -- one reward broadcast to all time-steps
+      return self:includingSharedClones(function()
+         return parent.reinforce(self, reward)
+      end)
+   end
+end
+
 function Recurrent:maskZero()
    error("Recurrent doesn't support maskZero as it uses a different "..
       "module for the first time-step. Use nn.Recurrence instead.")

test/test.lua

@@ -4047,6 +4047,23 @@ function rnntest.encoderdecoder()
    numgradtest()
 end
 
+function rnntest.reinforce()
+   -- test that AbstractRecurrent:reinforce(rewards) words
+   local seqLen = 4
+   local batchSize = 3
+   local rewards = {}
+   for i=1,seqLen do
+      rewards[i] = torch.randn(batchSize)
+   end
+   local rf = nn.ReinforceNormal(0.1)
+   local rnn = nn.Recursor(rf)
+   rnn:reinforce(rewards)
+   for i=1,seqLen do
+      local rm = rnn:getStepModule(i)
+      mytester:assertTensorEq(rm.reward, rewards[i], 0.000001, "Reinforce error")
+   end
+end
+
 function rnn.test(tests, benchmark_)
    mytester = torch.Tester()
    benchmark = benchmark_