def generate_trajectory(sequence_fn, space, n_steps):
  """Generates random actions and observations that follow sequence_fn."""
  act = [space.sample() for _ in range(n_steps)]
  obs = [space.sample()]

  for (o, a) in zip(
      obs,
      act[:-1],  # Don't generate the last observation.
  ):
    context = list(np.array([o, a]).flatten())
    symbols = []
    for _ in range(np.array(o).size):
      symbol = sequence_fn(context + symbols)
      symbols.append(symbol)
    obs.append(np.reshape(symbols, space.shape))

  obs = np.array([obs])
  act = np.array([act])
  return (obs, act)
Generates random actions and observations that follow sequence_fn.
trax/supervised/callbacks_test.py