makeGridWorldPOMDP.wppl

// function that builds a gridworld POMDP

// arg is a gridworld MDP, which is an object with features (an array of
// arrays of the features at each point), xLim, yLim, feature (a function of
// a (manifest) state that returns the features of that state), transition
// (a transition function on manifest states), actions, and stateToActions.
// noReverse must be false in the MDP.

// manifest states in this POMDP are objects with loc (an array of two
// numbers), timeLeft, terminateAfterAction, and timeAtRestraurant - the same as states in gridworld MDPs.
// latent states are an object listing whether each restaurant type is open.
// transition acts the same as in gridworld MDPs, not touching the latent state,
// unless the transition would take the agent into a closed restaurant, in which
// case the agent stays put.
// observe tells the agent which adjacent restaurants are open. If the agent is not
// adjacent to any restaurants, they get 'noObservation'.


var makeGridWorldPOMDP = function(options) {

  var gridworld = makeGridWorldMDP(options).world;

  var actions = gridworld.actions;
  var feature = gridworld.feature;

  // returns an array of restaurants that neighbour a state
  var neighbourRestaurants = function(manifestState) {
    var loc = manifestState.loc;
    var updateStateLoc = function(manifestState, newLoc) {
      return extend(manifestState, {
        loc: newLoc
      });
    };

    var potentialneighbourStates = [
      updateStateLoc(manifestState, [loc[0] - 1, loc[1]]),
      updateStateLoc(manifestState, [loc[0] + 1, loc[1]]),
      updateStateLoc(manifestState, [loc[0], loc[1] + 1]),
      updateStateLoc(manifestState, [loc[0], loc[1] - 1]),
    ];

    var neighbourStates = filter(function(manifestState) {
      return inGrid_(gridworld, manifestState.loc);
    }, potentialneighbourStates);

    var isRestaurant = function(manifestState) {
      return feature(manifestState).name;
    };

    return filter(isRestaurant, neighbourStates);
  };

  // returns a subobject of latentState that gives the status of adjacent restaurants
  var observe = function(state) {
    var neighbourRests = neighbourRestaurants(state.manifestState);
    if (_.isEmpty(neighbourRests)) {
      return 'noObservation';
    } else {
      var restaurantNames = map(function(manifestState) {
        return feature(manifestState).name;
      }, neighbourRests);

      return map(function(name) {
          return [name, state.latentState[name]];
        },
        restaurantNames);
    }
  };

  var _mdpTransition = gridworld.transition;
  var mdpTransition = function(state, action) {
    // transition the manifest state as it would be transitioned in the gridworld
    // mdp, leave the latent state alone
    var newManifestState = _mdpTransition(state.manifestState, action);
    return {
      manifestState: newManifestState,
      latentState: state.latentState
    };
  };

  var transition = function(state, action) {
    assert.ok(stateHasManifestLatent(state), 'transition state arg');

    var proposedNewState = mdpTransition(state, action);
    var newFeatureName = feature(proposedNewState.manifestState).name;

    // if proposedNewState is a restaurant that is closed, stay put, but increment
    // time. otherwise, change to proposed new state.
    if (newFeatureName && !proposedNewState.latentState[newFeatureName]) {
      return {
        manifestState: advanceStateTime(state.manifestState),
        latentState: state.latentState
      };
    } else {
      return proposedNewState;
    }
  };

  var manifestStateToActions = function(manifestState) {
    var neighbourRestaurantArray = neighbourRestaurants(manifestState);
    if (neighbourRestaurantArray.length == 0) {
      var manifestStateToActions_ = gridworld.stateToActions;
      return manifestStateToActions_(manifestState);
    } else {
      return gridworld.actions;
    }
  };

  return {
    manifestStateToActions: manifestStateToActions,
    transition: transition,
    observe: observe,
    feature: feature,
    MDPWorld: gridworld
  };
};


