questionsRSA.wppl

var questionStartSymbol = "$WH";
var answerStartSymbol = "$S";

var KL = function(p, q){
    return expectation(p, function(value) {
        var scoreP = p.score(value);
        var scoreQ = q.score(value);
        return scoreP === -Infinity ? 0 : scoreP - scoreQ;
    });
};

var utteranceDist = function(utterances, costFn, rationality) {
    var probabilities = map(function(u) {
        return Math.exp(-rationality * costFn(u));
    }, utterances);
    return Categorical({ ps: probabilities, vs: utterances });
}

var createVariationalModel = function(questions, answers, costFn, worldPrior, qudCandidates, 
                                      questionerRationality, answererRationality, grammarFn, 
                                      parametrizedNetworks, observationFn, opts, experimentName, runId, relearnParams,
                                      preTrainModel) {
    assert.ok(_.isNumber(answererRationality), "Pass in answererRationality")
    var runPath = makeRunPath(experimentName, runId);
    var modelFn = function() {
        var specifiedNetworks = mapObject(function(k, v) {
            var params = apply(v.paramConstructor, [k]);
            return v.runner(params);
        }, parametrizedNetworks)
        var grammar = grammarFn.call(null, specifiedNetworks);
        var parser = createParser(grammar, createParserWeights(grammar));
        var model = createModel(questions, answers, costFn, worldPrior, qudCandidates, 
                                questionerRationality, answererRationality, parser);
        if (observationFn) {
            observationFn(model);
        }
        return _.assign({ networks: specifiedNetworks }, model);
    }

    if (preTrainModel) {
        preTrainModel(sample(SampleGuide(modelFn)));
    }

    if (!relearnParams && fs.existsSync(runPath.paramsPath)) {
        // Restore previously learned params
        var params = webpplFs.read(runPath.paramsPath);
        setParams(deserializeParams(params));
    }
    else {
        // learn new params
        var optsWithFilenames = _.assign({ 
            model: modelFn,
            checkpointParamsFilename: runPath.paramsPath,
            logProgressFilename: runPath.logPath
        }, opts);
        Optimize(optsWithFilenames);
    }
                                     
    return sample(SampleGuide(modelFn));
}

var createModel = function(questions, answers, costFn, worldPrior, qudCandidates,
    questionerRationality, answererRationality, parser) {
    assert.ok(_.isNumber(answererRationality), "Pass in answererRationality")
    var questionsDist = utteranceDist(questions, costFn, questionerRationality);
    var answersDist = utteranceDist(answers, costFn, answererRationality);

    var marginalizeCached = cache(function(dist, qudName) {
        var qud = qudCandidates[qudName];
        return marginalize(dist, qud);
    })

    var interpretAnswer = cache(function(answer) {
        var answerMeaning = parser(answer, answerStartSymbol);
        return Infer({method: 'enumerate'}, function() {
            var world = sample(worldPrior);
            factor(Math.log(answerMeaning(world)));
            return world;
        });
    })

    var informationGainFn = function(answerer) {
        var interpretAnswerPragmatic = cache(function(question, answer) {
            return Infer({ method: 'enumerate' }, function() {
                var world = sample(worldPrior);
                observe(answerer(question, world), answer);
                return world;
            })
        })

        return cache(function(question, qudName, trueWorld) {
            var prior = marginalizeCached(worldPrior, qudName);
            var possibleAnswers = answerer(question, trueWorld);
            return expectation(possibleAnswers, function(answer) {
                var posterior = marginalizeCached(interpretAnswerPragmatic(question, answer), qudName);
                return KL(posterior, prior);
            });
        })
    }
    
    var similarity = cache(function(question, trueWorld, otherWorld, domain) {
        // trueWorldTruthValues and otherWorldTruthValues are each distributions
        // over subsets of the domain. Compute KL between them using factorization
        // for sequence of independent bernoulli trials.
        var questionMeaning = parser(question, questionStartSymbol);
        var trueWorldTruthValues = questionMeaning(trueWorld);
        var otherWorldTruthValues = questionMeaning(otherWorld);
        return Math.exp(-sum(map2(function(p, q) {
            var positiveTerm = p === 0 ? 0 : p * Math.log(p / q);
            var negativeTerm = p === 1 ? 0 : (1 - p) * Math.log((1 - p) / (1 - q));
            return positiveTerm + negativeTerm;
        }, trueWorldTruthValues, otherWorldTruthValues)));
    })

    var answerQualityForQuestion = cache(function(answer, question, trueWorld) {
        // Using soft similarity semantics
        var consistentWorlds = interpretAnswer(answer);
        return Math.log(expectation(consistentWorlds, function(w) {
            assert.ok(w.domain.length === trueWorld.domain.length)
            return similarity(question, trueWorld, w, trueWorld.domain);
        }))
    }, /*maxSize*/ 1.75e4)

    var answerQualityForQud = cache(function(answer, qudName, trueWorld) {
        // Using hard partition semantics
        var consistentWorlds = interpretAnswer(answer);
        var qud = qudCandidates[qudName];
        return marginalizeCached(consistentWorlds, qudName).score(qud(trueWorld));
    })

    var explicitAnswerer = cache(function(question, trueWorld) {
        return Infer({method: 'enumerate'}, function() {
            var answer = sample(answersDist);
            var answerMeaning = parser(answer, answerStartSymbol);
            factor(answerQualityForQuestion(answer, question, trueWorld)
                * answererRationality);
            return answer;
        })
    })

    var questioner = function(answerer) {
        var informationGain = informationGainFn(answerer);
        return cache(function(qudIndex) {
            return Infer({method: 'enumerate'}, function() {
                var question = sample(questionsDist)
                var expectedInformationGain = expectation(worldPrior, function(world) {
                    return informationGain(question, qudIndex, world);
                });
                factor(expectedInformationGain * questionerRationality);
                return question;
            });
        })
    }

    var explicitQuestioner = questioner(explicitAnswerer);

    var inferGoal = cache(function(question) {
        Infer({ method: 'enumerate' }, function() {
            var qudName = uniformDraw(_.keys(qudCandidates));
            observe(explicitQuestioner(qudName), question);
            return qudName;
        })
    })

    var pragmaticAnswerer = cache(function(question, trueWorld) {
        return Infer({method: 'enumerate'}, function() {
            var answer = sample(answersDist)
            var qudName = sample(inferGoal(question));
            factor(answerQualityForQud(answer, qudName, trueWorld) * answererRationality);
            return answer;
        });
    })

    var pragmaticQuestioner = questioner(pragmaticAnswerer);

    // Standard RSA
    var literalListener = function(utterance) {
        return interpretAnswer(utterance);
    }

    var speaker = function(world) {
        return Infer({ method: 'enumerate' }, function() {
            var utterance = sample(answersDist);
            var L = literalListener(utterance);
            factor(L.score(world) * answererRationality);
            return utterance;
        })
    }

    var pragmaticListener = function(utterance) {
        return Infer({ method: 'enumerate' }, function() {
            var world = sample(worldPrior);
            var S = speaker(world);
            factor(S.score(utterance));
            return world;
        })
    }

    return { 
        Q1: explicitQuestioner, 
        A1: explicitAnswerer, 
        Q2: pragmaticQuestioner, 
        A2: pragmaticAnswerer,
        L0: literalListener,
        S1: speaker,
        L1: pragmaticListener
    };
}