code: def split_df(df):
      '''
      Split a dataframe in two dataframes: one with the history of agents,
      and one with the environment history
      '''
      envmask = (df['agent_id'] == 'env')
      n_env = envmask.sum()
      if n_env == len(df):
          return df, None
      elif n_env == 0:
          return None, df
      agents, env = [x for _, x in df.groupby(envmask)]
      return env, agentf split_df(df):
      '''
      Split a dataframe in two dataframes: one with the history of agents,
      and one with the environment history
      '''
      envmask = (df['agent_id'] == 'env')
      n_env = envmask.sum()
      if n_env == len(df):
          return df, None
      elif n_env == 0:
          return None, df
      agents, env = [x for _, x in df.groupby(envmask)]
      return env, agents
docs: Split a dataframe in two dataframes: one with the history of agents,
      and one with the environment history