       try simulator.add(agent: agent)
       return State(simulator: simulator, agent: agent)
     }
-    let observation = Observation.stack(states.map { state in
+    let observation = Observation.stack(zip(configurations, states).map { (configuration, state) in
       let agentState = state.simulator.agentStates.values.first!
+      let rewardFunction = configuration.rewardSchedule.reward(forStep: state.simulator.time)
       return Observation(
         vision: Tensor<Float>(agentState.vision),
         scent: Tensor<Float>(agentState.scent),
-        moved: Tensor<Bool>(repeating: false, shape: [batchSize]))
+        moved: Tensor<Bool>(repeating: false, shape: [batchSize]),
+        rewardFunction: rewardFunction)
     })
     step =  Step(
       kind: StepKind.first(batchSize: batchSize),

        Observation(
          vision: visionDistribution.mode(),
          scent: scentDistribution.mode(),
          moved: movedDistribution.mode() .> 0)
      }

      @inlinable

      try simulator.add(agent: agent)
      return State(simulator: simulator, agent: agent)
    }
    let observation = Observation.stack(zip(configurations, states).map { (configuration, state) in
      let agentState = state.simulator.agentStates.values.first!
      let rewardFunction = configuration.rewardSchedule.reward(forStep: state.simulator.time)
      return Observation(
        vision: Tensor<Float>(agentState.vision),
        scent: Tensor<Float>(agentState.scent),
        moved: Tensor<Bool>(repeating: false, shape: [batchSize]),
        rewardFunction: rewardFunction)
    })
    self.step =  Step(
      kind: StepKind.first(batchSize: batchSize),