var inferGridWorldPOMDP = function(world, startState, baseParams, trueAgentParams, prior, agentTypeAndFunctions, trajectoryOrOffPolicy, numRejectionSamples) {

  var makeAgent = agentTypeAndFunctions.makeAgent;
  var simulate = agentTypeAndFunctions.simulate;
  var beliefOrBeliefDelay = agentTypeAndFunctions.type;

  var isJointPrior = prior.isJoint;

  // try just not defining things instead of having them as 'undefined'
  var jointPrior = isJointPrior ? prior.jointPrior : undefined;
  var priorUtilityTable = isJointPrior ? undefined : prior.priorUtilityTable;
  var priorAgentPrior = isJointPrior ? undefined : prior.priorAgentPrior;

  if (isJointPrior) {
    assert.ok(stateHasManifestLatent(startState) && jointPrior.sample,
      'inferGridworld args');
  } else {
    assert.ok(stateHasManifestLatent(startState) &&
      priorUtilityTable.sample &&
      priorAgentPrior.sample,
      'inferGridworld args');
  }

  assert.ok(trajectoryOrOffPolicy == 'trajectory' || trajectoryOrOffPolicy == 'offPolicy',
    'trajectoryOrOffPolicy is not valid');

  assert.ok(_.isNumber(numRejectionSamples) && isPOMDPWithManifestLatent(world), 'inferGridWorld args');

  // get observations using trueAgentParams
  var agent = makeAgent(trueAgentParams, world);
  var observedStateAction = simulatePOMDP(startState, world, agent, 'stateAction');
  var observe = agent.POMDPFunctions.observe;

  assert.ok(stateHasManifestLatent(observedStateAction[0][0]), 'fullstate in trajectory for inferGridWorld');

  return Infer({
    method: 'enumerate'
  }, function() {
    var utilityAndPrior = isJointPrior ? sample(jointPrior) : null;

    var utilityTable = isJointPrior ? utilityAndPrior.utilityTable :
      sample(priorUtilityTable);
    var utility = makeRestaurantUtilityFunction(world, utilityTable);

    var priorBelief = isJointPrior ? utilityAndPrior.priorBelief :
      sample(priorAgentPrior);

    var params = extend(baseParams, {
      utility: utility,
      priorBelief: priorBelief
    });

    var agent = makeAgent(params, world);
    var agentAct = agent.act;
    var agentUpdateBelief = agent.updateBelief;

    // Factor on whole sampled trajectory (SLOW IF TRANSITIONS NOT DETERMINISTIC OR IF NUM SAMPLES HIGH)
    var factorOnTrajectory = function() {
      var trajectoryDist = Infer({
        method: 'rejection',
        samples: numRejectionSamples
      }, function() {
        return simulatePOMDP(startState, world, agent, 'states')
      });
      factor(trajectoryDist.score(map(first, observedStateAction)));
    };

    // Move agent through observed sequence
    var factorSequenceOffPolicy = function(currentBelief, previousAction, timeIndex) {
      if (timeIndex < observedStateAction.length) {

        // Go to next world state and sample observation from that state
        var state = observedStateAction[timeIndex][0];
        var observation = observe(state);

        // Update agent's internal state and get action Dist
        var delay = 0;
        var nextBelief = beliefOrBeliefDelay == 'belief' ?
          agentUpdateBelief(currentBelief, observation, previousAction) :
          agentUpdateBelief(currentBelief, observation, previousAction, delay);


        var nextActionDist = beliefOrBeliefDelay == 'belief' ?
          agentAct(nextBelief) : agentAct(nextBelief, delay);

        var observedAction = observedStateAction[timeIndex][1];

        factor(nextActionDist.score(observedAction));

        // condition on next world state, passing through updated internal state
        factorSequenceOffPolicy(nextBelief, observedAction, timeIndex + 1);
      }
    };

    var doInfer = (trajectoryOrOffPolicy == 'trajectory') ? factorOnTrajectory() :
      factorSequenceOffPolicy(priorBelief, 'noAction', 0);

    return { utilityTable, priorBelief };
  });
};