   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) {
     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)
+    hiddenLSTMCell = LSTMCell<Float>(inputSize: 2 * hiddenSize + 1, hiddenSize: hiddenSize)
+    denseAction = Dense<Float>(inputSize: 3 * hiddenSize + 1, outputSize: 3)
+    denseValue = Dense<Float>(inputSize: 3 * 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])

}

public struct LSTMCell<Scalar: TensorFlowFloatingPoint>: Layer {
  public var inputWeight, updateWeight, forgetWeight, outputWeight: Tensor<Scalar>
  public var inputBias, updateBias, forgetBias, outputBias: Tensor<Scalar>

  @inlinable
  public func stateShape(batchSize: Int) -> TensorShape {
    TensorShape([batchSize, inputWeight.shape[1]])
  }

  @inlinable
  public func zeroState(batchSize: Int) -> LSTMState<Scalar> {
    let stateShape = self.stateShape(batchSize: batchSize)
    return LSTMState(cell: Tensor(zeros: stateShape), hidden: Tensor(zeros: stateShape))
  }

  @inlinable
  public init(
    inputSize: Int,
    hiddenSize: Int,
    seed: TensorFlowSeed = Context.local.randomSeed
  ) {
    let concatenatedInputSize = inputSize + hiddenSize
    let gateWeightShape = TensorShape([concatenatedInputSize, hiddenSize])
    let gateBiasShape = TensorShape([hiddenSize])
    self.inputWeight = Tensor(glorotUniform: gateWeightShape, seed: seed)
    self.inputBias = Tensor(zeros: gateBiasShape)
    self.updateWeight = Tensor(glorotUniform: gateWeightShape, seed: seed)
    self.updateBias = Tensor(zeros: gateBiasShape)
    self.forgetWeight = Tensor(glorotUniform: gateWeightShape, seed: seed)
    self.forgetBias = Tensor(ones: gateBiasShape)
    self.outputWeight = Tensor(glorotUniform: gateWeightShape, seed: seed)
    self.outputBias = Tensor(zeros: gateBiasShape)
  }

  @inlinable
  @differentiable
  public func callAsFunction(_ input: LSTMValue<Scalar>) -> LSTMValue<Scalar> {
    let gateInput = input.value.concatenated(with: input.state.hidden, alongAxis: 1)
    let inputGate = sigmoid(matmul(gateInput, inputWeight) + inputBias)
    let updateGate = tanh(matmul(gateInput, updateWeight) + updateBias)
    let forgetGate = sigmoid(matmul(gateInput, forgetWeight) + forgetBias)
    let outputGate = sigmoid(matmul(gateInput, outputWeight) + outputBias)
    let newCellState = input.state.cell * forgetGate + inputGate * updateGate
    let newHiddenState = tanh(newCellState) * outputGate
    let newState = LSTMState(cell: newCellState, hidden: newHiddenState)
    return LSTMValue(value: newCellState, state: newState)
  }
}

public struct LSTMValue<Scalar: TensorFlowFloatingPoint>: Differentiable, KeyPathIterable {
  public var value: Tensor<Scalar>
  public var state: LSTMState<Scalar>

  @inlinable
  @differentiable
  public init(value: Tensor<Scalar>, state: LSTMState<Scalar>) {
    self.value = value
    self.state = state
  }
}

public struct LSTMState<Scalar: TensorFlowFloatingPoint>: Differentiable, KeyPathIterable {
  public var cell: Tensor<Scalar>
  public var hidden: Tensor<Scalar>

  @inlinable
  @differentiable
  public init(cell: Tensor<Scalar>, hidden: Tensor<Scalar>) {
    self.cell = cell
    self.hidden = hidden
  }
}


  public var denseValue: Dense<Float>

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

  @inlinable
  public init(hiddenSize: Int = 128) {
    visionLayer = VisionLayer(outputSize: hiddenSize)
    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, 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 vision = observation.vision.flattenedBatch(outerDimCount: outerDimCount)
    let moved = 2 * observation.moved.flattenedBatch(
      outerDimCount: outerDimCount
    ).expandingShape(at: -1) - 1
    let hidden = selu(visionLayer(vision)).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)
