     public var vision: Tensor<Float>
     public var scent: Tensor<Float>
     @noDerivative public var moved: Tensor<Bool>
-
-    @inlinable
-    public init(vision: Tensor<Float>, scent: Tensor<Float>, moved: Tensor<Bool>) {
+    @noDerivative public var rewardFunction: Reward?
+
+    @inlinable
+    public init(
+      vision: Tensor<Float>,
+      scent: Tensor<Float>,
+      moved: Tensor<Bool>,
+      rewardFunction: Reward?
+    ) {
       self.vision = vision
       self.scent = scent
       self.moved = moved
+      self.rewardFunction = rewardFunction
     }
   }
 }

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

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