makeGridWorldMDP.wppl

// TODO: looks like if you have named features in a gridworld, then
// they have to be terminal (or at terminal after *timeAtRestaurant*
// is exceeded). we should change this so that features can be
// non-terminal.


var inGrid_ = function(gridMap, loc) {
  return (loc[0] >= 0 && loc[0] < gridMap.xLim &&
    loc[1] >= 0 && loc[1] < gridMap.yLim);
}


var isBlockedLoc = function(gridMap, loc) {
  var getFeature = gridMap.feature;
  var feature = getFeature({
    loc: loc
  });
  return feature === '#';
}


var isAllowedState_ = function(gridMap, loc) {
  return (inGrid_(gridMap, loc)) && !isBlockedLoc(gridMap, loc);
};


var advanceStateTime = function(state) {
  return extend(state, {
    timeLeft: state.timeLeft - 1,
    terminateAfterAction: state.timeLeft - 1 > 1 ? state.terminateAfterAction : true,
    previousLoc: state.loc
  });
};


var addPrevious = function(state) {
  return extend(state, {
    previousLoc: state.loc
  });
};


var advanceRestaurant = function(state, maxTimeAtRestaurant) {
  var timeAtRestaurant = ((state.timeAtRestaurant === undefined) ? 0 :
                          state.timeAtRestaurant + 1);
  var state = extend(state, { timeAtRestaurant });
  if (state.timeAtRestaurant < maxTimeAtRestaurant - 1) {
    return state;
  } else {
    return extend(state, { terminateAfterAction: true });
  }
};


var moveState = function(gridMap, state, action) {
  var loc = state.loc;
  var gridTransition = {
    l: [loc[0] - 1, loc[1]],
    r: [loc[0] + 1, loc[1]],
    u: [loc[0], loc[1] + 1],
    d: [loc[0], loc[1] - 1]
  };

  var possibleNextLoc = gridTransition[action];
  assert.ok(possibleNextLoc != undefined, 'action was not found');

  var nextLoc = isAllowedState_(gridMap, possibleNextLoc) ? possibleNextLoc : loc;
  return extend(state, {
    loc: nextLoc
  });
};


var makeGridTransition_ = function(gridMap, options) {
  return function(state, action) {
    var getFeature = gridMap.feature;
    var state = options.noReverse ? addPrevious(state) : state;
    var state = !getFeature(state).name ? advanceStateTime(state) : state;

    var state = !getFeature(state).name ? moveState(gridMap, state, action) : state;
    var state = getFeature(state).name ? advanceRestaurant(state, options.maxTimeAtRestaurant) : state;

    return state;
  };
};


var makeNoisyGridTransition_ = function(gridMap, options, transitionNoiseProb) {
  var deterministicTransition = makeGridTransition_(gridMap, options);
  // If agent selects *key*, with *transitionNoiseProb* they do one of the two
  // orthogonal actions
  var noiseActionTable = {
    u: ['l', 'r'],
    d: ['l', 'r'],
    l: ['u', 'd'],
    r: ['u', 'd']
  };

  return function(state, action) {
    return flip(1 - transitionNoiseProb) ? deterministicTransition(state, action) :
      deterministicTransition(state, uniformDraw(noiseActionTable[action]));
  };
};


var makeGridMap = function(rfeatures) {
  var features = rfeatures.reverse();
  return {
    features: features,
    xLim: features[0].length,
    yLim: features.length,
    feature: function(state) {
      return features[state.loc[1]][state.loc[0]];
    }
  };
};


var makeGridWorldDeterministic = function(features, options) {
  var defaults = {
    noReverse: false,
    maxTimeAtRestaurant: 1
  };
  var options = extend(defaults, options);

  var gridMap = makeGridMap(features);

  var transition = makeGridTransition_(gridMap, options);
  var actions = ['l', 'r', 'u', 'd'];

  var stateToActions = function(state) {

    var possibleActions = filter(function(action) {
      var newState = transition(state, action);
      if (options.noReverse &&
        state.previousLoc &&
        state.previousLoc[0] === newState.loc[0] && state.previousLoc[1] === newState.loc[1]) {
        return false;
      }
      return state.loc[0] !== newState.loc[0] || state.loc[1] !== newState.loc[1];
    }, actions);

    if (possibleActions.length > 0) {
      return possibleActions;
    } else {
      return [actions[0]];
    }
  };

  return extend(gridMap, { transition, actions, stateToActions });
};


var makeNoisyGridworld = function(features, inputOptions, transitionNoiseProb) {
  var defaults = {
    noReverse: false,
    maxTimeAtRestaurant: 1
  };

  var options = extend(defaults, inputOptions);
  var gridMap = makeGridMap(features);

  var transition = makeNoisyGridTransition_(gridMap, options,
    transitionNoiseProb);
  var actions = ['l', 'r', 'u', 'd'];

  var stateToActions = function(state) {
    var possibleActions = filter(function(action) {
      // if noReverse, check that MAP of newState doesn't reverse you
      var transitionDist = Infer({
        method: 'enumerate'
      }, function() {
        return transition(state, action);
      });
      var likelyNewState = MAP(transitionDist).val;
      if (options.noReverse && state.previousLoc &&
        state.previousLoc[0] === likelyNewState.loc[0] &&
        state.previousLoc[1] === likelyNewState.loc[1]) {
        return false;
      } else {
        return state.loc[0] !== likelyNewState.loc[0] ||
          state.loc[1] !== likelyNewState.loc[1];
      }
    }, actions);


    if (possibleActions.length > 0) {
      return possibleActions;
    } else {
      return [actions[0]];
    }
  };

  return extend(gridMap, { transition, actions, stateToActions });
};


var makeGridWorldMDP = function(options) {

  var defaults = {
    noReverse: false,
    maxTimeAtRestaurant: 1,
    transitionNoiseProbability: 0,
    start: [0, 0],
    totalTime: 2
  };

  var options = extend(defaults, options);

  map(function(el) {
    assert.ok(_.isString(el) || _.isString(el.name), 'grid is invalid')
  }, _.flatten(options.grid));

  assert.ok(options.totalTime >= 1,
            'totalTime is invalid');
  assert.ok(options.transitionNoiseProbability <= 1,
            'transitionNoiseProbability is invalid');
  assert.ok(options.start.length == 2,
            'start is invalid');

  var world = (
    (options.transitionNoiseProbability == 0) ?
    makeGridWorldDeterministic(options.grid, options) :
    makeNoisyGridworld(options.grid, options, options.transitionNoiseProbability));

  var startState = {
    loc: options.start,
    terminateAfterAction: false,
    timeLeft: options.totalTime
  };

  // assumes that utilities are only defined on named features and timeCost
  // doesnt work for the hyperbolic case
  var makeUtilityFunction = function(utilityTable) {
    assert.ok(utilityTable.hasOwnProperty('timeCost'),
              'makeUtilityFunction utilityTable lacks timeCost');
    return function(state, action) {
      var getFeature = world.feature;
      var feature = getFeature(state);

      if (feature.name) {
        return utilityTable[feature.name];
      }
      return utilityTable.timeCost;
    };
  };

  return { world, startState, makeUtilityFunction };

};