         Observation(
           vision: visionDistribution.mode(),
           scent: scentDistribution.mode(),
-          moved: movedDistribution.mode() .> 0)
+          moved: movedDistribution.mode() .> 0,
+          rewardFunction: nil)
       }
 
       @inlinable

      states[i].agent.nextAction = Int(actions[i].scalarized())
      try states[i].simulator.step()
      let agentState = states[i].simulator.agentStates.values.first!
      let observation = Observation(
        vision: Tensor<Float>(agentState.vision),
        scent: Tensor<Float>(agentState.scent),
        moved: Tensor<Bool>(agentState.position != previousAgentState.position))
      let rewardFunction = configurations[i].rewardSchedule.reward(
        forStep: states[i].simulator.time)
      let reward = Tensor<Float>(rewardFunction(for: AgentTransition(
        previousState: previousAgentState,
        currentState: agentState)))

      true // TODO: Check for the range of values.
    }

    // TODO: How do we sample a reward function?
    public struct ValueDistribution: DifferentiableDistribution, KeyPathIterable {
      @usableFromInline internal var visionDistribution: Uniform<Float> = Uniform<Float>(
        lowerBound: Tensor<Float>(0),
