   public var denseValue: Dense<Float>
 
   @inlinable
-  public func initialState(batchSize: Int) -> LSTMState<Float> {
+  public func initialState(batchSize: Int) -> LSTMCell<Float>.State {
     hiddenLSTMCell.zeroState(batchSize: batchSize)
   }
 
   @inlinable
-  public init(hiddenSize: Int = 64) {
+  public init(hiddenSize: Int = 128) {
     scentLayer = ScentLayer(outputSize: hiddenSize)
-    hiddenLSTMCell = LSTMCell<Float>(inputSize: hiddenSize, hiddenSize: hiddenSize)
-    denseAction = Dense<Float>(inputSize: 2 * hiddenSize, outputSize: 3)
-    denseValue = Dense<Float>(inputSize: 2 * hiddenSize, outputSize: 1)
+    hiddenLSTMCell = LSTMCell<Float>(inputSize: hiddenSize + 1, hiddenSize: hiddenSize)
+    denseAction = Dense<Float>(inputSize: 2 * hiddenSize + 1, outputSize: 3)
+    denseValue = Dense<Float>(inputSize: 2 * hiddenSize + 1, outputSize: 1)
   }
 
   @inlinable
   @differentiable
   public func callAsFunction(
-    _ input: AgentInput<JellyBeanWorld.Environment.Observation, LSTMState<Float>>
-  ) -> ActorCriticOutput<Categorical<Int32>, LSTMState<Float>> {
+    _ input: AgentInput<JellyBeanWorld.Environment.Observation, LSTMCell<Float>.State>
+  ) -> ActorCriticOutput<Categorical<Int32>, LSTMCell<Float>.State> {
     let observation = input.observation
     let outerDimCount = observation.vision.rank - 3
     let outerDims = [Int](observation.vision.shape.dimensions[0..<outerDimCount])
     let scent = observation.scent.flattenedBatch(outerDimCount: outerDimCount)
-    let hidden = selu(scentLayer(scent))
+    let moved = 2 * observation.moved.flattenedBatch(
+      outerDimCount: outerDimCount
+    ).expandingShape(at: -1) - 1
+    let hidden = selu(scentLayer(scent)).concatenated(with: moved, alongAxis: -1)
     let state = withoutDerivative(at: input.state).flattenedBatch(outerDimCount: outerDimCount)
-    let hiddenLSTMOutput = hiddenLSTMCell(LSTMValue<Float>(value: hidden, state: state))
+    let hiddenLSTMOutput = hiddenLSTMCell(LSTMCell<Float>.Value(value: hidden, state: state))
     let hiddenConcatenated = hidden.concatenated(with: hiddenLSTMOutput.value, alongAxis: -1)
     let actionLogits = denseAction(hiddenConcatenated)
     let flattenedValue = denseValue(hiddenConcatenated)

  public var denseValue: Dense<Float>

  @inlinable
  public func initialState(batchSize: Int) -> LSTMState<Float> {
    hiddenLSTMCell.zeroState(batchSize: batchSize)
  }

  @inlinable
  public init(hiddenSize: Int = 64) {
    visionLayer = VisionLayer(outputSize: hiddenSize)
    scentLayer = ScentLayer(outputSize: hiddenSize)
    hiddenLSTMCell = LSTMCell<Float>(inputSize: 2 * hiddenSize, hiddenSize: hiddenSize)
    denseAction = Dense<Float>(inputSize: 3 * hiddenSize, outputSize: 3)
    denseValue = Dense<Float>(inputSize: 3 * hiddenSize, outputSize: 1)
  }

  @inlinable
  @differentiable
  public func callAsFunction(
    _ input: AgentInput<JellyBeanWorld.Environment.Observation, LSTMState<Float>>
  ) -> ActorCriticOutput<Categorical<Int32>, LSTMState<Float>> {
    let observation = input.observation
    let outerDimCount = observation.vision.rank - 3
    let outerDims = [Int](observation.vision.shape.dimensions[0..<outerDimCount])

    let scent = observation.scent.flattenedBatch(outerDimCount: outerDimCount)
    let visionHidden = selu(visionLayer(vision))
    let scentHidden = selu(scentLayer(scent))
    let moved = 2 * observation.moved.flattenedBatch(
      outerDimCount: outerDimCount
    ).expandingShape(at: -1) - 1
    let hidden = visionHidden
      .concatenated(with: scentHidden, alongAxis: -1)
      .concatenated(with: moved, alongAxis: -1)
    let state = withoutDerivative(at: input.state).flattenedBatch(outerDimCount: outerDimCount)
    let hiddenLSTMOutput = hiddenLSTMCell(LSTMCell<Float>.Value(value: hidden, state: state))
    let hiddenConcatenated = hidden.concatenated(with: hiddenLSTMOutput.value, alongAxis: -1)
    let actionLogits = denseAction(hiddenConcatenated)
    let flattenedValue = denseValue(hiddenConcatenated)
