UpdateSpec rather than choicemap for update specifications

[georgematheos]

I think it can sometimes make sense for custom generative functions to know how to perform special types of updates which are not simply of the form of of changing the value of one of the choices in the choicemap.

For instance, in some code I'm writing, we have a type of "split" update which splits a generated object into two generated objects with the same properties as the original. One hacky thing we can do now is have an update call along the lines of:

update(tr, args, argdiffs, choicemap((:inner_gen_fn_addr => :split, object_id))

But I think that really the call signature for the update function should be something along the lines of

Gen.update(tr::Trace, args, argdiffs, update_spec::UpdateSpecification)

where we have a built-in ChoicemapUpdateSpecification <: UpdateSpecification that choicemaps can get cast to by default so the old form of the interface works.

Then specifying a split update as in the example above could look like

update(tr, args, argdiffs, updatespec((:inner_gen_fn_addr => SplitUpdate(object_id)))

or, if we wanted, we could even have some automatic way to convert Pairs into update specs, so we could have the interface simply look like

update(tr, args, argdiffs, :inner_gen_fn_addr => SplitUpdate(object_id))

(Really, I think that the update method interface should take a fifth optional kwargs argument, as I describe in https://github.com/probcomp/Gen/issues/182, but this is a separate issue.)

[georgematheos]

Another consideration is that if we use UpdateSpecifications rather than choicemaps to specify updates, perhaps it makes sense for the discard return value of Gen.update to be a "reversing" UpdateSpecification rather than a choicemap.

In other words, the current update interface is

new_tr::Trace, weight::Float64, retdiff::Diff, discard::ChoiceMap = Gen.update(tr, args, argdiffs, constraints::ChoiceMap)

but I perhaps we should change it to

new_tr::Trace, weight::Float64, retdiff::Diff, reverse_spec::UpdateSpecification = Gen.update(tr, args, argdiffs, spec::UpdateSpecification)

This is because, as far as I am aware, the only place we use the discard is to tell us how to make reverse moves in MCMC. Therefore if we make it possible to update traces using more general UpdateSpecification objects, the discard should be moved to an object of this type as well.

[georgematheos]

I want to give a sense of what it might look like to run MCMC inference with this new interface. In this example, say I am running "split/merge" MH updates on a trace. Say I have a custom generative function generate_objects which understands how to "split" and "merge" objects using SplitUpdate and MergeUpdate UpdateSpecifications. If I pass in a SplitUpdate(obj_id) which results in obj_id splitting to objects with obj_1 and obj_2 as IDs, then the discard from this update call will be a MergeUpdate(obj_1, obj_2), the move which reverses the split by merging the split objects into one object.

If I want to run MH, I can create a proposal distribution which will decide which objects to split and merge:

@gen function splitmerge_proposal(tr)
  if @trace(bernoulli(0.5), :do_split)
    @trace(sample_one(tr[:object_set], :target)
  else
    @trace(sample_one(tr[:object_set]), :target1)
    @trace(sample_one(tr[:object_set]), :target2)
  end
end

Then to use mh, we can use the involution variant to translate from the outputs of propose(splitmerge_proposal) to SplitUpdate and MergeUpdates. This can look something like

function involution(trace, fwd_choices::ChoiceMap, fwd_ret, _)
  args = get_args(trace)
  argdiffs = map(_ -> NoChange(), args)
  if fwd_choices[:do_split]
    new_tr, weight, retdiff, reverse_update_spec = update(trace, args, argdiffs, SplitUpdate(fwd_choices[:target]))
  bwd_choices = choicemap((:do_split, false), (:target1, reverse_update_spec.target1), (:target2, reverse_update_spec.target2))
  else
    new_tr, weight, retdiff, reverse_update_spec = update(trace, args, argdiffs, MergeUpdate(fwd_choices[:target1], fwd_choices[:target2))
    bwd_choices = choicemap((:do_split, true), (:target, reverse_update_spec.target))
  end
  return new_tr, bwd_choices, weight

Then to run MH, we can call

tr, acc = Gen.mh(tr, splitmerge_proposal, (), involution)

The involution function handles translating from the decisions made by the splitmerge_proposal to the UpdateSpecifications that should be passed into the update function, and handles translating from the reverse UpdateSpecifications output by update to what values in the proposal distribution could have produced them.

(If we wanted, we could actually implement a whole new version of mh where instead of passing in an involution, we pass in 2 functions: 1 to map from proposal outputs to UpdateSpecs, and one to map from UpdateSpecs to proposal outputs which should imply those UpdateSpecs.)