     }
   }
 
+  private static func getMemory(gameData: GameData) -> ShapedArray<UInt8> {
+    let handle = gameDataMemory(gameData.handle)
+    let cBlocks = memoryViewBlocks(handle)!.pointee
+    let blocks = Array(UnsafeBufferPointer(start: cBlocks.blocks, count: cBlocks.numBlocks))
+    var memoryBytes = [UInt8]()
+    var numBytesPerBlock = 0
+    blocks.sorted(by: { $0.address > $1.address }).forEach {
+      memoryBytes += Array(UnsafeBufferPointer(start: $0.bytes, count: $0.numBytes))
+      numBytesPerBlock = $0.numBytes
+    }
+    return ShapedArray(shape: [blocks.count, numBytesPerBlock], scalars: memoryBytes)
+  }
+
   public func recordMovie(at url: URL) {
     movie = Movie(at: url, recording: true, numPlayers: numPlayers)
     movie!.configure(for: self)

    self.state = state
  }

  public func memory() -> ShapedArray<UInt8> {
    return Environment.memory(gameData: gameData)
  }

  public func screen(forPlayer player: UInt32 = 0) -> ShapedArray<UInt8> {
    return Environment.screen(
      gameData: gameData,
      emulatorHandle: emulatorHandle,
      forPlayer: player)
  }

  private static func memory(gameData: GameData) -> ShapedArray<UInt8> {
    let handle = gameDataMemory(gameData.handle)
    let cBlocks = memoryViewBlocks(handle)!.pointee
    let blocks = Array(UnsafeBufferPointer(start: cBlocks.blocks, count: cBlocks.numBlocks))
    var memoryBytes = [UInt8]()
    var numBytesPerBlock = 0
    blocks.sorted(by: { $0.address > $1.address }).forEach {
      memoryBytes += Array(UnsafeBufferPointer(start: $0.bytes, count: $0.numBytes))
      numBytesPerBlock = $0.numBytes
    }
    return ShapedArray(shape: [blocks.count, numBytesPerBlock], scalars: memoryBytes)
  }

  private static func screen(
    gameData: GameData,
    emulatorHandle: UnsafeMutablePointer<CEmulator>?, 
    forPlayer player: UInt32 = 0

    // Configure the observation space.
    switch config.observationSpaceType {
    case .screen:
      self.screen = Environment.getScreen(
        gameData: self.gameData,
        emulatorHandle: self.emulatorHandle,
        forPlayer: 0)
      self.observationSpace = Box(low: 0, high: 255, shape: screen!.shape)
    case .memory:
      self.memory = Environment.getMemory(gameData: self.gameData)
      self.observationSpace = Box(low: 0, high: 255, shape: memory!.shape)
    }

    if let state = self.state {
