/*******************************************************************************
 * DIANNE  - Framework for distributed artificial neural networks
 * Copyright (C) 2015  iMinds - IBCN - UGent
 *
 * This file is part of DIANNE.
 *
 * DIANNE is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 * Contributors:
 *     Tim Verbelen, Steven Bohez
 *******************************************************************************/
package be.iminds.iot.dianne.rl.agent.strategy;

import java.util.Map;

import be.iminds.iot.dianne.api.nn.NeuralNetwork;
import be.iminds.iot.dianne.api.rl.agent.ActionStrategy;
import be.iminds.iot.dianne.api.rl.agent.AgentProgress;
import be.iminds.iot.dianne.api.rl.environment.Environment;
import be.iminds.iot.dianne.nn.util.DianneConfigHandler;
import be.iminds.iot.dianne.rl.agent.strategy.config.BoltzmannConfig;
import be.iminds.iot.dianne.tensor.ModuleOps;
import be.iminds.iot.dianne.tensor.Tensor;
import be.iminds.iot.dianne.tensor.TensorOps;

/**
 * Discrete sampling ActionStrategy that takes the output of a (discrete) policy network
 * ending with a softmax/logsoftmax module and samples an actual discrete action from it.  
 * @author tverbele
 *
 */
public class DiscreteSamplingActionStrategy implements ActionStrategy {
	
	private BoltzmannConfig config;
	private NeuralNetwork nn;
	
	@Override
	public void setup(Map<String, String> config, Environment env, NeuralNetwork... nns) throws Exception {
		this.config = DianneConfigHandler.getConfig(config, BoltzmannConfig.class);
		this.nn = nns[0];
	}

	@Override
	public AgentProgress processIteration(long i, Tensor state) throws Exception {
		Tensor output = nn.forward(state);
		
		Tensor action = new Tensor(output.size());
		action.fill(0);
		
		if(TensorOps.max(output) < 0){
			// assume logsoftmax output, take exp
			output = TensorOps.exp(output, output);
		}
		
		double s = 0, r = Math.random();
		int a = 0;
		
		while((s += output.get(a)) < r && a < action.size())
			a++;
		
		action.set(1, a);
		
		return new AgentProgress(i, action);
	}

}

@@ -1446,54 +1446,8 @@
 ig;%0A
-import be.iminds.iot.dianne.tensor.ModuleOps;%0A
 impo
@@ -1485,16 +1485,16 @@
 Tensor;%0A
+
 import b
@@ -1822,41 +1822,8 @@
 %7B%0A%09%0A
-%09private BoltzmannConfig config;%0A
 %09pri
@@ -1967,86 +1967,8 @@
 n %7B%0A
-%09%09this.config = DianneConfigHandler.getConfig(config, BoltzmannConfig.class);%0A
 %09%09th
@@ -2200,18 +2200,18 @@
 sorOps.m
-ax
+in
 (output)
@@ -2360,19 +2360,16 @@
  0;%0A
-%09%09%0A
 %09%09while(
 (s +
@@ -2364,16 +2364,37 @@
 %09%09while(
+a %3C output.size() && 
 (s += ou
@@ -2413,30 +2413,10 @@
  %3C r
- && a %3C action.size()
 )
+%7B
 %0A%09%09%09
@@ -2420,16 +2420,20 @@
 %09%09%09a++;%0A
+%09%09%7D%0A
 %09%09%0A%09%09act

fix discrete sampling ActionStrategy