OptimizedTransitionCalculator.java

public class OptimizedTransitionCalculator {
    private static final Logger logger = LoggerFactory.getLogger(
        OptimizedTransitionCalculator.class
    );

    public Policy generateNewPolicy(
        Set<Action> actions,
        Map<State, Double> expectedUtilities,
        Map<StateAction, ? extends Transition> transitions) {
        Map.Entry<State, Double>[] expectedUtilityEntries = expectedUtilities
            .entrySet()
            .toArray(new Map.Entry[0]);
        Action[] actionArray = actions.toArray(new Action[0]);
        return new Policy(
            expectedUtilities.keySet().parallelStream().collect(
                toConcurrentMap(
                    (state) -> state,
                    (state) -> maxAction(
                        state, actionArray, transitions, expectedUtilityEntries
                    )
                )
            )
        );
    }

    /**
     * This function is the most used function in this program.
     * Every cycle saved here counts for a lot. Optimizations are welcome.
     *
     * I have smashed all subroutines into this function. I've commented the
     * delineation.
     *
     * @param state
     * @param actions
     * @param transitions
     * @param expectedUtilities
     * @return
     */
    private Action maxAction(
        final State state,
        final Action[] actions,
        final Map<StateAction, ? extends Transition> transitions,
        final Map.Entry<State, Double>[] expectedUtilities) {
        double maxExpectedUtility = Double.NEGATIVE_INFINITY;
        Action maxAction = null;
        for (Action act : actions) {
            // Calculate the action's expected utility in this state.
            Transition t = transitions.get(new StateAction(state, act));
            double expectedUtility = 0.0;
            for (Map.Entry<State, Double> entry : expectedUtilities) {
                expectedUtility += t.getStateProbability(entry.getKey()) * entry.getValue();
            }

            // Update maximal expected utility
            if (expectedUtility > maxExpectedUtility) {
                maxAction = act;
                maxExpectedUtility = expectedUtility;
            }
        }
        return maxAction;
    }


}