source: branches/HeuristicLab.Problems.GrammaticalOptimization/HeuristicLab.Algorithms.Bandits/GrammarPolicies/GenericTDPolicy.cs @ 11832

using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Linq;
using System.Text;
using System.Threading.Tasks;
using HeuristicLab.Algorithms.Bandits.BanditPolicies;
using HeuristicLab.Common;
using HeuristicLab.Problems.GrammaticalOptimization;

namespace HeuristicLab.Algorithms.Bandits.GrammarPolicies {
  public sealed class GenericTDPolicy : IGrammarPolicy {
    private Dictionary<string, IBanditPolicyActionInfo> stateInfo; // stores the necessary information for bandit policies for each state (=canonical phrase)
    private HashSet<string> done;
    private readonly bool useCanonicalPhrases;
    private readonly IProblem problem;
    private readonly IBanditPolicy banditPolicy;

    public GenericTDPolicy(IProblem problem, bool useCanonicalPhrases = false) {
      this.useCanonicalPhrases = useCanonicalPhrases;
      this.problem = problem;
      this.banditPolicy = new EpsGreedyPolicy(0.1);
      this.stateInfo = new Dictionary<string, IBanditPolicyActionInfo>();
      this.done = new HashSet<string>();
    }

    private IBanditPolicyActionInfo[] activeAfterStates; // don't allocate each time
    private int[] actionIndexMap; // don't allocate each time

    public bool TrySelect(Random random, string curState, IEnumerable<string> afterStates, out int selectedStateIdx) {
      // fail if all states are done (corresponding state infos are disabled)
      if (afterStates.All(s => Done(s))) {
        // fail because all follow states have already been visited => also disable the current state (if we can be sure that it has been fully explored)
        MarkAsDone(curState);

        selectedStateIdx = -1;
        return false;
      }

      // determine active actions (not done yet) and create an array to map the selected index back to original actions
      if (activeAfterStates == null || activeAfterStates.Length < afterStates.Count()) {
        activeAfterStates = new IBanditPolicyActionInfo[afterStates.Count()];
        actionIndexMap = new int[afterStates.Count()];
      }
      var idx = 0; int originalIdx = 0;
      foreach (var afterState in afterStates) {
        if (!Done(afterState)) {
          activeAfterStates[idx] = GetStateInfo(afterState);
          actionIndexMap[idx] = originalIdx;
          idx++;
        }
        originalIdx++;
      }

      selectedStateIdx = actionIndexMap[banditPolicy.SelectAction(random, activeAfterStates.Take(idx))];
      
      
      GetStateInfo(curState).UpdateReward(afterStates.Select(s=>GetStateInfo(s).Value).Max());

      return true;
    }



    private IBanditPolicyActionInfo GetStateInfo(string state) {
      var s = CanonicalState(state);
      IBanditPolicyActionInfo info;
      if (!stateInfo.TryGetValue(s, out info)) {
        info = banditPolicy.CreateActionInfo();
        stateInfo[s] = info;
      }
      return info;
    }

    public void UpdateReward(IEnumerable<string> stateTrajectory, double reward) {

      foreach (var state in stateTrajectory.Zip(stateTrajectory.Skip(1), Tuple.Create)) {
        var parent = state.Item1;
        var child = state.Item2;

        // only the last state can be terminal
        if (problem.Grammar.IsTerminal(child)) {
          MarkAsDone(child);
        }
        // GetStateInfo(parent).UpdateReward(GetValue(child));

      }

      GetStateInfo(stateTrajectory.Last()).UpdateReward(reward);
    }


    public void Reset() {
      stateInfo.Clear();
      done.Clear();
    }

    public int GetTries(string state) {
      var s = CanonicalState(state);
      if (stateInfo.ContainsKey(s)) return stateInfo[s].Tries;
      else return 0;
    }

    public double GetValue(string state) {
      var s = CanonicalState(state);
      if (stateInfo.ContainsKey(s)) return stateInfo[s].Value;
      else return 0.0; // TODO: check alternatives
    }

    // the canonical states for the value function (banditInfos) and the done set must be distinguished
    // sequences of different length could have the same canonical representation and can have the same value (banditInfo)
    // however, if the canonical representation of a state is shorter than we must not mark the canonical state as done when all possible derivations from the initial state have been explored
    // eg. in the ant problem the canonical representation for ...lllA is ...rA
    // even though all possible derivations (of limited length) of lllA have been visited we must not mark the state rA as done
    private void MarkAsDone(string state) {
      var s = CanonicalState(state);
      // when the lengths of the canonical string and the original string are the same we also disable the actions
      // always disable terminals
      Debug.Assert(s.Length <= state.Length);
      if (s.Length == state.Length || problem.Grammar.IsTerminal(state)) {
        Debug.Assert(!done.Contains(s));
        done.Add(s);
      } else {
        // for non-terminals where the canonical string is shorter than the original string we can only disable the canonical representation for all states in the same level
        Debug.Assert(!done.Contains(s + state.Length));
        done.Add(s + state.Length); // encode the original length of the state, states in the same level of the tree are treated as equivalent
      }
    }

    // symmetric to MarkDone
    private bool Done(string state) {
      var s = CanonicalState(state);
      if (s.Length == state.Length || problem.Grammar.IsTerminal(state)) {
        return done.Contains(s);
      } else {
        // it is not necessary to visit states if the canonical representation has already been fully explored
        if (done.Contains(s)) return true;
        if (done.Contains(s + state.Length)) return true;
        for (int i = 1; i < state.Length; i++) {
          if (done.Contains(s + i)) return true;
        }
        return false;
      }
    }

    private string CanonicalState(string state) {
      if (useCanonicalPhrases) {
        return problem.CanonicalRepresentation(state);
      } else
        return state;
    }
  }
}