   func value() -> Value
 }
 
-// TODO: Improve this temporary and bad implementation. :)
 public struct AverageEpisodeLength<Observation, Action, Reward, State>: Metric {
   public let batchSize: Int
   public let bufferSize: Int
 
-  private var buffer: [Float]
+  private var buffer: [Int]
   private var index: Int
   private var full: Bool
   private var episodeSteps: Tensor<Int32>

  }

  public func value() -> Float {
    if full {
      return buffer.reduce(0, +) / Float(buffer.count)
    }
    return buffer[0..<index].reduce(0, +) / Float(index)
  }
}

  public init(batchSize: Int, bufferSize: Int) {
    self.batchSize = batchSize
    self.bufferSize = bufferSize
    self.buffer = [Int](repeating: 0, count: bufferSize)
    self.index = 0
    self.full = false
    self.episodeSteps = Tensor<Int32>(repeating: 0, shape: [batchSize])
  }

  public mutating func update(using trajectory: Trajectory<Observation, Action, Reward, State>) {
    let isLast = trajectory.stepKind.isLast()
    let isNotLast = 1 - Tensor<Int32>(isLast)
    episodeSteps += isNotLast
    for length in episodeSteps.gathering(where: isLast).scalars {
      buffer[index] = Int(length)
      index += 1
      full = full || index == buffer.count
      index = index % buffer.count
