From Sampling to Reinforce

Discrete Random Variables

We are going to start from scratch by implementing a mini-language for working with random variables. Our language will allow us to apply some basic operations on these variables and propagate key properties. These are random in the statistical sense, but will not require drawing pseudorandom values in the CS sense. Instead, we keep all possible values and their probabilities.

The main object in the language is a finite-sized, discrete random variable. It is represented by aligned arrays of values and probabilities. We visualize it as a histogram.

eight = Var(np.arange(1, 9), 
            np.array([1, 3, 1, 3, 1, 3, 1, 3])/16)

histogram(eight)

We can transform random variables to create new ones. Mathematically, this is written as

However, to make it clearer what is going on, we write it as an explicit map over values. Similarly, we define a filter over values that renormalizes the distribution.

histogram_row(eight.op(lambda a: a * 2), eight.cond(lambda a: a < 6),
              labels=("Double the value", "Keep values below 6"))

We will be interested in estimating expectations of random variables. The expectation is represented by the red line in the histogram and computed as

def expect(x: Var) -> float:
    return float(x.values @ x.probs)

print(f"E[X] = {expect(eight):.2f}")

E[X] = 4.75

As we build our way to Reinforce, our primary goal will be to build low-variance, unbiased estimators. Unbiased means that they have the same expectation as a target quantity. Variance measures the spread of a random variable around its mean. It's the quantity we aim to minimize.

The visualization for variance will be a weighted sum of squares. The base of each box is a square whose side length is the distance from the mean. Its height is the same as in the histogram.

If it is helpful, you can think of a random variable as expectation + noise. The variance is a way of quantifying the noise.

def variance(x: Var) -> float:
    mean = expect(x)
    return expect(x.op(lambda value: (value - mean)**2))

print(f"Var(X) = {variance(eight):.4f}")
display(scaled_variance_slider(eight))

Var(X) = 5.1875

Var(bX) = 5.19

Joint Random Variables

Next, we will consider the relationship between random variables. We'll be overly pedantic here to show how it works. Our object here will keep track of the probability of every shared state in a table.

joint_example = Joint([0, 2, 5], [0, 1],
                      [[1/12, 1/12], [1/6, 1/4], [1/4, 1/6]])
histogram(joint_example)

Just as with random variables, we can map over the values.

transformed = joint_example.op(lambda a: 2 * a, lambda b: -1 * b + 3)
histogram(transformed)

We can also combine them.

histogram(transformed.add())

Or construct joints from independent random variables.

die = uniform(1, 7)
two_dice = indep(die, die) 
histogram(two_dice)

Or from the same shared random variable.

two_dice_shared = shared(die)
histogram(two_dice_shared)

Notably, these joints have different sums.

histogram_row(two_dice.add(), two_dice_shared.add(),
              labels=("Independent", "Shared"))

The key property of joints will be their covariance.

Covariance measures how each paired outcome in our joint distribution relates to the two means. Unlike variance contributions, covariance contributions can be positive or negative, depending on which side of their respective means the paired values lie.

In the sense that variance represents the squared norm of noise, covariance is the dot product. For the dice above, the two ways of drawing give these rules:

In this example we interpolate between shared and independent dice and show the covariance. Magenta indicates a negative contribution.

assert np.isclose(covar(two_dice), 0)
assert np.isclose(covar(two_dice_shared), variance(die))
display(covariance_interpolation_slider(die))

Cov(X,Y) = -0.000;  Var(X) = 2.917