MOReinforceNormal.lua

--
-- Created by IntelliJ IDEA.
-- User: petrfiala
-- Date: 17/02/2016
-- Time: 14:09
-- To change this template use File | Settings | File Templates.
--

------------------------------------------------------------------------
--[[ ReinforceNormal ]] --
-- Ref A. http://incompleteideas.net/sutton/williams-92.pdf
-- Inputs are mean (mu) of multivariate normal distribution.
-- Ouputs are samples drawn from these distributions.
-- Standard deviation is provided as constructor argument.
-- Uses the REINFORCE algorithm (ref. A sec 6. p.237-239) which is
-- implemented through the nn.Module:reinforce(r,b) interface.
-- gradOutputs are ignored (REINFORCE algorithm).
------------------------------------------------------------------------
local MOReinforceNormal, parent = torch.class("MOReinforceNormal", "MOReinforce")

function MOReinforceNormal:__init(stdev, stochastic)
    parent.__init(self, stochastic)
    self.stdev = stdev
    if not stdev then
        self.gradInput = { torch.Tensor(), torch.Tensor() }
    end
    -- TODO
    reinforce_step = 0
end

function MOReinforceNormal:updateOutput(input)
    local mean, stdev = input, self.stdev
    if torch.type(input) == 'table' then
        -- input is {mean, stdev}
        assert(#input == 2)
        mean, stdev = unpack(input)
    end
    assert(stdev)

    self.output:resizeAs(mean)

    if self.stochastic or self.train ~= false then
        self.output:normal()
        -- multiply by standard deviations
        if torch.type(stdev) == 'number' then
            self.output:mul(stdev)
        elseif torch.isTensor(stdev) then
            if stdev:dim() == mean:dim() then
                assert(stdev:isSameSizeAs(mean))
                self.output:cmul(stdev)
            else
                assert(stdev:dim() + 1 == mean:dim())
                self._stdev = self._stdev or stdev.new()
                self._stdev:view(stdev, 1, table.unpack(stdev:size():totable()))
                self.__stdev = self.__stdev or stdev.new()
                self.__stdev:expandAs(self._stdev, mean)
                self.output:cmul(self.__stdev)
            end
        else
            error "unsupported mean type"
        end
        if self.train then
            reinforce_step = reinforce_step + 1
        elseif self.train == nil then
            error "model.train has to be  defined!"
        end

        -- re-center the means to the mean
        self.output:add(mean)
    else
        -- use maximum a posteriori (MAP) estimate
        self.output:copy(mean)
    end
    return self.output
end

function MOReinforceNormal:updateGradInput(input, gradOutput)
    -- Note that gradOutput is ignored
    -- f : normal probability density function
    -- x : the sampled values (self.output)
    -- u : mean (mu) (mean)
    -- s : standard deviation (sigma) (stdev)

    local mean, stdev = input, self.stdev
    local gradMean, gradStdev = self.gradInput, nil
    if torch.type(input) == 'table' then
        mean, stdev = unpack(input)
        gradMean, gradStdev = unpack(self.gradInput)
    end
    assert(stdev)

    -- Derivative of log normal w.r.t. mean :
    -- d ln(f(x,u,s))   (x - u)
    -- -------------- = -------
    --      d u           s^2

    gradMean:resizeAs(mean)
    -- (x - u)
    gradMean:copy(self.output):add(-1, mean)

    -- divide by squared standard deviations
    if torch.type(stdev) == 'number' then
        gradMean:div(stdev ^ 2)
    else
        if stdev:dim() == mean:dim() then
            gradMean:cdiv(stdev):cdiv(stdev)
        else
            gradMean:cdiv(self.__stdev):cdiv(self.__stdev)
        end
    end
    -- multiply by reward
    gradMean:cmul(self:rewardAs(mean, reinforce_step))
    reinforce_step = reinforce_step - 1

    -- multiply by -1 ( gradient descent on mean )
    gradMean:mul(-1)

    -- Derivative of log normal w.r.t. stdev :
    -- d ln(f(x,u,s))   (x - u)^2 - s^2
    -- -------------- = ---------------
    --      d s              s^3

    if gradStdev then
        gradStdev:resizeAs(stdev)
        -- (x - u)^2
        gradStdev:copy(self.output):add(-1, mean):pow(2)
        -- subtract s^2
        self._stdev2 = self._stdev2 or stdev.new()
        self._stdev2:resizeAs(stdev):copy(stdev):cmul(stdev)
        gradStdev:add(-1, self._stdev2)
        -- divide by s^3
        self._stdev2:cmul(stdev):add(0.00000001)
        gradStdev:cdiv(self._stdev2)
        -- multiply by reward
        gradStdev:cmul(self:rewardAs(stdev))
        -- multiply by -1 ( gradient descent on stdev )
        gradStdev:mul(-1)
    end

    return self.gradInput
end