source: branches/HeuristicLab.Problems.GrammaticalOptimization/HeuristicLab.Algorithms.GrammaticalOptimization/MctsSampler.cs @ 11728

using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Linq;
using System.Text;
using HeuristicLab.Algorithms.Bandits;
using HeuristicLab.Problems.GrammaticalOptimization;

namespace HeuristicLab.Algorithms.GrammaticalOptimization {
  public class MctsSampler {
    private class TreeNode {
      public int randomTries;
      public IPolicy policy;
      public TreeNode[] children;
      public bool done = false;

      public override string ToString() {
        return string.Format("Node(random-tries: {0}, done: {1}, policy: {2})", randomTries, done, policy);
      }
    }

    public event Action<string, double> FoundNewBestSolution;
    public event Action<string, double> SolutionEvaluated;

    private readonly int maxLen;
    private readonly IProblem problem;
    private readonly Random random;
    private readonly int randomTries;
    private readonly Func<int, IPolicy> policyFactory;

    private List<Tuple<TreeNode, int>> updateChain;
    private TreeNode rootNode;

    public MctsSampler(IProblem problem, int maxLen, Random random) :
      this(problem, maxLen, random, 10, (numActions) => new EpsGreedyPolicy(random, numActions, 0.1)) {

    }

    public MctsSampler(IProblem problem, int maxLen, Random random, int randomTries, Func<int, IPolicy> policyFactory) {
      this.maxLen = maxLen;
      this.problem = problem;
      this.random = random;
      this.randomTries = randomTries;
      this.policyFactory = policyFactory;
    }

    public void Run(int maxIterations) {
      double bestQuality = double.MinValue;
      InitPolicies();
      for (int i = 0; !rootNode.done && i < maxIterations; i++) {
        var sentence = SampleSentence(problem.Grammar);
        var quality = problem.Evaluate(sentence) / problem.GetBestKnownQuality(maxLen);
        Debug.Assert(quality >= 0 && quality <= 1.0);
        DistributeReward(quality);

        RaiseSolutionEvaluated(sentence, quality);

        if (quality > bestQuality) {
          bestQuality = quality;
          RaiseFoundNewBestSolution(sentence, quality);
        }
      }
    }

    private void InitPolicies() {
      this.updateChain = new List<Tuple<TreeNode, int>>();
      rootNode = new TreeNode();
    }

    private string SampleSentence(IGrammar grammar) {
      updateChain.Clear();
      return CompleteSentence(grammar, grammar.SentenceSymbol.ToString());
    }

    public string CompleteSentence(IGrammar g, string phrase) {
      if (phrase.Length > maxLen) throw new ArgumentException();
      if (g.MinPhraseLength(phrase) > maxLen) throw new ArgumentException();
      TreeNode n = rootNode;
      bool done = phrase.All(g.IsTerminal); // terminal phrase means we are done
      int selectedAltIdx = -1;
      while (!done) {
        int ntIdx; char nt;
        Grammar.FindFirstNonTerminal(g, phrase, out nt, out ntIdx);

        int maxLenOfReplacement = maxLen - (phrase.Length - 1); // replacing aAb with maxLen 4 means we can only use alternatives with a minPhraseLen <= 2
        Debug.Assert(maxLenOfReplacement > 0);

        var alts = g.GetAlternatives(nt).Where(alt => g.MinPhraseLength(alt) <= maxLenOfReplacement);

        if (n.randomTries < randomTries) {
          n.randomTries++;
          return g.CompleteSentenceRandomly(random, phrase, maxLen);
        } else if (n.randomTries == randomTries && n.policy == null) {
          n.policy = policyFactory(alts.Count());
          n.children = alts.Select(_ => new TreeNode()).ToArray(); // create a new node for each alternative
        }

        // => select using bandit policy
        selectedAltIdx = n.policy.SelectAction();
        string selectedAlt = alts.ElementAt(selectedAltIdx);
        // replace nt with alt
        phrase = phrase.Remove(ntIdx, 1);
        phrase = phrase.Insert(ntIdx, selectedAlt);

        updateChain.Add(Tuple.Create(n, selectedAltIdx));

        done = phrase.All(g.IsTerminal); // terminal phrase means we are done
        if (!done) {
          // prepare for next iteration
          n = n.children[selectedAltIdx];
          Debug.Assert(!n.done);
        }
      } // while

      // the last node is a leaf node (sentence is done), so we never need to visit this node again
      n.children[selectedAltIdx].done = true;

      return phrase;
    }

    private void DistributeReward(double reward) {
      // iterate in reverse order (bottom up)
      updateChain.Reverse();

      foreach (var e in updateChain) {
        var node = e.Item1;
        var policy = node.policy;
        var action = e.Item2;
        policy.UpdateReward(action, reward);

        if (node.children[action].done) node.policy.DisableAction(action);
        if (node.children.All(c => c.done)) node.done = true;
      }
    }

    private void RaiseSolutionEvaluated(string sentence, double quality) {
      var handler = SolutionEvaluated;
      if (handler != null) handler(sentence, quality);
    }
    private void RaiseFoundNewBestSolution(string sentence, double quality) {
      var handler = FoundNewBestSolution;
      if (handler != null) handler(sentence, quality);
    }
  }
}