Treatment assignment as extra demographic variable

[chathasphere]

Necessary for AIPW outcome regression model

Currently the tensor w is just age and sex. We want to (optionally) modify this to have age, sex, and treatment assignment

[chathasphere]

Not entirely certain whether padding out instances of treatment in the records actually "works." Probably better practice to throw "unknown token" at it or something