Skip to article frontmatterSkip to article content
Site not loading correctly?

This may be due to an incorrect BASE_URL configuration. See the MyST Documentation for reference.

Open In Colab   Open in Kaggle

Macrocircuits

Macrocircuits: Leveraging neural architectural priors and modularity in embodied agents

By Neuromatch Academy

Content creators: Divyansha Lachi, Kseniia Shilova

Content reviewers: Eva Dyer, Hannah Choi

Production editors: Konstantine Tsafatinos, Ella Batty, Spiros Chavlis, Samuele Bolotta, Hlib Solodzhuk


Background

This project explores how we can build a biologically inspired artificial neural network (ANN) architecture, derived from the C. Elegans motor circuit, for the control of a simulated Swimmer agent. Traditional motor control ANNs often rely on generic, fully connected multilayer perceptrons (MLPs), which demand extensive training data, offer limited transferability, and possess complex internal dynamics that challenge interpretability. The project aims to understand how the biologically motivated ANN, which is shaped by evolution to be highly structured and sparse, could help to solve these problems and provide advantages in the domain of motor control. We will train MLPs using algorithms such as PPO, DDPG, and ES, and compare their performance in terms of rewards and sample efficiency with our bio-inspired ANN. The project also includes visualizing the C. Elegans connectome and building the network using this circuitry. We will conduct various ablation analyses by removing sign and weight-sharing constraints, and altering environmental parameters like the swimmer’s length or viscosity. These investigations aim to understand how architecture and modularity impact performance and learning across different environments. Finally, the project aims at building an agent that is robust to environmental variations, navigating towards specific targets, and enhancing our understanding of bio-inspired motor control.

Relevant references:

This notebook uses code from the following GitHub repository: ncap by Nikhil X. Bhattasali and Anthony M. Zador and Tatiana A. Engel.

Infrastructure note: This notebook contains GPU install guide as well as CPU ones for different OS.

Install and import feedback gadget

Source
# @title Install and import feedback gadget

!pip install vibecheck datatops --quiet

from vibecheck import DatatopsContentReviewContainer
def content_review(notebook_section: str):
    return DatatopsContentReviewContainer(
        "",  # No text prompt
        notebook_section,
        {
            "url": "https://pmyvdlilci.execute-api.us-east-1.amazonaws.com/klab",
            "name": "neuromatch_neuroai",
            "user_key": "wb2cxze8",
        },
    ).render()

feedback_prefix = "Project_Macrocircuits"

Project Background

Youtube
Bilibili

Project slides

If you want to download the slides: https://osf.io/download/cz93b/
Loading...

Project Template

Source
#@title Project Template
from IPython.display import Image, display
import os
from pathlib import Path

url = "https://github.com/neuromatch/NeuroAI_Course/blob/main/projects/project-notebooks/static/NCAPProjectTemplate.png?raw=true"

display(Image(url=url))
Loading...

Tutorial links

This particular project connects a couple of distinct ideas explored throughout the course. Firstly, the innate ability to learn a certain set of actions quickly is the main topic of Tutorial 4 for W2D4 on biological meta-learning. The focus comes with the observation that the brain is not of a generic architecture but is a highly structured and optimized hierarchy of modules, the importance of which is highlighted in Tutorial 3 for W2D1, forming inductive bias for efficient motor control. The default model for the agent used here is already known Actor-Critic; you had the opportunity to observe in already mentioned tutorials as well as in Tutorial 3 for W1D2.


Scetion 0: Initial setup

Disclaimer: As an alternative to Google Colab, Kaggle, and local installation, we have prepared a Dockerfile with the instructions on the virtual environment setup. It will allow you to work locally with no interventions into already installed packages (you will just need to install Docker and run two commands). Grab a Dockerfile, put it in the same folder as this notebook, and follow the instructions in README file.

IF USING COLAB (recommended):

Uncomment the cell below and run it.

Installing Dependencies (Colab GPU case, uncomment if you want to use this one)

Source
#@title Installing Dependencies (Colab GPU case, uncomment if you want to use this one)

import distutils.util
import os
import subprocess
if subprocess.run('nvidia-smi').returncode:
  raise RuntimeError(
      'Cannot communicate with GPU. '
      'Make sure you are using a GPU Colab runtime. '
      'Go to the Runtime menu and select Choose runtime type.')

# Add an ICD config so that glvnd can pick up the Nvidia EGL driver.
# This is usually installed as part of an Nvidia driver package, but the Colab
# kernel doesn't install its driver via APT, and as a result the ICD is missing.
# (https://github.com/NVIDIA/libglvnd/blob/master/src/EGL/icd_enumeration.md)
NVIDIA_ICD_CONFIG_PATH = '/usr/share/glvnd/egl_vendor.d/10_nvidia.json'
if not os.path.exists(NVIDIA_ICD_CONFIG_PATH):
  with open(NVIDIA_ICD_CONFIG_PATH, 'w') as f:
    f.write("""{
    "file_format_version" : "1.0.0",
    "ICD" : {
        "library_path" : "libEGL_nvidia.so.0"
    }
}
""")

print('Installing dm_control...')
!pip install -q dm_control>=1.0.16

# Configure dm_control to use the EGL rendering backend (requires GPU)
%env MUJOCO_GL=egl

!echo Installed dm_control $(pip show dm_control | grep -Po "(?<=Version: ).+")
!pip install -q dm-acme
!mkdir output_videos
---------------------------------------------------------------------------
FileNotFoundError                         Traceback (most recent call last)
Cell In[4], line 6
      4 import os
      5 import subprocess
----> 6 if subprocess.run('nvidia-smi').returncode:
      7   raise RuntimeError(
      8       'Cannot communicate with GPU. '
      9       'Make sure you are using a GPU Colab runtime. '
     10       'Go to the Runtime menu and select Choose runtime type.')
     12 # Add an ICD config so that glvnd can pick up the Nvidia EGL driver.
     13 # This is usually installed as part of an Nvidia driver package, but the Colab
     14 # kernel doesn't install its driver via APT, and as a result the ICD is missing.
     15 # (https://github.com/NVIDIA/libglvnd/blob/master/src/EGL/icd_enumeration.md)

File /opt/hostedtoolcache/Python/3.10.20/x64/lib/python3.10/subprocess.py:503, in run(input, capture_output, timeout, check, *popenargs, **kwargs)
    500     kwargs['stdout'] = PIPE
    501     kwargs['stderr'] = PIPE
--> 503 with Popen(*popenargs, **kwargs) as process:
    504     try:
    505         stdout, stderr = process.communicate(input, timeout=timeout)

File /opt/hostedtoolcache/Python/3.10.20/x64/lib/python3.10/subprocess.py:971, in Popen.__init__(self, args, bufsize, executable, stdin, stdout, stderr, preexec_fn, close_fds, shell, cwd, env, universal_newlines, startupinfo, creationflags, restore_signals, start_new_session, pass_fds, user, group, extra_groups, encoding, errors, text, umask, pipesize)
    967         if self.text_mode:
    968             self.stderr = io.TextIOWrapper(self.stderr,
    969                     encoding=encoding, errors=errors)
--> 971     self._execute_child(args, executable, preexec_fn, close_fds,
    972                         pass_fds, cwd, env,
    973                         startupinfo, creationflags, shell,
    974                         p2cread, p2cwrite,
    975                         c2pread, c2pwrite,
    976                         errread, errwrite,
    977                         restore_signals,
    978                         gid, gids, uid, umask,
    979                         start_new_session)
    980 except:
    981     # Cleanup if the child failed starting.
    982     for f in filter(None, (self.stdin, self.stdout, self.stderr)):

File /opt/hostedtoolcache/Python/3.10.20/x64/lib/python3.10/subprocess.py:1863, in Popen._execute_child(self, args, executable, preexec_fn, close_fds, pass_fds, cwd, env, startupinfo, creationflags, shell, p2cread, p2cwrite, c2pread, c2pwrite, errread, errwrite, restore_signals, gid, gids, uid, umask, start_new_session)
   1861     if errno_num != 0:
   1862         err_msg = os.strerror(errno_num)
-> 1863     raise child_exception_type(errno_num, err_msg, err_filename)
   1864 raise child_exception_type(err_msg)

FileNotFoundError: [Errno 2] No such file or directory: 'nvidia-smi'

IF USING KAGGLE (recommended):

Uncomment the cell below and run it.

Installing Dependencies (Kaggle GPU case, uncomment if you want to use this one)

Source
#@title Installing Dependencies (Kaggle GPU case, uncomment if you want to use this one)

# import subprocess

# subprocess.run(["sudo", "apt-get", "install", "-y", "libgl1-mesa-glx", "libosmesa6"])
# subprocess.run(["pip", "install", "-q", "imageio[ffmpeg]"])

# print('Installing dm_control...')
# !pip install -q dm_control>=1.0.16

# %env MUJOCO_GL=osmesa

# !echo Installed dm_control $(pip show dm_control | grep -Po "(?<=Version: ).+")
# !pip install -q dm-acme[envs]
# !mkdir output_videos

IF RUNNING LOCALLY

Uncomment the relevant lines of code depending on your OS.

Installing Dependencies (CPU case, comment if you want to use GPU one)

Source
#@title Installing Dependencies (CPU case, comment if you want to use GPU one)

import subprocess
import os

############### For Linux #####################
# subprocess.run(["sudo", "apt-get", "install", "-y", "libglew-dev"])
# subprocess.run(["sudo", "apt-get", "install", "-y", "libglfw3"])
# subprocess.run(["sudo", "apt", "install", "ffmpeg"])
###############################################

############### For MacOS #####################
# subprocess.run(["brew", "install", "glew"])
# subprocess.run(["brew", "install", "glfw"])
###############################################

subprocess.run(["pip", "install", "-q", "ffmpeg"])
subprocess.run(["pip", "install", "-q", "dm-acme[envs]"])
subprocess.run(["pip", "install", "-q", "dm_control>=1.0.16"])

!mkdir output_videos

Imports and Utility Functions

Importing Libraries

Source
#@title Importing Libraries
import numpy as np
import collections
import argparse
import os
import yaml
import typing as T
import imageio
import matplotlib
import matplotlib.pyplot as plt
import matplotlib.animation as animation
import pandas as pd
import seaborn as sns
from IPython.display import HTML

import dm_control as dm
import dm_control.suite.swimmer as swimmer
from dm_control.rl import control
from dm_control.utils import rewards
from dm_control import suite
# from dm_control.suite.wrappers import pixels

# from acme import wrappers

from torch import nn

Utility code for displaying videos

Source
#@title Utility code for displaying videos
def write_video(
  filepath: os.PathLike,
  frames: T.Iterable[np.ndarray],
  fps: int = 60,
  macro_block_size: T.Optional[int] = None,
  quality: int = 10,
  verbose: bool = False,
  **kwargs,
):
  """
  Saves a sequence of frames as a video file.

  Parameters:
  - filepath (os.PathLike): Path to save the video file.
  - frames (Iterable[np.ndarray]): An iterable of frames, where each frame is a numpy array.
  - fps (int, optional): Frames per second, defaults to 60.
  - macro_block_size (Optional[int], optional): Macro block size for video encoding, can affect compression efficiency.
  - quality (int, optional): Quality of the output video, higher values indicate better quality.
  - verbose (bool, optional): If True, prints the file path where the video is saved.
  - **kwargs: Additional keyword arguments passed to the imageio.get_writer function.

  Returns:
  None. The video is written to the specified filepath.
  """

  with imageio.get_writer(filepath,
                        fps=fps,
                        macro_block_size=macro_block_size,
                        quality=quality,
                        **kwargs) as video:
    if verbose: print('Saving video to:', filepath)
    for frame in frames:
      video.append_data(frame)


def display_video(
  frames: T.Iterable[np.ndarray],
  filename='output_videos/temp.mp4',
  fps=60,
  **kwargs,
):
  """
  Displays a video within a Jupyter Notebook from an iterable of frames.

  Parameters:
  - frames (Iterable[np.ndarray]): An iterable of frames, where each frame is a numpy array.
  - filename (str, optional): Temporary filename to save the video before display, defaults to 'output_videos/temp.mp4'.
  - fps (int, optional): Frames per second for the video display, defaults to 60.
  - **kwargs: Additional keyword arguments passed to the write_video function.

  Returns:
  HTML object: An HTML video element that can be displayed in a Jupyter Notebook.
  """

  # Write video to a temporary file.
  filepath = os.path.abspath(filename)
  write_video(filepath, frames, fps=fps, verbose=False, **kwargs)

  height, width, _ = frames[0].shape
  dpi = 70
  orig_backend = matplotlib.get_backend()
  matplotlib.use('Agg')  # Switch to headless 'Agg' to inhibit figure rendering.
  fig, ax = plt.subplots(1, 1, figsize=(width / dpi, height / dpi), dpi=dpi)
  matplotlib.use(orig_backend)  # Switch back to the original backend.
  ax.set_axis_off()
  ax.set_aspect('equal')
  ax.set_position([0, 0, 1, 1])
  im = ax.imshow(frames[0])
  def update(frame):
    im.set_data(frame)
    return [im]
  interval = 1000/fps
  anim = animation.FuncAnimation(fig=fig, func=update, frames=frames,
                                  interval=interval, blit=True, repeat=False)
  return HTML(anim.to_html5_video())

In this notebook we will explore the major components essential for this project.

  • Understanding the DeepMind Control Suite Swimmer Agent: We will begin by exploring the swimmer agent provided by the DeepMind Control Suite. This section includes a detailed exploration of the agent’s API, task customization capabilities, and how to adapt the environment to fit our experimental needs.

  • Training Models Using Various Reinforcement Learning Algorithms: Next, we move on to learn how can we train models for the agents we created. We will be using Tonic_RL library to train our model. We will first train a standard MLP model using the Proximal Policy Optimization (PPO) algorithm.

  • Training the NCAP model: Finally we will define the NCAP model from Neural Circuit Architectural Priors for Embodied Control paper. We will train it using PPO and compare it against the MLP model we trained before.


Section 1: Exploring the DeepMind Swimmer

1.1 Create a basic swim task for the swimmer environment

First, we’ll initialize a basic swimmer agent consisting of 6 links. Each agent requires a defined task and its corresponding reward function. In this instance, we’ve designed a swim forward task that involves the agent swimming forward in any direction.

The environment is flexible, allowing for modifications to introduce additional tasks such as “swim only in the x-direction” or “move towards a ball.”

_SWIM_SPEED = 0.1

@swimmer.SUITE.add()
def swim(
  n_links=6,
  desired_speed=_SWIM_SPEED,
  time_limit=swimmer._DEFAULT_TIME_LIMIT,
  random=None,
  environment_kwargs={},
):
  """Returns the Swim task for a n-link swimmer."""
  model_string, assets = swimmer.get_model_and_assets(n_links)
  physics = swimmer.Physics.from_xml_string(model_string, assets=assets)
  task = Swim(desired_speed=desired_speed, random=random)
  return control.Environment(
    physics,
    task,
    time_limit=time_limit,
    control_timestep=swimmer._CONTROL_TIMESTEP,
    **environment_kwargs,
  )


class Swim(swimmer.Swimmer):
  """Task to swim forwards at the desired speed."""
  def __init__(self, desired_speed=_SWIM_SPEED, **kwargs):
    super().__init__(**kwargs)
    self._desired_speed = desired_speed

  def initialize_episode(self, physics):
    super().initialize_episode(physics)
    # Hide target by setting alpha to 0.
    physics.named.model.mat_rgba['target', 'a'] = 0
    physics.named.model.mat_rgba['target_default', 'a'] = 0
    physics.named.model.mat_rgba['target_highlight', 'a'] = 0

  def get_observation(self, physics):
    """Returns an observation of joint angles and body velocities."""
    obs = collections.OrderedDict()
    obs['joints'] = physics.joints()
    obs['body_velocities'] = physics.body_velocities()
    return obs

  def get_reward(self, physics):
    """Returns a smooth reward that is 0 when stopped or moving backwards, and rises linearly to 1
    when moving forwards at the desired speed."""
    forward_velocity = -physics.named.data.sensordata['head_vel'][1]
    return rewards.tolerance(
      forward_velocity,
      bounds=(self._desired_speed, float('inf')),
      margin=self._desired_speed,
      value_at_margin=0.,
      sigmoid='linear',
    )

1.2 Vizualizing an agent that takes random actions in the environment

Let’s visualize the environment by executing a sequence of random actions on a swimmer agent. This involves applying random actions over a series of steps and compiling the rendered frames into a video to visualize the agent’s behavior.

""" Renders the current environment state to an image """
def render(env):
    return env.physics.render(camera_id=0, width=640, height=480)

""" Tests a DeepMind control suite environment by executing a series of random actions """
def test_dm_control(env):
    spec = env.action_spec()
    timestep = env.reset()
    frames = [render(env)]

    for _ in range(60):
        action = np.random.uniform(
            low=spec.minimum,
            high=spec.maximum,
            size=spec.shape
        )
        timestep = env.step(action)
        frames.append(render(env))

    return display_video(frames)

env = suite.load('swimmer', 'swim', task_kwargs={'random': 1})
test_dm_control(env)
/tmp/ipykernel_8472/4073379236.py:65: MatplotlibDeprecationWarning: Auto-close()ing of figures upon backend switching is deprecated since 3.8 and will be removed two minor releases later.  To suppress this warning, explicitly call plt.close('all') first.
  matplotlib.use(orig_backend)  # Switch back to the original backend.
Loading...

1.3 Swimmer Agent API

The observation space consists of 25 total dimensions, combining joint positions and body velocities, while the action space involves 5 dimensions representing normalized joint forces.

Observation Space: 4k - 1 total (k = 6 \rightarrow 23)

  • k - 1: joint positions qi[π,π]q_i \in [-\pi, \pi] (joints)

  • 3k: link linear velocities vxi,vyiRvx_i, vy_i \in \mathbb{R} and rotational velocity wziRwz_i \in \mathbb{R} (body_velocities)

env.observation_spec()
OrderedDict([('joints', Array(shape=(5,), dtype=dtype('float64'), name='joints')), ('body_velocities', Array(shape=(18,), dtype=dtype('float64'), name='body_velocities'))])

Action Space: k - 1 total (k = 6 \rightarrow 5)

  • k - 1: joint normalized force q¨i[1,1]\ddot{q}_i \in [-1, 1]

env.action_spec()
BoundedArray(shape=(5,), dtype=dtype('float64'), name=None, minimum=[-1. -1. -1. -1. -1.], maximum=[1. 1. 1. 1. 1.])

1.4 Example of simple modification to the agent

Let’s make a new swimmer agent with 12 links instead of 6, introducing complexity. Additionally, we have the flexibility to adjust various other parameters.

@swimmer.SUITE.add()
def swim_12_links(
  n_links=12,
  desired_speed=_SWIM_SPEED,
  time_limit=swimmer._DEFAULT_TIME_LIMIT,
  random=None,
  environment_kwargs={},
):
  """Returns the Swim task for a n-link swimmer."""
  model_string, assets = swimmer.get_model_and_assets(n_links)
  physics = swimmer.Physics.from_xml_string(model_string, assets=assets)
  task = Swim(desired_speed=desired_speed, random=random)
  return control.Environment(
    physics,
    task,
    time_limit=time_limit,
    control_timestep=swimmer._CONTROL_TIMESTEP,
    **environment_kwargs,
  )

env = suite.load('swimmer', 'swim_12_links', task_kwargs={'random': 1})
test_dm_control(env)
/tmp/ipykernel_8472/4073379236.py:65: MatplotlibDeprecationWarning: Auto-close()ing of figures upon backend switching is deprecated since 3.8 and will be removed two minor releases later.  To suppress this warning, explicitly call plt.close('all') first.
  matplotlib.use(orig_backend)  # Switch back to the original backend.
Loading...

We can visualize this longer agent using our previously defined test_dm_control function.

Using the API provided by Deepmind we can create any kind of changes to the agent and the environment.

Try to make the following changes to make yourself more familiar with the swimmer.

  • Adding a target (like a ball) to this environment at some x distance away from the agent.

  • Increasing the viscosity of the environment.

Have a look at the following links to see what kind of assets you will need to modify to make these changes.

# add your code

Section 2: Training models on the swim task

To train the agents we defined in the previous section, we will utilize standard reinforcement learning (RL) algorithms. For the purposes of this tutorial, we will employ the tonic_rl library, which provides a robust framework for training RL agents. Throughout most of this project, you will primarily be modifying the environment or the model architecture. Therefore, I suggest treating these algorithms as a “black box” for now. Simply put, you input an untrained model, and the algorithm processes and returns a well-trained model. This approach allows us to focus on the impact of different architectures and environmental settings without delving deeply into the algorithmic complexities at this stage.

Download and install tonic library for training agents

Source
#@title Download and install tonic library for training agents

import contextlib
import io

with contextlib.redirect_stdout(io.StringIO()): #to suppress output
    !git clone https://github.com/neuromatch/tonic
    %cd tonic

Section 2.1 Defining the train function

First we defined a general training function to train any agent on any given environment with a variety of available algorithms. Given below are some of the parameter definitions of the function. You’ll likely want to adjust these parameters to customize the training process for an agent in a specific environment using your chosen algorithm from the tonic library:

  • Header: Python code required to run before training begins, primarily for importing essential libraries or modules.

  • Agent: The agent that will undergo training; refer to section 3.2 and 4.2 for definitions of MLP and NCAP respectively.

  • Environment: The training environment for the agent. Ensure it is registered with the DeepMind Control Suite as detailed in section 2.

  • Name: The experiment’s name, which will be utilized for log and model saving purposes.

  • Trainer: The trainer instance selected for use. It allows the configuration of the training steps, model saving frequency, and other training-related parameters.

import tonic
import tonic.torch

def train(
  header,
  agent,
  environment,
  name = 'test',
  trainer = 'tonic.Trainer()',
  before_training = None,
  after_training = None,
  parallel = 1,
  sequential = 1,
  seed = 0
):
  """
  Some additional parameters:

  - before_training: Python code to execute immediately before the training loop commences, suitable for setup actions needed after initialization but prior to training.
  - after_training: Python code to run once the training loop concludes, ideal for teardown or analytical purposes.
  - parallel: The count of environments to execute in parallel. Limited to 1 in a Colab notebook, but if additional resources are available, this number can be increased to expedite training.
  - sequential: The number of sequential steps the environment runs before sending observations back to the agent. This setting is useful for temporal batching. It can be disregarded for this tutorial's purposes.
  - seed: The experiment's random seed, guaranteeing the reproducibility of the training process.

  """
  # Capture the arguments to save them, e.g. to play with the trained agent.
  args = dict(locals())

  # Run the header first, e.g. to load an ML framework.
  if header:
    exec(header)

  # Build the train and test environments.
  _environment = environment
  environment = tonic.environments.distribute(lambda: eval(_environment), parallel, sequential)
  test_environment = tonic.environments.distribute(lambda: eval(_environment))


  # Build the agent.
  agent = eval(agent)
  agent.initialize(
    observation_space=test_environment.observation_space,
    action_space=test_environment.action_space, seed=seed)

  # Choose a name for the experiment.
  if hasattr(test_environment, 'name'):
    environment_name = test_environment.name
  else:
    environment_name = test_environment.__class__.__name__
  if not name:
    if hasattr(agent, 'name'):
      name = agent.name
    else:
      name = agent.__class__.__name__
    if parallel != 1 or sequential != 1:
      name += f'-{parallel}x{sequential}'

  # Initialize the logger to save data to the path environment/name/seed.
  path = os.path.join('data', 'local', 'experiments', 'tonic', environment_name, name)
  tonic.logger.initialize(path, script_path=None, config=args)

  # Build the trainer.
  trainer = eval(trainer)
  trainer.initialize(
    agent=agent,
    environment=environment,
    test_environment=test_environment,
  )
  # Run some code before training.
  if before_training:
    exec(before_training)

  # Train.
  trainer.run()

  # Run some code after training.
  if after_training:
    exec(after_training)
Gym has been unmaintained since 2022 and does not support NumPy 2.0 amongst other critical functionality.
Please upgrade to Gymnasium, the maintained drop-in replacement of Gym, or contact the authors of your software and request that they upgrade.
Users of this version of Gym should be able to simply replace 'import gym' with 'import gymnasium as gym' in the vast majority of cases.
See the migration guide at https://gymnasium.farama.org/introduction/migration_guide/ for additional information.

Section 2.2 Training MLP model on swim task

Now we are going to define a function for creating an actor-critic model suitable for Proximal Policy Optimization (PPO) using a Multi-Layer Perceptron (MLP) architecture.

from tonic.torch import models, normalizers
import torch

def ppo_mlp_model(
  actor_sizes=(64, 64),
  actor_activation=torch.nn.Tanh,
  critic_sizes=(64, 64),
  critic_activation=torch.nn.Tanh,
):

  """
  Constructs an ActorCritic model with specified architectures for the actor and critic networks.

  Parameters:
  - actor_sizes (tuple): Sizes of the layers in the actor MLP.
  - actor_activation (torch activation): Activation function used in the actor MLP.
  - critic_sizes (tuple): Sizes of the layers in the critic MLP.
  - critic_activation (torch activation): Activation function used in the critic MLP.

  Returns:
  - models.ActorCritic: An ActorCritic model comprising an actor and a critic with MLP torsos,
    equipped with a Gaussian policy head for the actor and a value head for the critic,
    along with observation normalization.
  """

  return models.ActorCritic(
    actor=models.Actor(
      encoder=models.ObservationEncoder(),
      torso=models.MLP(actor_sizes, actor_activation),
      head=models.DetachedScaleGaussianPolicyHead(),
    ),
    critic=models.Critic(
      encoder=models.ObservationEncoder(),
      torso=models.MLP(critic_sizes, critic_activation),
      head=models.ValueHead(),
    ),
    observation_normalizer=normalizers.MeanStd(),
  )

Next we call the train function which initiates the training process for the provided agent using the Tonic library. It specifies the components necessary for training, including the model, environment, and training parameters:

Agent: A Proximal Policy Optimization (PPO) agent with a custom Multi-Layer Perceptron (MLP) model architecture, configured with 256 units in each of two layers for both the actor and the critic.

Environment: The training environment is set to “swimmer-swim” from the Control Suite, a benchmark suite for continuous control tasks.

Name: The experiment is named ‘mlp_256’, which is useful for identifying logs and saved models associated with this training run.

Trainer: Specifies the training configuration, including the total number of steps (5e5) and the frequency of saving the model (1e5 steps).

Note: The model will checkpoint every ‘save_steps’ amount of training steps*

The model can take some time to train so feel free to skip the training for now. We have provided the pretrained model for you to play with. Move on to the next section to vizualize a agent with the pretrained model.

Uncomment the cell below if you want to perform the training.

train('import tonic.torch',
      'tonic.torch.agents.PPO(model=ppo_mlp_model(actor_sizes=(256, 256), critic_sizes=(256,256)))',
      'tonic.environments.ControlSuite("swimmer-swim")',
      name = 'mlp_256',
      trainer = 'tonic.Trainer(steps=int(5e5),save_steps=int(1e5))')
/opt/hostedtoolcache/Python/3.10.20/x64/lib/python3.10/site-packages/gym/spaces/box.py:127: UserWarning: WARN: Box bound precision lowered by casting to float32
  logger.warn(f"Box bound precision lowered by casting to {self.dtype}")
Config file saved to data/local/experiments/tonic/swimmer-swim/mlp_256/config.yaml
          Time left:  epoch 0:00:10  total 0:04:21          
          Time left:  epoch 0:00:09  total 0:04:16          
          Time left:  epoch 0:00:09  total 0:04:13          
          Time left:  epoch 0:00:09  total 0:04:13          
          Time left:  epoch 0:00:08  total 0:04:12          
          Time left:  epoch 0:00:08  total 0:04:12          
          Time left:  epoch 0:00:08  total 0:04:12          
          Time left:  epoch 0:00:16  total 0:08:39          
          Time left:  epoch 0:00:14  total 0:08:05          
          Time left:  epoch 0:00:13  total 0:07:38          
          Time left:  epoch 0:00:11  total 0:07:17          
          Time left:  epoch 0:00:10  total 0:06:59          
          Time left:  epoch 0:00:09  total 0:06:45          
          Time left:  epoch 0:00:10  total 0:07:44          
          Time left:  epoch 0:00:09  total 0:07:28          
          Time left:  epoch 0:00:08  total 0:07:14          
          Time left:  epoch 0:00:08  total 0:07:02          
          Time left:  epoch 0:00:07  total 0:06:52          
          Time left:  epoch 0:00:06  total 0:06:43          
          Time left:  epoch 0:00:06  total 0:07:24          
          Time left:  epoch 0:00:05  total 0:07:14          
          Time left:  epoch 0:00:05  total 0:07:04          
          Time left:  epoch 0:00:04  total 0:06:56          
          Time left:  epoch 0:00:03  total 0:06:48          
          Time left:  epoch 0:00:03  total 0:06:40          
          Time left:  epoch 0:00:03  total 0:06:37          
          Time left:  epoch 0:00:02  total 0:07:05          
          Time left:  epoch 0:00:02  total 0:06:58          
          Time left:  epoch 0:00:01  total 0:06:51          
          Time left:  epoch 0:00:00  total 0:06:45          
          Time left:  epoch 0:00:00  total 0:06:38          
          Time left:  epoch 0:00:00  total 0:06:36          

actor                                                       
  clip fraction                                        0.149
  entropy                                               1.05
  iterations                                            23.8
  kl                                                 0.00834
  loss                                               -0.0172
  std                                                  0.693
  stop                                                0.0316
critic                                                      
  iterations                                              80
  loss                                                  2.67
  v                                                     5.76
test                                                        
  action                                                    
    max                                                 2.95
    mean                                             -0.0265
    min                                                -2.66
    size                                               5,000
    std                                                0.716
  episode length                                            
    max                                                1,000
    mean                                               1e+03
    min                                                1,000
    size                                                   5
    std                                                    0
  episode score                                             
    max                                                  250
    mean                                                 226
    min                                                  210
    size                                                   5
    std                                                 16.5
train                                                       
  action                                                    
    max                                                 2.87
    mean                                             -0.0349
    min                                                -3.14
    size                                              20,000
    std                                                0.701
  episode length                                            
    max                                                1,000
    mean                                               1e+03
    min                                                1,000
    size                                                  20
    std                                                    0
  episode score                                             
    max                                                  248
    mean                                                 145
    min                                                 75.4
    size                                                  20
    std                                                   49
  episodes                                                20
  epoch seconds                                         18.6
  epoch steps                                          2e+04
  epochs                                                   1
  seconds                                               18.6
  steps                                                2e+04
  steps per second                                  1.08e+03
  worker steps                                         2e+04

Logging data to data/local/experiments/tonic/swimmer-swim/mlp_256/log.csv
          Time left:  epoch 0:00:18  total 0:07:23          
          Time left:  epoch 0:00:18  total 0:07:44          
          Time left:  epoch 0:00:17  total 0:07:36          
          Time left:  epoch 0:00:16  total 0:07:29          
          Time left:  epoch 0:00:15  total 0:07:22          
          Time left:  epoch 0:00:14  total 0:07:16          
          Time left:  epoch 0:00:14  total 0:07:10          
          Time left:  epoch 0:00:14  total 0:07:28          
          Time left:  epoch 0:00:13  total 0:07:22          
          Time left:  epoch 0:00:12  total 0:07:16          
          Time left:  epoch 0:00:11  total 0:07:11          
          Time left:  epoch 0:00:11  total 0:07:05          
          Time left:  epoch 0:00:10  total 0:07:00          
          Time left:  epoch 0:00:10  total 0:06:58          
          Time left:  epoch 0:00:09  total 0:07:13          
          Time left:  epoch 0:00:09  total 0:07:08          
          Time left:  epoch 0:00:08  total 0:07:03          
          Time left:  epoch 0:00:07  total 0:06:59          
          Time left:  epoch 0:00:07  total 0:06:54          
          Time left:  epoch 0:00:06  total 0:06:50          
          Time left:  epoch 0:00:06  total 0:07:03          
          Time left:  epoch 0:00:05  total 0:06:59          
          Time left:  epoch 0:00:04  total 0:06:55          
          Time left:  epoch 0:00:04  total 0:06:50          
          Time left:  epoch 0:00:03  total 0:06:47          
          Time left:  epoch 0:00:02  total 0:06:43          
          Time left:  epoch 0:00:02  total 0:06:55          
          Time left:  epoch 0:00:01  total 0:06:51          
          Time left:  epoch 0:00:01  total 0:06:47          
          Time left:  epoch 0:00:00  total 0:06:44          
          Time left:  epoch 0:00:00  total 0:06:40          

actor                                                       
  clip fraction                                        0.118
  entropy                                               1.05
  iterations                                             4.4
  kl                                                 0.00744
  loss                                              -0.00654
  std                                                  0.691
  stop                                                 0.227
critic                                                      
  iterations                                              80
  loss                                                  3.29
  v                                                     21.9
test                                                        
  action                                                    
    max                                                  2.9
    mean                                             -0.0215
    min                                                -3.42
    size                                               5,000
    std                                                0.798
  episode length                                            
    max                                                1,000
    mean                                               1e+03
    min                                                1,000
    size                                                   5
    std                                                    0
  episode score                                             
    max                                                  491
    mean                                                 473
    min                                                  461
    size                                                   5
    std                                                 11.1
train                                                       
  action                                                    
    max                                                 3.46
    mean                                              -0.022
    min                                                -3.24
    size                                              20,000
    std                                                0.762
  episode length                                            
    max                                                1,000
    mean                                               1e+03
    min                                                1,000
    size                                                  20
    std                                                    0
  episode score                                             
    max                                                  459
    mean                                                 364
    min                                                  236
    size                                                  20
    std                                                 72.8
  episodes                                                40
  epoch seconds                                         18.3
  epoch steps                                          2e+04
  epochs                                                   2
  seconds                                               36.9
  steps                                                4e+04
  steps per second                                  1.09e+03
  worker steps                                         4e+04

          Time left:  epoch 0:00:18  total 0:07:03          
          Time left:  epoch 0:00:17  total 0:07:01          
          Time left:  epoch 0:00:17  total 0:07:11          
          Time left:  epoch 0:00:16  total 0:07:07          
          Time left:  epoch 0:00:16  total 0:07:03          
          Time left:  epoch 0:00:15  total 0:07:00          
          Time left:  epoch 0:00:14  total 0:06:56          
          Time left:  epoch 0:00:13  total 0:06:53          
          Time left:  epoch 0:00:13  total 0:06:51          
          Time left:  epoch 0:00:13  total 0:07:00          
          Time left:  epoch 0:00:12  total 0:06:57          
          Time left:  epoch 0:00:11  total 0:06:53          
          Time left:  epoch 0:00:11  total 0:06:50          
          Time left:  epoch 0:00:10  total 0:06:47          
          Time left:  epoch 0:00:09  total 0:06:44          
          Time left:  epoch 0:00:09  total 0:06:52          
          Time left:  epoch 0:00:08  total 0:06:49          
          Time left:  epoch 0:00:08  total 0:06:46          
          Time left:  epoch 0:00:07  total 0:06:43          
          Time left:  epoch 0:00:06  total 0:06:40          
          Time left:  epoch 0:00:06  total 0:06:37          
          Time left:  epoch 0:00:05  total 0:06:45          
          Time left:  epoch 0:00:05  total 0:06:42          
          Time left:  epoch 0:00:04  total 0:06:40          
          Time left:  epoch 0:00:03  total 0:06:37          
          Time left:  epoch 0:00:03  total 0:06:34          
          Time left:  epoch 0:00:02  total 0:06:31          
          Time left:  epoch 0:00:02  total 0:06:30          
          Time left:  epoch 0:00:01  total 0:06:37          
          Time left:  epoch 0:00:01  total 0:06:34          
          Time left:  epoch 0:00:00  total 0:06:32          
          Time left:  epoch 0:00:00  total 0:06:29          

actor                                                       
  clip fraction                                       0.0987
  entropy                                               1.05
  iterations                                             4.6
  kl                                                 0.00775
  loss                                              -0.00415
  std                                                  0.688
  stop                                                 0.217
critic                                                      
  iterations                                              80
  loss                                                  4.02
  v                                                     42.2
test                                                        
  action                                                    
    max                                                 3.45
    mean                                            -0.00164
    min                                                -3.45
    size                                               5,000
    std                                                 0.88
  episode length                                            
    max                                                1,000
    mean                                               1e+03
    min                                                1,000
    size                                                   5
    std                                                    0
  episode score                                             
    max                                                  659
    mean                                                 626
    min                                                  600
    size                                                   5
    std                                                 21.8
train                                                       
  action                                                    
    max                                                 3.27
    mean                                             -0.0133
    min                                                -3.38
    size                                              20,000
    std                                                 0.83
  episode length                                            
    max                                                1,000
    mean                                               1e+03
    min                                                1,000
    size                                                  20
    std                                                    0
  episode score                                             
    max                                                  624
    mean                                                 553
    min                                                  451
    size                                                  20
    std                                                 49.1
  episodes                                                60
  epoch seconds                                         18.2
  epoch steps                                          2e+04
  epochs                                                   3
  seconds                                               55.2
  steps                                                6e+04
  steps per second                                   1.1e+03
  worker steps                                         6e+04

          Time left:  epoch 0:00:18  total 0:06:43          
          Time left:  epoch 0:00:17  total 0:06:41          
          Time left:  epoch 0:00:17  total 0:06:39          
          Time left:  epoch 0:00:16  total 0:06:45          
          Time left:  epoch 0:00:15  total 0:06:43          
          Time left:  epoch 0:00:15  total 0:06:40          
          Time left:  epoch 0:00:14  total 0:06:39          
          Time left:  epoch 0:00:14  total 0:06:38          
          Time left:  epoch 0:00:13  total 0:06:35          
          Time left:  epoch 0:00:13  total 0:06:34          
          Time left:  epoch 0:00:12  total 0:06:40          
          Time left:  epoch 0:00:12  total 0:06:37          
          Time left:  epoch 0:00:11  total 0:06:35          
          Time left:  epoch 0:00:10  total 0:06:32          
          Time left:  epoch 0:00:10  total 0:06:30          
          Time left:  epoch 0:00:09  total 0:06:28          
          Time left:  epoch 0:00:09  total 0:06:33          
          Time left:  epoch 0:00:08  total 0:06:31          
          Time left:  epoch 0:00:07  total 0:06:29          
          Time left:  epoch 0:00:07  total 0:06:26          
          Time left:  epoch 0:00:06  total 0:06:24          
          Time left:  epoch 0:00:05  total 0:06:22          
          Time left:  epoch 0:00:05  total 0:06:21          
          Time left:  epoch 0:00:05  total 0:06:27          
          Time left:  epoch 0:00:04  total 0:06:25          
          Time left:  epoch 0:00:03  total 0:06:22          
          Time left:  epoch 0:00:03  total 0:06:20          
          Time left:  epoch 0:00:02  total 0:06:18          
          Time left:  epoch 0:00:02  total 0:06:16          
          Time left:  epoch 0:00:01  total 0:06:21          
          Time left:  epoch 0:00:00  total 0:06:19          
          Time left:  epoch 0:00:00  total 0:06:17          
          Time left:  epoch 0:00:00  total 0:06:16          

actor                                                       
  clip fraction                                       0.0912
  entropy                                               1.04
  iterations                                               5
  kl                                                 0.00792
  loss                                                -0.006
  std                                                  0.685
  stop                                                   0.2
critic                                                      
  iterations                                              80
  loss                                                  2.62
  v                                                     59.6
test                                                        
  action                                                    
    max                                                 3.63
    mean                                               -0.01
    min                                                -3.52
    size                                               5,000
    std                                                0.985
  episode length                                            
    max                                                1,000
    mean                                               1e+03
    min                                                1,000
    size                                                   5
    std                                                    0
  episode score                                             
    max                                                  772
    mean                                                 752
    min                                                  735
    size                                                   5
    std                                                 12.1
train                                                       
  action                                                    
    max                                                 3.75
    mean                                            -0.00041
    min                                                -3.33
    size                                              20,000
    std                                                0.929
  episode length                                            
    max                                                1,000
    mean                                               1e+03
    min                                                1,000
    size                                                  20
    std                                                    0
  episode score                                             
    max                                                  735
    mean                                                 686
    min                                                  618
    size                                                  20
    std                                                 32.5
  episodes                                                80
  epoch seconds                                         18.6
  epoch steps                                          2e+04
  epochs                                                   4
  seconds                                               73.9
  steps                                                8e+04
  steps per second                                  1.08e+03
  worker steps                                         8e+04

          Time left:  epoch 0:00:18  total 0:06:26          
          Time left:  epoch 0:00:17  total 0:06:24          
          Time left:  epoch 0:00:16  total 0:06:22          
          Time left:  epoch 0:00:16  total 0:06:27          
          Time left:  epoch 0:00:15  total 0:06:25          
          Time left:  epoch 0:00:15  total 0:06:23          
          Time left:  epoch 0:00:14  total 0:06:21          
          Time left:  epoch 0:00:13  total 0:06:19          
          Time left:  epoch 0:00:13  total 0:06:17          
          Time left:  epoch 0:00:12  total 0:06:16          
          Time left:  epoch 0:00:12  total 0:06:26          
          Time left:  epoch 0:00:11  total 0:06:24          
          Time left:  epoch 0:00:11  total 0:06:22          
          Time left:  epoch 0:00:10  total 0:06:20          
          Time left:  epoch 0:00:09  total 0:06:18          
          Time left:  epoch 0:00:09  total 0:06:16          
          Time left:  epoch 0:00:08  total 0:06:20          
          Time left:  epoch 0:00:08  total 0:06:18          
          Time left:  epoch 0:00:07  total 0:06:17          
          Time left:  epoch 0:00:06  total 0:06:15          
          Time left:  epoch 0:00:06  total 0:06:13          
          Time left:  epoch 0:00:05  total 0:06:11          
          Time left:  epoch 0:00:05  total 0:06:20          
          Time left:  epoch 0:00:04  total 0:06:18          
          Time left:  epoch 0:00:03  total 0:06:16          
          Time left:  epoch 0:00:03  total 0:06:14          
          Time left:  epoch 0:00:02  total 0:06:13          
          Time left:  epoch 0:00:01  total 0:06:11          
          Time left:  epoch 0:00:01  total 0:06:19          
          Time left:  epoch 0:00:00  total 0:06:17          
          Time left:  epoch 0:00:00  total 0:06:15          

actor                                                       
  clip fraction                                        0.113
  entropy                                               1.01
  iterations                                            50.4
  kl                                                 0.00915
  loss                                               -0.0158
  std                                                  0.665
  stop                                               0.00794
critic                                                      
  iterations                                              80
  loss                                                  1.64
  v                                                     72.1
test                                                        
  action                                                    
    max                                                 3.17
    mean                                            -0.00699
    min                                                -3.67
    size                                               5,000
    std                                                    1
  episode length                                            
    max                                                1,000
    mean                                               1e+03
    min                                                1,000
    size                                                   5
    std                                                    0
  episode score                                             
    max                                                  815
    mean                                                 809
    min                                                  798
    size                                                   5
    std                                                 6.37
train                                                       
  action                                                    
    max                                                 3.62
    mean                                            -0.00871
    min                                                 -3.5
    size                                              20,000
    std                                                    1
  episode length                                            
    max                                                1,000
    mean                                               1e+03
    min                                                1,000
    size                                                  20
    std                                                    0
  episode score                                             
    max                                                  811
    mean                                                 774
    min                                                  730
    size                                                  20
    std                                                 19.2
  episodes                                               100
  epoch seconds                                         22.1
  epoch steps                                          2e+04
  epochs                                                   5
  seconds                                                 96
  steps                                                1e+05
  steps per second                                       904
  worker steps                                         1e+05


Saving weights to data/local/experiments/tonic/swimmer-swim/mlp_256/checkpoints/step_100000.pt
          Time left:  epoch 0:00:18  total 0:06:23          
          Time left:  epoch 0:00:18  total 0:06:21          
          Time left:  epoch 0:00:17  total 0:06:19          
          Time left:  epoch 0:00:16  total 0:06:17          
          Time left:  epoch 0:00:16  total 0:06:20          
          Time left:  epoch 0:00:15  total 0:06:18          
          Time left:  epoch 0:00:14  total 0:06:17          
          Time left:  epoch 0:00:14  total 0:06:15          
          Time left:  epoch 0:00:13  total 0:06:13          
          Time left:  epoch 0:00:12  total 0:06:11          
          Time left:  epoch 0:00:12  total 0:06:19          
          Time left:  epoch 0:00:11  total 0:06:17          
          Time left:  epoch 0:00:11  total 0:06:15          
          Time left:  epoch 0:00:10  total 0:06:13          
          Time left:  epoch 0:00:09  total 0:06:12          
          Time left:  epoch 0:00:09  total 0:06:10          
          Time left:  epoch 0:00:08  total 0:06:17          
          Time left:  epoch 0:00:08  total 0:06:15          
          Time left:  epoch 0:00:07  total 0:06:13          
          Time left:  epoch 0:00:06  total 0:06:12          
          Time left:  epoch 0:00:06  total 0:06:10          
          Time left:  epoch 0:00:05  total 0:06:08          
          Time left:  epoch 0:00:05  total 0:06:07          
          Time left:  epoch 0:00:04  total 0:06:14          
          Time left:  epoch 0:00:03  total 0:06:12          
          Time left:  epoch 0:00:03  total 0:06:10          
          Time left:  epoch 0:00:02  total 0:06:09          
          Time left:  epoch 0:00:01  total 0:06:07          
          Time left:  epoch 0:00:01  total 0:06:05          
          Time left:  epoch 0:00:00  total 0:06:11          
          Time left:  epoch 0:00:00  total 0:06:10          

actor                                                       
  clip fraction                                        0.109
  entropy                                              0.969
  iterations                                            64.8
  kl                                                 0.00918
  loss                                               -0.0163
  std                                                  0.638
  stop                                               0.00309
critic                                                      
  iterations                                              80
  loss                                                  0.94
  v                                                     79.8
test                                                        
  action                                                    
    max                                                 3.26
    mean                                            -0.00105
    min                                                -3.13
    size                                               5,000
    std                                                 1.04
  episode length                                            
    max                                                1,000
    mean                                               1e+03
    min                                                1,000
    size                                                   5
    std                                                    0
  episode score                                             
    max                                                  862
    mean                                                 850
    min                                                  821
    size                                                   5
    std                                                 14.9
train                                                       
  action                                                    
    max                                                 3.69
    mean                                            -0.00322
    min                                                -3.73
    size                                              20,000
    std                                                 1.03
  episode length                                            
    max                                                1,000
    mean                                               1e+03
    min                                                1,000
    size                                                  20
    std                                                    0
  episode score                                             
    max                                                  868
    mean                                                 825
    min                                                  781
    size                                                  20
    std                                                 22.6
  episodes                                               120
  epoch seconds                                         22.9
  epoch steps                                          2e+04
  epochs                                                   6
  seconds                                                119
  steps                                              1.2e+05
  steps per second                                       873
  worker steps                                       1.2e+05

          Time left:  epoch 0:00:19  total 0:06:16          
          Time left:  epoch 0:00:18  total 0:06:14          
          Time left:  epoch 0:00:18  total 0:06:12          
          Time left:  epoch 0:00:17  total 0:06:11          
          Time left:  epoch 0:00:17  total 0:06:10          
          Time left:  epoch 0:00:16  total 0:06:16          
          Time left:  epoch 0:00:15  total 0:06:14          
          Time left:  epoch 0:00:15  total 0:06:12          
          Time left:  epoch 0:00:14  total 0:06:11          
          Time left:  epoch 0:00:13  total 0:06:09          
          Time left:  epoch 0:00:13  total 0:06:07          
          Time left:  epoch 0:00:12  total 0:06:13          
          Time left:  epoch 0:00:11  total 0:06:11          
          Time left:  epoch 0:00:11  total 0:06:10          
          Time left:  epoch 0:00:10  total 0:06:08          
          Time left:  epoch 0:00:09  total 0:06:07          
          Time left:  epoch 0:00:09  total 0:06:05          
          Time left:  epoch 0:00:08  total 0:06:04          
          Time left:  epoch 0:00:08  total 0:06:10          
          Time left:  epoch 0:00:07  total 0:06:08          
          Time left:  epoch 0:00:06  total 0:06:06          
          Time left:  epoch 0:00:06  total 0:06:05          
          Time left:  epoch 0:00:05  total 0:06:03          
          Time left:  epoch 0:00:04  total 0:06:02          
          Time left:  epoch 0:00:04  total 0:06:07          
          Time left:  epoch 0:00:03  total 0:06:05          
          Time left:  epoch 0:00:03  total 0:06:04          
          Time left:  epoch 0:00:02  total 0:06:02          
          Time left:  epoch 0:00:01  total 0:06:01          
          Time left:  epoch 0:00:00  total 0:05:59          
          Time left:  epoch 0:00:00  total 0:06:03          
          Time left:  epoch 0:00:00  total 0:06:02          

actor                                                       
  clip fraction                                        0.103
  entropy                                              0.905
  iterations                                            75.2
  kl                                                 0.00787
  loss                                               -0.0172
  std                                                  0.598
  stop                                               0.00266
critic                                                      
  iterations                                              80
  loss                                                 0.515
  v                                                     85.8
test                                                        
  action                                                    
    max                                                 3.26
    mean                                             0.00292
    min                                                -2.94
    size                                               5,000
    std                                                 1.03
  episode length                                            
    max                                                1,000
    mean                                               1e+03
    min                                                1,000
    size                                                   5
    std                                                    0
  episode score                                             
    max                                                  898
    mean                                                 892
    min                                                  879
    size                                                   5
    std                                                 6.53
train                                                       
  action                                                    
    max                                                 3.95
    mean                                             0.00156
    min                                                -3.77
    size                                              20,000
    std                                                 1.03
  episode length                                            
    max                                                1,000
    mean                                               1e+03
    min                                                1,000
    size                                                  20
    std                                                    0
  episode score                                             
    max                                                  896
    mean                                                 873
    min                                                  838
    size                                                  20
    std                                                 15.6
  episodes                                               140
  epoch seconds                                         24.1
  epoch steps                                          2e+04
  epochs                                                   7
  seconds                                                143
  steps                                              1.4e+05
  steps per second                                       831
  worker steps                                       1.4e+05

          Time left:  epoch 0:00:20  total 0:06:07          
          Time left:  epoch 0:00:19  total 0:06:05          
          Time left:  epoch 0:00:18  total 0:06:04          
          Time left:  epoch 0:00:17  total 0:06:02          
          Time left:  epoch 0:00:17  total 0:06:01          
          Time left:  epoch 0:00:16  total 0:06:00          
          Time left:  epoch 0:00:16  total 0:06:04          
          Time left:  epoch 0:00:15  total 0:06:03          
          Time left:  epoch 0:00:14  total 0:06:01          
          Time left:  epoch 0:00:14  total 0:05:59          
          Time left:  epoch 0:00:13  total 0:05:58          
          Time left:  epoch 0:00:12  total 0:05:56          
          Time left:  epoch 0:00:12  total 0:05:58          
          Time left:  epoch 0:00:11  total 0:05:56          
          Time left:  epoch 0:00:10  total 0:05:55          
          Time left:  epoch 0:00:10  total 0:05:53          
          Time left:  epoch 0:00:09  total 0:05:52          
          Time left:  epoch 0:00:08  total 0:05:50          
          Time left:  epoch 0:00:08  total 0:05:52          
          Time left:  epoch 0:00:07  total 0:05:51          
          Time left:  epoch 0:00:06  total 0:05:49          
          Time left:  epoch 0:00:06  total 0:05:48          
          Time left:  epoch 0:00:05  total 0:05:46          
          Time left:  epoch 0:00:04  total 0:05:45          
          Time left:  epoch 0:00:04  total 0:05:47          
          Time left:  epoch 0:00:03  total 0:05:45          
          Time left:  epoch 0:00:02  total 0:05:44          
          Time left:  epoch 0:00:02  total 0:05:42          
          Time left:  epoch 0:00:01  total 0:05:41          
          Time left:  epoch 0:00:00  total 0:05:39          
          Time left:  epoch 0:00:00  total 0:05:39          
          Time left:  epoch 0:00:00  total 0:05:43          

actor                                                       
  clip fraction                                        0.101
  entropy                                              0.837
  iterations                                              34
  kl                                                  0.0103
  loss                                               -0.0161
  std                                                  0.559
  stop                                                0.0235
critic                                                      
  iterations                                              80
  loss                                                  0.29
  v                                                     89.9
test                                                        
  action                                                    
    max                                                 3.35
    mean                                             0.00267
    min                                                -3.22
    size                                               5,000
    std                                                 1.03
  episode length                                            
    max                                                1,000
    mean                                               1e+03
    min                                                1,000
    size                                                   5
    std                                                    0
  episode score                                             
    max                                                  922
    mean                                                 905
    min                                                  897
    size                                                   5
    std                                                    9
train                                                       
  action                                                    
    max                                                 3.21
    mean                                             0.00226
    min                                                -3.51
    size                                              20,000
    std                                                 1.03
  episode length                                            
    max                                                1,000
    mean                                               1e+03
    min                                                1,000
    size                                                  20
    std                                                    0
  episode score                                             
    max                                                  919
    mean                                                 902
    min                                                  878
    size                                                  20
    std                                                 8.57
  episodes                                               160
  epoch seconds                                           21
  epoch steps                                          2e+04
  epochs                                                   8
  seconds                                                164
  steps                                              1.6e+05
  steps per second                                       954
  worker steps                                       1.6e+05

          Time left:  epoch 0:00:20  total 0:05:48          
          Time left:  epoch 0:00:19  total 0:05:46          
          Time left:  epoch 0:00:18  total 0:05:45          
          Time left:  epoch 0:00:17  total 0:05:43          
          Time left:  epoch 0:00:17  total 0:05:42          
          Time left:  epoch 0:00:16  total 0:05:40          
          Time left:  epoch 0:00:16  total 0:05:44          
          Time left:  epoch 0:00:15  total 0:05:43          
          Time left:  epoch 0:00:14  total 0:05:41          
          Time left:  epoch 0:00:13  total 0:05:40          
          Time left:  epoch 0:00:13  total 0:05:39          
          Time left:  epoch 0:00:12  total 0:05:37          
          Time left:  epoch 0:00:12  total 0:05:41          
          Time left:  epoch 0:00:11  total 0:05:39          
          Time left:  epoch 0:00:10  total 0:05:38          
          Time left:  epoch 0:00:09  total 0:05:37          
          Time left:  epoch 0:00:09  total 0:05:35          
          Time left:  epoch 0:00:08  total 0:05:34          
          Time left:  epoch 0:00:08  total 0:05:33          
          Time left:  epoch 0:00:07  total 0:05:37          
          Time left:  epoch 0:00:06  total 0:05:35          
          Time left:  epoch 0:00:06  total 0:05:34          
          Time left:  epoch 0:00:05  total 0:05:33          
          Time left:  epoch 0:00:04  total 0:05:31          
          Time left:  epoch 0:00:04  total 0:05:30          
          Time left:  epoch 0:00:03  total 0:05:32          
          Time left:  epoch 0:00:02  total 0:05:30          
          Time left:  epoch 0:00:02  total 0:05:29          
          Time left:  epoch 0:00:01  total 0:05:28          
          Time left:  epoch 0:00:00  total 0:05:26          
          Time left:  epoch 0:00:00  total 0:05:25          

actor                                                       
  clip fraction                                        0.105
  entropy                                              0.791
  iterations                                            65.5
  kl                                                  0.0104
  loss                                               -0.0146
  std                                                  0.534
  stop                                               0.00382
critic                                                      
  iterations                                              80
  loss                                                 0.188
  v                                                       92
test                                                        
  action                                                    
    max                                                 2.96
    mean                                            -0.00476
    min                                                -2.76
    size                                               5,000
    std                                                 1.02
  episode length                                            
    max                                                1,000
    mean                                               1e+03
    min                                                1,000
    size                                                   5
    std                                                    0
  episode score                                             
    max                                                  925
    mean                                                 920
    min                                                  915
    size                                                   5
    std                                                 4.66
train                                                       
  action                                                    
    max                                                 3.41
    mean                                            -0.00266
    min                                                -3.15
    size                                              20,000
    std                                                 1.02
  episode length                                            
    max                                                1,000
    mean                                               1e+03
    min                                                1,000
    size                                                  20
    std                                                    0
  episode score                                             
    max                                                  925
    mean                                                 918
    min                                                  906
    size                                                  20
    std                                                 4.89
  episodes                                               180
  epoch seconds                                           21
  epoch steps                                          2e+04
  epochs                                                   9
  seconds                                                185
  steps                                              1.8e+05
  steps per second                                       953
  worker steps                                       1.8e+05

          Time left:  epoch 0:00:20  total 0:05:29          
          Time left:  epoch 0:00:20  total 0:05:32          
          Time left:  epoch 0:00:19  total 0:05:30          
          Time left:  epoch 0:00:18  total 0:05:29          
          Time left:  epoch 0:00:17  total 0:05:28          
          Time left:  epoch 0:00:17  total 0:05:26          
          Time left:  epoch 0:00:16  total 0:05:25          
          Time left:  epoch 0:00:15  total 0:05:28          
          Time left:  epoch 0:00:15  total 0:05:27          
          Time left:  epoch 0:00:14  total 0:05:25          
          Time left:  epoch 0:00:13  total 0:05:24          
          Time left:  epoch 0:00:13  total 0:05:23          
          Time left:  epoch 0:00:12  total 0:05:21          
          Time left:  epoch 0:00:12  total 0:05:21          
          Time left:  epoch 0:00:11  total 0:05:23          
          Time left:  epoch 0:00:10  total 0:05:22          
          Time left:  epoch 0:00:10  total 0:05:21          
          Time left:  epoch 0:00:09  total 0:05:19          
          Time left:  epoch 0:00:08  total 0:05:18          
          Time left:  epoch 0:00:07  total 0:05:17          
          Time left:  epoch 0:00:07  total 0:05:17          
          Time left:  epoch 0:00:06  total 0:05:16          
          Time left:  epoch 0:00:05  total 0:05:15          
          Time left:  epoch 0:00:05  total 0:05:14          
          Time left:  epoch 0:00:04  total 0:05:12          
          Time left:  epoch 0:00:03  total 0:05:11          
          Time left:  epoch 0:00:03  total 0:05:14          
          Time left:  epoch 0:00:02  total 0:05:13          
          Time left:  epoch 0:00:01  total 0:05:11          
          Time left:  epoch 0:00:01  total 0:05:10          
          Time left:  epoch 0:00:00  total 0:05:09          
          Time left:  epoch 0:00:00  total 0:05:08          

actor                                                       
  clip fraction                                        0.098
  entropy                                              0.721
  iterations                                              60
  kl                                                 0.00927
  loss                                               -0.0155
  std                                                  0.498
  stop                                               0.00667
critic                                                      
  iterations                                              80
  loss                                                 0.134
  v                                                     93.8
test                                                        
  action                                                    
    max                                                 3.42
    mean                                            0.000662
    min                                                -2.84
    size                                               5,000
    std                                                 1.01
  episode length                                            
    max                                                1,000
    mean                                               1e+03
    min                                                1,000
    size                                                   5
    std                                                    0
  episode score                                             
    max                                                  954
    mean                                                 945
    min                                                  936
    size                                                   5
    std                                                 6.22
train                                                       
  action                                                    
    max                                                 3.18
    mean                                            -0.00384
    min                                                -2.98
    size                                              20,000
    std                                                 1.01
  episode length                                            
    max                                                1,000
    mean                                               1e+03
    min                                                1,000
    size                                                  20
    std                                                    0
  episode score                                             
    max                                                  954
    mean                                                 934
    min                                                  916
    size                                                  20
    std                                                 9.55
  episodes                                               200
  epoch seconds                                         22.7
  epoch steps                                          2e+04
  epochs                                                  10
  seconds                                                208
  steps                                                2e+05
  steps per second                                       880
  worker steps                                         2e+05


Saving weights to data/local/experiments/tonic/swimmer-swim/mlp_256/checkpoints/step_200000.pt
          Time left:  epoch 0:00:20  total 0:05:11          
          Time left:  epoch 0:00:20  total 0:05:10          
          Time left:  epoch 0:00:19  total 0:05:11          
          Time left:  epoch 0:00:18  total 0:05:10          
          Time left:  epoch 0:00:18  total 0:05:08          
          Time left:  epoch 0:00:17  total 0:05:07          
          Time left:  epoch 0:00:16  total 0:05:06          
          Time left:  epoch 0:00:15  total 0:05:05          
          Time left:  epoch 0:00:15  total 0:05:07          
          Time left:  epoch 0:00:14  total 0:05:06          
          Time left:  epoch 0:00:13  total 0:05:05          
          Time left:  epoch 0:00:13  total 0:05:03          
          Time left:  epoch 0:00:12  total 0:05:02          
          Time left:  epoch 0:00:11  total 0:05:01          
          Time left:  epoch 0:00:11  total 0:05:03          
          Time left:  epoch 0:00:10  total 0:05:02          
          Time left:  epoch 0:00:09  total 0:05:01          
          Time left:  epoch 0:00:09  total 0:05:00          
          Time left:  epoch 0:00:08  total 0:04:58          
          Time left:  epoch 0:00:07  total 0:04:57          
          Time left:  epoch 0:00:06  total 0:04:59          
          Time left:  epoch 0:00:06  total 0:04:58          
          Time left:  epoch 0:00:05  total 0:04:57          
          Time left:  epoch 0:00:04  total 0:04:55          
          Time left:  epoch 0:00:04  total 0:04:54          
          Time left:  epoch 0:00:03  total 0:04:53          
          Time left:  epoch 0:00:03  total 0:04:53          
          Time left:  epoch 0:00:02  total 0:04:55          
          Time left:  epoch 0:00:01  total 0:04:53          
          Time left:  epoch 0:00:01  total 0:04:52          
          Time left:  epoch 0:00:00  total 0:04:51          
          Time left:  epoch 0:00:00  total 0:04:50          

actor                                                       
  clip fraction                                       0.0983
  entropy                                              0.644
  iterations                                              63
  kl                                                 0.00938
  loss                                               -0.0161
  std                                                  0.461
  stop                                               0.00635
critic                                                      
  iterations                                              80
  loss                                                0.0992
  v                                                     95.2
test                                                        
  action                                                    
    max                                                 2.86
    mean                                            -0.00174
    min                                                -2.82
    size                                               5,000
    std                                                0.991
  episode length                                            
    max                                                1,000
    mean                                               1e+03
    min                                                1,000
    size                                                   5
    std                                                    0
  episode score                                             
    max                                                  957
    mean                                                 950
    min                                                  945
    size                                                   5
    std                                                  4.1
train                                                       
  action                                                    
    max                                                    3
    mean                                           -0.000557
    min                                                -3.11
    size                                              20,000
    std                                                0.997
  episode length                                            
    max                                                1,000
    mean                                               1e+03
    min                                                1,000
    size                                                  20
    std                                                    0
  episode score                                             
    max                                                  957
    mean                                                 943
    min                                                  929
    size                                                  20
    std                                                 7.35
  episodes                                               220
  epoch seconds                                         22.7
  epoch steps                                          2e+04
  epochs                                                  11
  seconds                                                231
  steps                                              2.2e+05
  steps per second                                       880
  worker steps                                       2.2e+05

          Time left:  epoch 0:00:20  total 0:04:53          
          Time left:  epoch 0:00:19  total 0:04:51          
          Time left:  epoch 0:00:19  total 0:04:52          
          Time left:  epoch 0:00:18  total 0:04:51          
          Time left:  epoch 0:00:17  total 0:04:50          
          Time left:  epoch 0:00:17  total 0:04:49          
          Time left:  epoch 0:00:16  total 0:04:48          
          Time left:  epoch 0:00:15  total 0:04:46          
          Time left:  epoch 0:00:15  total 0:04:47          
          Time left:  epoch 0:00:14  total 0:04:46          
          Time left:  epoch 0:00:13  total 0:04:45          
          Time left:  epoch 0:00:12  total 0:04:43          
          Time left:  epoch 0:00:12  total 0:04:42          
          Time left:  epoch 0:00:11  total 0:04:41          
          Time left:  epoch 0:00:11  total 0:04:41          
          Time left:  epoch 0:00:10  total 0:04:41          
          Time left:  epoch 0:00:09  total 0:04:40          
          Time left:  epoch 0:00:09  total 0:04:39          
          Time left:  epoch 0:00:08  total 0:04:38          
          Time left:  epoch 0:00:07  total 0:04:37          
          Time left:  epoch 0:00:06  total 0:04:36          
          Time left:  epoch 0:00:06  total 0:04:37          
          Time left:  epoch 0:00:05  total 0:04:36          
          Time left:  epoch 0:00:04  total 0:04:35          
          Time left:  epoch 0:00:04  total 0:04:34          
          Time left:  epoch 0:00:03  total 0:04:33          
          Time left:  epoch 0:00:02  total 0:04:31          
          Time left:  epoch 0:00:02  total 0:04:33          
          Time left:  epoch 0:00:01  total 0:04:32          
          Time left:  epoch 0:00:00  total 0:04:31          
          Time left:  epoch 0:00:00  total 0:04:30          

actor                                                       
  clip fraction                                        0.104
  entropy                                               0.59
  iterations                                            38.8
  kl                                                    0.01
  loss                                               -0.0128
  std                                                  0.437
  stop                                                0.0206
critic                                                      
  iterations                                              80
  loss                                                0.0814
  v                                                     96.3
test                                                        
  action                                                    
    max                                                 2.56
    mean                                            -0.00367
    min                                                -2.52
    size                                               5,000
    std                                                0.988
  episode length                                            
    max                                                1,000
    mean                                               1e+03
    min                                                1,000
    size                                                   5
    std                                                    0
  episode score                                             
    max                                                  969
    mean                                                 959
    min                                                  947
    size                                                   5
    std                                                 7.46
train                                                       
  action                                                    
    max                                                 2.81
    mean                                            -0.00121
    min                                                 -2.7
    size                                              20,000
    std                                                0.989
  episode length                                            
    max                                                1,000
    mean                                               1e+03
    min                                                1,000
    size                                                  20
    std                                                    0
  episode score                                             
    max                                                  962
    mean                                                 952
    min                                                  940
    size                                                  20
    std                                                 5.73
  episodes                                               240
  epoch seconds                                         20.8
  epoch steps                                          2e+04
  epochs                                                  12
  seconds                                                251
  steps                                              2.4e+05
  steps per second                                       963
  worker steps                                       2.4e+05

          Time left:  epoch 0:00:20  total 0:04:31          
          Time left:  epoch 0:00:19  total 0:04:30          
          Time left:  epoch 0:00:19  total 0:04:30          
          Time left:  epoch 0:00:18  total 0:04:30          
          Time left:  epoch 0:00:18  total 0:04:29          
          Time left:  epoch 0:00:17  total 0:04:28          
          Time left:  epoch 0:00:16  total 0:04:27          
          Time left:  epoch 0:00:15  total 0:04:26          
          Time left:  epoch 0:00:15  total 0:04:25          
          Time left:  epoch 0:00:14  total 0:04:24          
          Time left:  epoch 0:00:14  total 0:04:24          
          Time left:  epoch 0:00:13  total 0:04:23          
          Time left:  epoch 0:00:12  total 0:04:22          
          Time left:  epoch 0:00:12  total 0:04:21          
          Time left:  epoch 0:00:11  total 0:04:20          
          Time left:  epoch 0:00:10  total 0:04:19          
          Time left:  epoch 0:00:10  total 0:04:19          
          Time left:  epoch 0:00:09  total 0:04:18          
          Time left:  epoch 0:00:08  total 0:04:17          
          Time left:  epoch 0:00:07  total 0:04:16          
          Time left:  epoch 0:00:07  total 0:04:15          
          Time left:  epoch 0:00:06  total 0:04:14          
          Time left:  epoch 0:00:05  total 0:04:15          
          Time left:  epoch 0:00:05  total 0:04:14          
          Time left:  epoch 0:00:04  total 0:04:13          
          Time left:  epoch 0:00:03  total 0:04:12          
          Time left:  epoch 0:00:03  total 0:04:11          
          Time left:  epoch 0:00:02  total 0:04:10          
          Time left:  epoch 0:00:02  total 0:04:09          
          Time left:  epoch 0:00:01  total 0:04:10          
          Time left:  epoch 0:00:00  total 0:04:09          
          Time left:  epoch 0:00:00  total 0:04:08          

actor                                                       
  clip fraction                                       0.0908
  entropy                                              0.555
  iterations                                              21
  kl                                                  0.0104
  loss                                               -0.0088
  std                                                  0.422
  stop                                                0.0476
critic                                                      
  iterations                                              80
  loss                                                0.0626
  v                                                     96.9
test                                                        
  action                                                    
    max                                                  2.5
    mean                                            -0.00633
    min                                                -2.65
    size                                               5,000
    std                                                0.977
  episode length                                            
    max                                                1,000
    mean                                               1e+03
    min                                                1,000
    size                                                   5
    std                                                    0
  episode score                                             
    max                                                  971
    mean                                                 957
    min                                                  943
    size                                                   5
    std                                                 10.4
train                                                       
  action                                                    
    max                                                 2.62
    mean                                            -0.00646
    min                                                -2.79
    size                                              20,000
    std                                                0.982
  episode length                                            
    max                                                1,000
    mean                                               1e+03
    min                                                1,000
    size                                                  20
    std                                                    0
  episode score                                             
    max                                                  974
    mean                                                 955
    min                                                  935
    size                                                  20
    std                                                 8.69
  episodes                                               260
  epoch seconds                                         19.5
  epoch steps                                          2e+04
  epochs                                                  13
  seconds                                                271
  steps                                              2.6e+05
  steps per second                                  1.03e+03
  worker steps                                       2.6e+05

          Time left:  epoch 0:00:20  total 0:04:09          
          Time left:  epoch 0:00:19  total 0:04:08          
          Time left:  epoch 0:00:19  total 0:04:07          
          Time left:  epoch 0:00:18  total 0:04:07          
          Time left:  epoch 0:00:18  total 0:04:07          
          Time left:  epoch 0:00:17  total 0:04:06          
          Time left:  epoch 0:00:16  total 0:04:05          
          Time left:  epoch 0:00:15  total 0:04:04          
          Time left:  epoch 0:00:15  total 0:04:03          
          Time left:  epoch 0:00:14  total 0:04:02          
          Time left:  epoch 0:00:13  total 0:04:03          
          Time left:  epoch 0:00:13  total 0:04:02          
          Time left:  epoch 0:00:12  total 0:04:01          
          Time left:  epoch 0:00:11  total 0:04:00          
          Time left:  epoch 0:00:11  total 0:03:59          
          Time left:  epoch 0:00:10  total 0:03:58          
          Time left:  epoch 0:00:10  total 0:03:58          
          Time left:  epoch 0:00:09  total 0:03:59          
          Time left:  epoch 0:00:08  total 0:03:58          
          Time left:  epoch 0:00:07  total 0:03:57          
          Time left:  epoch 0:00:07  total 0:03:56          
          Time left:  epoch 0:00:06  total 0:03:55          
          Time left:  epoch 0:00:05  total 0:03:54          
          Time left:  epoch 0:00:05  total 0:03:54          
          Time left:  epoch 0:00:04  total 0:03:53          
          Time left:  epoch 0:00:03  total 0:03:52          
          Time left:  epoch 0:00:03  total 0:03:51          
          Time left:  epoch 0:00:02  total 0:03:50          
          Time left:  epoch 0:00:01  total 0:03:49          
          Time left:  epoch 0:00:01  total 0:03:49          
          Time left:  epoch 0:00:00  total 0:03:48          
          Time left:  epoch 0:00:00  total 0:03:48          

actor                                                       
  clip fraction                                       0.0949
  entropy                                              0.518
  iterations                                            40.6
  kl                                                  0.0101
  loss                                               -0.0107
  std                                                  0.407
  stop                                                0.0197
critic                                                      
  iterations                                              80
  loss                                                0.0463
  v                                                     97.4
test                                                        
  action                                                    
    max                                                 2.45
    mean                                            0.000325
    min                                                -2.63
    size                                               5,000
    std                                                0.968
  episode length                                            
    max                                                1,000
    mean                                               1e+03
    min                                                1,000
    size                                                   5
    std                                                    0
  episode score                                             
    max                                                  974
    mean                                                 960
    min                                                  950
    size                                                   5
    std                                                 8.39
train                                                       
  action                                                    
    max                                                 2.74
    mean                                             0.00176
    min                                                -2.78
    size                                              20,000
    std                                                0.974
  episode length                                            
    max                                                1,000
    mean                                               1e+03
    min                                                1,000
    size                                                  20
    std                                                    0
  episode score                                             
    max                                                  974
    mean                                                 962
    min                                                  950
    size                                                  20
    std                                                 7.27
  episodes                                               280
  epoch seconds                                         21.4
  epoch steps                                          2e+04
  epochs                                                  14
  seconds                                                292
  steps                                              2.8e+05
  steps per second                                       935
  worker steps                                       2.8e+05

          Time left:  epoch 0:00:20  total 0:03:49          
          Time left:  epoch 0:00:19  total 0:03:48          
          Time left:  epoch 0:00:19  total 0:03:47          
          Time left:  epoch 0:00:18  total 0:03:46          
          Time left:  epoch 0:00:17  total 0:03:47          
          Time left:  epoch 0:00:17  total 0:03:46          
          Time left:  epoch 0:00:16  total 0:03:45          
          Time left:  epoch 0:00:15  total 0:03:44          
          Time left:  epoch 0:00:14  total 0:03:43          
          Time left:  epoch 0:00:14  total 0:03:42          
          Time left:  epoch 0:00:13  total 0:03:41          
          Time left:  epoch 0:00:13  total 0:03:42          
          Time left:  epoch 0:00:12  total 0:03:41          
          Time left:  epoch 0:00:11  total 0:03:40          
          Time left:  epoch 0:00:11  total 0:03:39          
          Time left:  epoch 0:00:10  total 0:03:38          
          Time left:  epoch 0:00:09  total 0:03:37          
          Time left:  epoch 0:00:09  total 0:03:38          
          Time left:  epoch 0:00:08  total 0:03:37          
          Time left:  epoch 0:00:07  total 0:03:36          
          Time left:  epoch 0:00:06  total 0:03:35          
          Time left:  epoch 0:00:06  total 0:03:34          
          Time left:  epoch 0:00:05  total 0:03:33          
          Time left:  epoch 0:00:04  total 0:03:33          
          Time left:  epoch 0:00:04  total 0:03:32          
          Time left:  epoch 0:00:03  total 0:03:31          
          Time left:  epoch 0:00:02  total 0:03:31          
          Time left:  epoch 0:00:02  total 0:03:30          
          Time left:  epoch 0:00:01  total 0:03:29          
          Time left:  epoch 0:00:01  total 0:03:28          
          Time left:  epoch 0:00:00  total 0:03:29          
          Time left:  epoch 0:00:00  total 0:03:28          

actor                                                       
  clip fraction                                        0.102
  entropy                                              0.465
  iterations                                            62.8
  kl                                                 0.00941
  loss                                               -0.0141
  std                                                  0.386
  stop                                               0.00637
critic                                                      
  iterations                                              80
  loss                                                0.0328
  v                                                     97.8
test                                                        
  action                                                    
    max                                                 2.45
    mean                                             0.00504
    min                                                 -2.5
    size                                               5,000
    std                                                 0.96
  episode length                                            
    max                                                1,000
    mean                                               1e+03
    min                                                1,000
    size                                                   5
    std                                                    0
  episode score                                             
    max                                                  981
    mean                                                 965
    min                                                  951
    size                                                   5
    std                                                 9.67
train                                                       
  action                                                    
    max                                                 2.67
    mean                                             -0.0029
    min                                                   -3
    size                                              20,000
    std                                                0.961
  episode length                                            
    max                                                1,000
    mean                                               1e+03
    min                                                1,000
    size                                                  20
    std                                                    0
  episode score                                             
    max                                                  976
    mean                                                 963
    min                                                  950
    size                                                  20
    std                                                  5.5
  episodes                                               300
  epoch seconds                                         22.7
  epoch steps                                          2e+04
  epochs                                                  15
  seconds                                                315
  steps                                                3e+05
  steps per second                                       883
  worker steps                                         3e+05


Saving weights to data/local/experiments/tonic/swimmer-swim/mlp_256/checkpoints/step_300000.pt
          Time left:  epoch 0:00:20  total 0:03:29          
          Time left:  epoch 0:00:19  total 0:03:28          
          Time left:  epoch 0:00:19  total 0:03:27          
          Time left:  epoch 0:00:18  total 0:03:26          
          Time left:  epoch 0:00:17  total 0:03:25          
          Time left:  epoch 0:00:17  total 0:03:25          
          Time left:  epoch 0:00:16  total 0:03:24          
          Time left:  epoch 0:00:15  total 0:03:23          
          Time left:  epoch 0:00:14  total 0:03:22          
          Time left:  epoch 0:00:14  total 0:03:22          
          Time left:  epoch 0:00:13  total 0:03:21          
          Time left:  epoch 0:00:12  total 0:03:21          
          Time left:  epoch 0:00:12  total 0:03:20          
          Time left:  epoch 0:00:11  total 0:03:19          
          Time left:  epoch 0:00:10  total 0:03:18          
          Time left:  epoch 0:00:10  total 0:03:17          
          Time left:  epoch 0:00:09  total 0:03:16          
          Time left:  epoch 0:00:08  total 0:03:16          
          Time left:  epoch 0:00:07  total 0:03:15          
          Time left:  epoch 0:00:07  total 0:03:14          
          Time left:  epoch 0:00:06  total 0:03:13          
          Time left:  epoch 0:00:05  total 0:03:12          
          Time left:  epoch 0:00:05  total 0:03:11          
          Time left:  epoch 0:00:04  total 0:03:11          
          Time left:  epoch 0:00:04  total 0:03:11          
          Time left:  epoch 0:00:03  total 0:03:10          
          Time left:  epoch 0:00:02  total 0:03:09          
          Time left:  epoch 0:00:02  total 0:03:08          
          Time left:  epoch 0:00:01  total 0:03:07          
          Time left:  epoch 0:00:00  total 0:03:06          
          Time left:  epoch 0:00:00  total 0:03:06          

actor                                                       
  clip fraction                                        0.088
  entropy                                              0.425
  iterations                                            11.4
  kl                                                 0.00978
  loss                                              -0.00738
  std                                                  0.371
  stop                                                0.0877
critic                                                      
  iterations                                              80
  loss                                                0.0282
  v                                                     98.2
test                                                        
  action                                                    
    max                                                 2.28
    mean                                             0.00865
    min                                                -2.81
    size                                               5,000
    std                                                0.959
  episode length                                            
    max                                                1,000
    mean                                               1e+03
    min                                                1,000
    size                                                   5
    std                                                    0
  episode score                                             
    max                                                  976
    mean                                                 965
    min                                                  954
    size                                                   5
    std                                                 8.46
train                                                       
  action                                                    
    max                                                  2.5
    mean                                             0.00407
    min                                                 -2.7
    size                                              20,000
    std                                                 0.96
  episode length                                            
    max                                                1,000
    mean                                               1e+03
    min                                                1,000
    size                                                  20
    std                                                    0
  episode score                                             
    max                                                  978
    mean                                                 966
    min                                                  952
    size                                                  20
    std                                                  7.6
  episodes                                               320
  epoch seconds                                         18.7
  epoch steps                                          2e+04
  epochs                                                  16
  seconds                                                334
  steps                                              3.2e+05
  steps per second                                  1.07e+03
  worker steps                                       3.2e+05

          Time left:  epoch 0:00:20  total 0:03:07          
          Time left:  epoch 0:00:19  total 0:03:06          
          Time left:  epoch 0:00:19  total 0:03:05          
          Time left:  epoch 0:00:18  total 0:03:04          
          Time left:  epoch 0:00:17  total 0:03:03          
          Time left:  epoch 0:00:17  total 0:03:03          
          Time left:  epoch 0:00:16  total 0:03:03          
          Time left:  epoch 0:00:15  total 0:03:02          
          Time left:  epoch 0:00:15  total 0:03:01          
          Time left:  epoch 0:00:14  total 0:03:00          
          Time left:  epoch 0:00:13  total 0:02:59          
          Time left:  epoch 0:00:13  total 0:02:58          
          Time left:  epoch 0:00:12  total 0:02:58          
          Time left:  epoch 0:00:12  total 0:02:58          
          Time left:  epoch 0:00:11  total 0:02:57          
          Time left:  epoch 0:00:10  total 0:02:56          
          Time left:  epoch 0:00:09  total 0:02:55          
          Time left:  epoch 0:00:09  total 0:02:54          
          Time left:  epoch 0:00:08  total 0:02:53          
          Time left:  epoch 0:00:07  total 0:02:53          
          Time left:  epoch 0:00:07  total 0:02:52          
          Time left:  epoch 0:00:06  total 0:02:51          
          Time left:  epoch 0:00:05  total 0:02:50          
          Time left:  epoch 0:00:05  total 0:02:49          
          Time left:  epoch 0:00:04  total 0:02:49          
          Time left:  epoch 0:00:03  total 0:02:48          
          Time left:  epoch 0:00:03  total 0:02:48          
          Time left:  epoch 0:00:02  total 0:02:47          
          Time left:  epoch 0:00:01  total 0:02:46          
          Time left:  epoch 0:00:01  total 0:02:45          
          Time left:  epoch 0:00:00  total 0:02:44          
          Time left:  epoch 0:00:00  total 0:02:45          

actor                                                       
  clip fraction                                       0.0935
  entropy                                              0.408
  iterations                                            19.8
  kl                                                 0.00774
  loss                                               -0.0108
  std                                                  0.364
  stop                                                0.0404
critic                                                      
  iterations                                              80
  loss                                                0.0238
  v                                                     98.4
test                                                        
  action                                                    
    max                                                 2.38
    mean                                              0.0126
    min                                                -2.35
    size                                               5,000
    std                                                0.955
  episode length                                            
    max                                                1,000
    mean                                               1e+03
    min                                                1,000
    size                                                   5
    std                                                    0
  episode score                                             
    max                                                  974
    mean                                                 963
    min                                                  952
    size                                                   5
    std                                                 7.75
train                                                       
  action                                                    
    max                                                 2.44
    mean                                             0.00408
    min                                                -2.34
    size                                              20,000
    std                                                0.957
  episode length                                            
    max                                                1,000
    mean                                               1e+03
    min                                                1,000
    size                                                  20
    std                                                    0
  episode score                                             
    max                                                  985
    mean                                                 968
    min                                                  957
    size                                                  20
    std                                                 6.64
  episodes                                               340
  epoch seconds                                         19.3
  epoch steps                                          2e+04
  epochs                                                  17
  seconds                                                353
  steps                                              3.4e+05
  steps per second                                  1.03e+03
  worker steps                                       3.4e+05

          Time left:  epoch 0:00:20  total 0:02:45          
          Time left:  epoch 0:00:19  total 0:02:44          
          Time left:  epoch 0:00:19  total 0:02:44          
          Time left:  epoch 0:00:18  total 0:02:43          
          Time left:  epoch 0:00:17  total 0:02:42          
          Time left:  epoch 0:00:16  total 0:02:41          
          Time left:  epoch 0:00:16  total 0:02:41          
          Time left:  epoch 0:00:15  total 0:02:41          
          Time left:  epoch 0:00:15  total 0:02:40          
          Time left:  epoch 0:00:14  total 0:02:39          
          Time left:  epoch 0:00:13  total 0:02:38          
          Time left:  epoch 0:00:13  total 0:02:37          
          Time left:  epoch 0:00:12  total 0:02:36          
          Time left:  epoch 0:00:11  total 0:02:36          
          Time left:  epoch 0:00:11  total 0:02:35          
          Time left:  epoch 0:00:10  total 0:02:34          
          Time left:  epoch 0:00:09  total 0:02:34          
          Time left:  epoch 0:00:08  total 0:02:33          
          Time left:  epoch 0:00:08  total 0:02:32          
          Time left:  epoch 0:00:07  total 0:02:32          
          Time left:  epoch 0:00:06  total 0:02:31          
          Time left:  epoch 0:00:06  total 0:02:30          
          Time left:  epoch 0:00:05  total 0:02:29          
          Time left:  epoch 0:00:04  total 0:02:28          
          Time left:  epoch 0:00:04  total 0:02:27          
          Time left:  epoch 0:00:03  total 0:02:27          
          Time left:  epoch 0:00:03  total 0:02:27          
          Time left:  epoch 0:00:02  total 0:02:26          
          Time left:  epoch 0:00:01  total 0:02:25          
          Time left:  epoch 0:00:01  total 0:02:24          
          Time left:  epoch 0:00:00  total 0:02:23          
          Time left:  epoch 0:00:00  total 0:02:23          

actor                                                       
  clip fraction                                       0.0847
  entropy                                              0.392
  iterations                                            14.2
  kl                                                  0.0102
  loss                                               -0.0103
  std                                                  0.359
  stop                                                0.0702
critic                                                      
  iterations                                              80
  loss                                                0.0286
  v                                                     98.4
test                                                        
  action                                                    
    max                                                 2.44
    mean                                             0.00693
    min                                                -2.19
    size                                               5,000
    std                                                0.953
  episode length                                            
    max                                                1,000
    mean                                               1e+03
    min                                                1,000
    size                                                   5
    std                                                    0
  episode score                                             
    max                                                  981
    mean                                                 974
    min                                                  964
    size                                                   5
    std                                                 5.77
train                                                       
  action                                                    
    max                                                 2.39
    mean                                             0.00459
    min                                                -2.56
    size                                              20,000
    std                                                0.956
  episode length                                            
    max                                                1,000
    mean                                               1e+03
    min                                                1,000
    size                                                  20
    std                                                    0
  episode score                                             
    max                                                  983
    mean                                                 969
    min                                                  958
    size                                                  20
    std                                                 8.25
  episodes                                               360
  epoch seconds                                         17.6
  epoch steps                                          2e+04
  epochs                                                  18
  seconds                                                371
  steps                                              3.6e+05
  steps per second                                  1.14e+03
  worker steps                                       3.6e+05

          Time left:  epoch 0:00:20  total 0:02:23          
          Time left:  epoch 0:00:19  total 0:02:23          
          Time left:  epoch 0:00:18  total 0:02:23          
          Time left:  epoch 0:00:18  total 0:02:22          
          Time left:  epoch 0:00:17  total 0:02:21          
          Time left:  epoch 0:00:16  total 0:02:20          
          Time left:  epoch 0:00:16  total 0:02:19          
          Time left:  epoch 0:00:15  total 0:02:19          
          Time left:  epoch 0:00:14  total 0:02:18          
          Time left:  epoch 0:00:14  total 0:02:17          
          Time left:  epoch 0:00:13  total 0:02:17          
          Time left:  epoch 0:00:12  total 0:02:16          
          Time left:  epoch 0:00:12  total 0:02:15          
          Time left:  epoch 0:00:11  total 0:02:15          
          Time left:  epoch 0:00:10  total 0:02:14          
          Time left:  epoch 0:00:09  total 0:02:13          
          Time left:  epoch 0:00:09  total 0:02:13          
          Time left:  epoch 0:00:08  total 0:02:12          
          Time left:  epoch 0:00:07  total 0:02:11          
          Time left:  epoch 0:00:07  total 0:02:11          
          Time left:  epoch 0:00:06  total 0:02:11          
          Time left:  epoch 0:00:06  total 0:02:10          
          Time left:  epoch 0:00:05  total 0:02:09          
          Time left:  epoch 0:00:04  total 0:02:08          
          Time left:  epoch 0:00:04  total 0:02:08          
          Time left:  epoch 0:00:03  total 0:02:07          
          Time left:  epoch 0:00:02  total 0:02:06          
          Time left:  epoch 0:00:02  total 0:02:06          
          Time left:  epoch 0:00:01  total 0:02:05          
          Time left:  epoch 0:00:00  total 0:02:04          
          Time left:  epoch 0:00:00  total 0:02:03          

actor                                                       
  clip fraction                                       0.0913
  entropy                                              0.363
  iterations                                            58.2
  kl                                                 0.00949
  loss                                                -0.012
  std                                                  0.349
  stop                                                0.0103
critic                                                      
  iterations                                              80
  loss                                                0.0205
  v                                                     98.7
test                                                        
  action                                                    
    max                                                 2.39
    mean                                              0.0109
    min                                                -2.41
    size                                               5,000
    std                                                0.953
  episode length                                            
    max                                                1,000
    mean                                               1e+03
    min                                                1,000
    size                                                   5
    std                                                    0
  episode score                                             
    max                                                  984
    mean                                                 971
    min                                                  961
    size                                                   5
    std                                                 8.88
train                                                       
  action                                                    
    max                                                 2.43
    mean                                             0.00508
    min                                                 -2.5
    size                                              20,000
    std                                                0.953
  episode length                                            
    max                                                1,000
    mean                                               1e+03
    min                                                1,000
    size                                                  20
    std                                                    0
  episode score                                             
    max                                                  984
    mean                                                 971
    min                                                  957
    size                                                  20
    std                                                 6.36
  episodes                                               380
  epoch seconds                                           23
  epoch steps                                          2e+04
  epochs                                                  19
  seconds                                                394
  steps                                              3.8e+05
  steps per second                                       870
  worker steps                                       3.8e+05

          Time left:  epoch 0:00:20  total 0:02:04          
          Time left:  epoch 0:00:20  total 0:02:03          
          Time left:  epoch 0:00:19  total 0:02:03          
          Time left:  epoch 0:00:18  total 0:02:02          
          Time left:  epoch 0:00:17  total 0:02:01          
          Time left:  epoch 0:00:17  total 0:02:00          
          Time left:  epoch 0:00:16  total 0:02:00          
          Time left:  epoch 0:00:15  total 0:01:59          
          Time left:  epoch 0:00:15  total 0:01:58          
          Time left:  epoch 0:00:14  total 0:01:58          
          Time left:  epoch 0:00:14  total 0:01:57          
          Time left:  epoch 0:00:13  total 0:01:56          
          Time left:  epoch 0:00:12  total 0:01:56          
          Time left:  epoch 0:00:12  total 0:01:55          
          Time left:  epoch 0:00:11  total 0:01:54          
          Time left:  epoch 0:00:10  total 0:01:54          
          Time left:  epoch 0:00:09  total 0:01:53          
          Time left:  epoch 0:00:09  total 0:01:52          
          Time left:  epoch 0:00:08  total 0:01:51          
          Time left:  epoch 0:00:07  total 0:01:50          
          Time left:  epoch 0:00:07  total 0:01:50          
          Time left:  epoch 0:00:06  total 0:01:49          
          Time left:  epoch 0:00:05  total 0:01:48          
          Time left:  epoch 0:00:05  total 0:01:48          
          Time left:  epoch 0:00:04  total 0:01:47          
          Time left:  epoch 0:00:03  total 0:01:46          
          Time left:  epoch 0:00:03  total 0:01:45          
          Time left:  epoch 0:00:02  total 0:01:45          
          Time left:  epoch 0:00:01  total 0:01:44          
          Time left:  epoch 0:00:01  total 0:01:43          
          Time left:  epoch 0:00:00  total 0:01:42          
          Time left:  epoch 0:00:00  total 0:01:42          

actor                                                       
  clip fraction                                       0.0933
  entropy                                              0.331
  iterations                                             3.4
  kl                                                 0.00887
  loss                                               -0.0013
  std                                                  0.338
  stop                                                 0.294
critic                                                      
  iterations                                              80
  loss                                                0.0199
  v                                                     98.8
test                                                        
  action                                                    
    max                                                  2.1
    mean                                             0.00375
    min                                                -2.26
    size                                               5,000
    std                                                0.955
  episode length                                            
    max                                                1,000
    mean                                               1e+03
    min                                                1,000
    size                                                   5
    std                                                    0
  episode score                                             
    max                                                  984
    mean                                                 971
    min                                                  956
    size                                                   5
    std                                                 9.92
train                                                       
  action                                                    
    max                                                 2.35
    mean                                             0.00258
    min                                                -2.61
    size                                              20,000
    std                                                0.949
  episode length                                            
    max                                                1,000
    mean                                               1e+03
    min                                                1,000
    size                                                  20
    std                                                    0
  episode score                                             
    max                                                  984
    mean                                                 973
    min                                                  959
    size                                                  20
    std                                                 7.54
  episodes                                               400
  epoch seconds                                         18.4
  epoch steps                                          2e+04
  epochs                                                  20
  seconds                                                412
  steps                                                4e+05
  steps per second                                  1.09e+03
  worker steps                                         4e+05


Saving weights to data/local/experiments/tonic/swimmer-swim/mlp_256/checkpoints/step_400000.pt
          Time left:  epoch 0:00:20  total 0:01:42          
          Time left:  epoch 0:00:19  total 0:01:41          
          Time left:  epoch 0:00:19  total 0:01:41          
          Time left:  epoch 0:00:18  total 0:01:41          
          Time left:  epoch 0:00:17  total 0:01:40          
          Time left:  epoch 0:00:17  total 0:01:39          
          Time left:  epoch 0:00:16  total 0:01:38          
          Time left:  epoch 0:00:15  total 0:01:38          
          Time left:  epoch 0:00:15  total 0:01:37          
          Time left:  epoch 0:00:14  total 0:01:36          
          Time left:  epoch 0:00:13  total 0:01:36          
          Time left:  epoch 0:00:13  total 0:01:35          
          Time left:  epoch 0:00:12  total 0:01:34          
          Time left:  epoch 0:00:11  total 0:01:33          
          Time left:  epoch 0:00:10  total 0:01:32          
          Time left:  epoch 0:00:10  total 0:01:32          
          Time left:  epoch 0:00:09  total 0:01:31          
          Time left:  epoch 0:00:08  total 0:01:30          
          Time left:  epoch 0:00:08  total 0:01:30          
          Time left:  epoch 0:00:07  total 0:01:29          
          Time left:  epoch 0:00:06  total 0:01:28          
          Time left:  epoch 0:00:06  total 0:01:28          
          Time left:  epoch 0:00:05  total 0:01:27          
          Time left:  epoch 0:00:05  total 0:01:27          
          Time left:  epoch 0:00:04  total 0:01:26          
          Time left:  epoch 0:00:03  total 0:01:25          
          Time left:  epoch 0:00:03  total 0:01:24          
          Time left:  epoch 0:00:02  total 0:01:24          
          Time left:  epoch 0:00:01  total 0:01:23          
          Time left:  epoch 0:00:01  total 0:01:23          
          Time left:  epoch 0:00:00  total 0:01:22          
          Time left:  epoch 0:00:00  total 0:01:21          

actor                                                       
  clip fraction                                        0.108
  entropy                                              0.322
  iterations                                            21.2
  kl                                                 0.00864
  loss                                               -0.0117
  std                                                  0.334
  stop                                                0.0377
critic                                                      
  iterations                                              80
  loss                                                0.0213
  v                                                     98.7
test                                                        
  action                                                    
    max                                                 2.29
    mean                                             0.00332
    min                                                -2.25
    size                                               5,000
    std                                                0.951
  episode length                                            
    max                                                1,000
    mean                                               1e+03
    min                                                1,000
    size                                                   5
    std                                                    0
  episode score                                             
    max                                                  988
    mean                                                 973
    min                                                  966
    size                                                   5
    std                                                 7.93
train                                                       
  action                                                    
    max                                                 2.36
    mean                                             0.00432
    min                                                -2.28
    size                                              20,000
    std                                                0.951
  episode length                                            
    max                                                1,000
    mean                                               1e+03
    min                                                1,000
    size                                                  20
    std                                                    0
  episode score                                             
    max                                                  984
    mean                                                 970
    min                                                  945
    size                                                  20
    std                                                 9.62
  episodes                                               420
  epoch seconds                                         19.6
  epoch steps                                          2e+04
  epochs                                                  21
  seconds                                                432
  steps                                              4.2e+05
  steps per second                                  1.02e+03
  worker steps                                       4.2e+05

          Time left:  epoch 0:00:20  total 0:01:21          
          Time left:  epoch 0:00:19  total 0:01:21          
          Time left:  epoch 0:00:18  total 0:01:20          
          Time left:  epoch 0:00:18  total 0:01:19          
          Time left:  epoch 0:00:17  total 0:01:19          
          Time left:  epoch 0:00:16  total 0:01:18          
          Time left:  epoch 0:00:16  total 0:01:17          
          Time left:  epoch 0:00:15  total 0:01:16          
          Time left:  epoch 0:00:14  total 0:01:16          
          Time left:  epoch 0:00:14  total 0:01:15          
          Time left:  epoch 0:00:13  total 0:01:14          
          Time left:  epoch 0:00:12  total 0:01:14          
          Time left:  epoch 0:00:11  total 0:01:13          
          Time left:  epoch 0:00:11  total 0:01:12          
          Time left:  epoch 0:00:10  total 0:01:11          
          Time left:  epoch 0:00:10  total 0:01:11          
          Time left:  epoch 0:00:09  total 0:01:11          
          Time left:  epoch 0:00:08  total 0:01:10          
          Time left:  epoch 0:00:08  total 0:01:09          
          Time left:  epoch 0:00:07  total 0:01:09          
          Time left:  epoch 0:00:06  total 0:01:08          
          Time left:  epoch 0:00:06  total 0:01:07          
          Time left:  epoch 0:00:05  total 0:01:07          
          Time left:  epoch 0:00:04  total 0:01:06          
          Time left:  epoch 0:00:04  total 0:01:05          
          Time left:  epoch 0:00:03  total 0:01:04          
          Time left:  epoch 0:00:02  total 0:01:04          
          Time left:  epoch 0:00:02  total 0:01:03          
          Time left:  epoch 0:00:01  total 0:01:02          
          Time left:  epoch 0:00:00  total 0:01:02          
          Time left:  epoch 0:00:00  total 0:01:01          

actor                                                       
  clip fraction                                        0.104
  entropy                                              0.306
  iterations                                            22.4
  kl                                                  0.0101
  loss                                               -0.0141
  std                                                  0.329
  stop                                                0.0357
critic                                                      
  iterations                                              80
  loss                                                0.0126
  v                                                       99
test                                                        
  action                                                    
    max                                                 2.44
    mean                                             0.00292
    min                                                -2.23
    size                                               5,000
    std                                                0.943
  episode length                                            
    max                                                1,000
    mean                                               1e+03
    min                                                1,000
    size                                                   5
    std                                                    0
  episode score                                             
    max                                                  980
    mean                                                 976
    min                                                  966
    size                                                   5
    std                                                 4.92
train                                                       
  action                                                    
    max                                                 2.37
    mean                                             0.00601
    min                                                -2.49
    size                                              20,000
    std                                                0.948
  episode length                                            
    max                                                1,000
    mean                                               1e+03
    min                                                1,000
    size                                                  20
    std                                                    0
  episode score                                             
    max                                                  984
    mean                                                 976
    min                                                  966
    size                                                  20
    std                                                 5.48
  episodes                                               440
  epoch seconds                                         19.7
  epoch steps                                          2e+04
  epochs                                                  22
  seconds                                                452
  steps                                              4.4e+05
  steps per second                                  1.01e+03
  worker steps                                       4.4e+05

          Time left:  epoch 0:00:20  total 0:01:01          
          Time left:  epoch 0:00:19  total 0:01:00          
          Time left:  epoch 0:00:18  total 0:00:59          
          Time left:  epoch 0:00:18  total 0:00:59          
          Time left:  epoch 0:00:17  total 0:00:58          
          Time left:  epoch 0:00:16  total 0:00:57          
          Time left:  epoch 0:00:16  total 0:00:57          
          Time left:  epoch 0:00:15  total 0:00:56          
          Time left:  epoch 0:00:14  total 0:00:55          
          Time left:  epoch 0:00:13  total 0:00:54          
          Time left:  epoch 0:00:13  total 0:00:54          
          Time left:  epoch 0:00:12  total 0:00:53          
          Time left:  epoch 0:00:11  total 0:00:52          
          Time left:  epoch 0:00:11  total 0:00:52          
          Time left:  epoch 0:00:10  total 0:00:51          
          Time left:  epoch 0:00:09  total 0:00:50          
          Time left:  epoch 0:00:09  total 0:00:50          
          Time left:  epoch 0:00:08  total 0:00:49          
          Time left:  epoch 0:00:07  total 0:00:48          
          Time left:  epoch 0:00:07  total 0:00:47          
          Time left:  epoch 0:00:06  total 0:00:47          
          Time left:  epoch 0:00:05  total 0:00:46          
          Time left:  epoch 0:00:05  total 0:00:45          
          Time left:  epoch 0:00:04  total 0:00:45          
          Time left:  epoch 0:00:03  total 0:00:44          
          Time left:  epoch 0:00:03  total 0:00:43          
          Time left:  epoch 0:00:02  total 0:00:43          
          Time left:  epoch 0:00:01  total 0:00:42          
          Time left:  epoch 0:00:01  total 0:00:42          
          Time left:  epoch 0:00:00  total 0:00:41          
          Time left:  epoch 0:00:00  total 0:00:40          

actor                                                       
  clip fraction                                       0.0906
  entropy                                              0.286
  iterations                                            19.4
  kl                                                  0.0104
  loss                                              -0.00813
  std                                                  0.323
  stop                                                0.0515
critic                                                      
  iterations                                              80
  loss                                                0.0151
  v                                                       99
test                                                        
  action                                                    
    max                                                 2.28
    mean                                             0.00589
    min                                                -2.16
    size                                               5,000
    std                                                0.938
  episode length                                            
    max                                                1,000
    mean                                               1e+03
    min                                                1,000
    size                                                   5
    std                                                    0
  episode score                                             
    max                                                  980
    mean                                                 968
    min                                                  953
    size                                                   5
    std                                                 11.4
train                                                       
  action                                                    
    max                                                 2.33
    mean                                             0.00231
    min                                                -2.39
    size                                              20,000
    std                                                0.945
  episode length                                            
    max                                                1,000
    mean                                               1e+03
    min                                                1,000
    size                                                  20
    std                                                    0
  episode score                                             
    max                                                  986
    mean                                                 972
    min                                                  954
    size                                                  20
    std                                                 8.29
  episodes                                               460
  epoch seconds                                         19.7
  epoch steps                                          2e+04
  epochs                                                  23
  seconds                                                472
  steps                                              4.6e+05
  steps per second                                  1.01e+03
  worker steps                                       4.6e+05

          Time left:  epoch 0:00:20  total 0:00:40          
          Time left:  epoch 0:00:19  total 0:00:39          
          Time left:  epoch 0:00:18  total 0:00:39          
          Time left:  epoch 0:00:18  total 0:00:38          
          Time left:  epoch 0:00:17  total 0:00:38          
          Time left:  epoch 0:00:17  total 0:00:37          
          Time left:  epoch 0:00:16  total 0:00:36          
          Time left:  epoch 0:00:15  total 0:00:36          
          Time left:  epoch 0:00:14  total 0:00:35          
          Time left:  epoch 0:00:14  total 0:00:34          
          Time left:  epoch 0:00:13  total 0:00:34          
          Time left:  epoch 0:00:12  total 0:00:33          
          Time left:  epoch 0:00:12  total 0:00:32          
          Time left:  epoch 0:00:11  total 0:00:31          
          Time left:  epoch 0:00:10  total 0:00:31          
          Time left:  epoch 0:00:10  total 0:00:30          
          Time left:  epoch 0:00:09  total 0:00:29          
          Time left:  epoch 0:00:09  total 0:00:29          
          Time left:  epoch 0:00:08  total 0:00:28          
          Time left:  epoch 0:00:07  total 0:00:28          
          Time left:  epoch 0:00:07  total 0:00:27          
          Time left:  epoch 0:00:06  total 0:00:26          
          Time left:  epoch 0:00:05  total 0:00:26          
          Time left:  epoch 0:00:05  total 0:00:25          
          Time left:  epoch 0:00:04  total 0:00:24          
          Time left:  epoch 0:00:03  total 0:00:24          
          Time left:  epoch 0:00:03  total 0:00:23          
          Time left:  epoch 0:00:02  total 0:00:22          
          Time left:  epoch 0:00:01  total 0:00:22          
          Time left:  epoch 0:00:01  total 0:00:21          
          Time left:  epoch 0:00:00  total 0:00:20          
          Time left:  epoch 0:00:00  total 0:00:20          

actor                                                       
  clip fraction                                        0.109
  entropy                                              0.263
  iterations                                              19
  kl                                                  0.0101
  loss                                               -0.0109
  std                                                  0.315
  stop                                                0.0421
critic                                                      
  iterations                                              80
  loss                                                0.0128
  v                                                       99
test                                                        
  action                                                    
    max                                                 2.45
    mean                                              0.0015
    min                                                -2.12
    size                                               5,000
    std                                                0.938
  episode length                                            
    max                                                1,000
    mean                                               1e+03
    min                                                1,000
    size                                                   5
    std                                                    0
  episode score                                             
    max                                                  979
    mean                                                 972
    min                                                  961
    size                                                   5
    std                                                 6.86
train                                                       
  action                                                    
    max                                                 2.21
    mean                                             0.00243
    min                                                -2.26
    size                                              20,000
    std                                                0.942
  episode length                                            
    max                                                1,000
    mean                                               1e+03
    min                                                1,000
    size                                                  20
    std                                                    0
  episode score                                             
    max                                                  985
    mean                                                 973
    min                                                  958
    size                                                  20
    std                                                 8.35
  episodes                                               480
  epoch seconds                                         19.8
  epoch steps                                          2e+04
  epochs                                                  24
  seconds                                                491
  steps                                              4.8e+05
  steps per second                                  1.01e+03
  worker steps                                       4.8e+05

          Time left:  epoch 0:00:20  total 0:00:20          
          Time left:  epoch 0:00:19  total 0:00:19          
          Time left:  epoch 0:00:18  total 0:00:18          
          Time left:  epoch 0:00:18  total 0:00:18          
          Time left:  epoch 0:00:17  total 0:00:17          
          Time left:  epoch 0:00:16  total 0:00:16          
          Time left:  epoch 0:00:16  total 0:00:16          
          Time left:  epoch 0:00:15  total 0:00:15          
          Time left:  epoch 0:00:14  total 0:00:14          
          Time left:  epoch 0:00:13  total 0:00:13          
          Time left:  epoch 0:00:13  total 0:00:13          
          Time left:  epoch 0:00:12  total 0:00:12          
          Time left:  epoch 0:00:12  total 0:00:12          
          Time left:  epoch 0:00:11  total 0:00:11          
          Time left:  epoch 0:00:10  total 0:00:10          
          Time left:  epoch 0:00:10  total 0:00:10          
          Time left:  epoch 0:00:09  total 0:00:09          
          Time left:  epoch 0:00:08  total 0:00:08          
          Time left:  epoch 0:00:08  total 0:00:08          
          Time left:  epoch 0:00:07  total 0:00:07          
          Time left:  epoch 0:00:06  total 0:00:06          
          Time left:  epoch 0:00:06  total 0:00:06          
          Time left:  epoch 0:00:05  total 0:00:05          
          Time left:  epoch 0:00:04  total 0:00:04          
          Time left:  epoch 0:00:04  total 0:00:04          
          Time left:  epoch 0:00:03  total 0:00:03          
          Time left:  epoch 0:00:02  total 0:00:02          
          Time left:  epoch 0:00:02  total 0:00:02          
          Time left:  epoch 0:00:01  total 0:00:01          
          Time left:  epoch 0:00:00  total 0:00:00          
          Time left:  epoch 0:00:00  total 0:00:00          
             Time left:  epoch 0:00:00  total               

actor                                                       
  clip fraction                                       0.0934
  entropy                                              0.244
  iterations                                            16.6
  kl                                                 0.00931
  loss                                              -0.00989
  std                                                  0.309
  stop                                                0.0602
critic                                                      
  iterations                                              80
  loss                                               0.00959
  v                                                     99.1
test                                                        
  action                                                    
    max                                                  2.3
    mean                                           -6.33e-06
    min                                                 -2.2
    size                                               5,000
    std                                                 0.94
  episode length                                            
    max                                                1,000
    mean                                               1e+03
    min                                                1,000
    size                                                   5
    std                                                    0
  episode score                                             
    max                                                  979
    mean                                                 974
    min                                                  965
    size                                                   5
    std                                                 4.99
train                                                       
  action                                                    
    max                                                  2.2
    mean                                             0.00124
    min                                                -2.27
    size                                              20,000
    std                                                 0.94
  episode length                                            
    max                                                1,000
    mean                                               1e+03
    min                                                1,000
    size                                                  20
    std                                                    0
  episode score                                             
    max                                                  988
    mean                                                 973
    min                                                  959
    size                                                  20
    std                                                 6.67
  episodes                                               500
  epoch seconds                                         19.4
  epoch steps                                          2e+04
  epochs                                                  25
  seconds                                                511
  steps                                                5e+05
  steps per second                                  1.03e+03
  worker steps                                         5e+05


Saving weights to data/local/experiments/tonic/swimmer-swim/mlp_256/checkpoints/step_500000.pt

Try playing with the parameters of the trainer and the MLP model and see how it affects the performance.

  • How do the actor and the critic model size affect the performance.

  • Consider increasing the number of steps in trainer to train the model for longer.

  • Explore Tonic library to see what algorithms we can use to train our agents. (D4PG is usually faster than PPO)

# add your code

Section 2.3 Function to run any model on the environment and generate video

One of the most fun things about these environments is their visualization. We don’t want to just look at the reward to know how good our model is we want to see how well the agent swims. This is particularly important to avoid “reward hacking,” where an agent learns to exploit the reward system in ways that are unintended and potentially detrimental to the desired outcomes. Moreover visualizing the agent also help us understand where the model is going wrong.

Here we have defined a function that will generate the videos of the agent using the input model. The function requires path to the checkpoint folder and the environment you wanna run the trained model on.

def play_model(path, checkpoint='last',environment='default',seed=None, header=None):

  """
    Plays a model within an environment and renders the gameplay to a video.

    Parameters:
    - path (str): Path to the directory containing the model and checkpoints.
    - checkpoint (str): Specifies which checkpoint to use ('last', 'first', or a specific ID). 'none' indicates no checkpoint.
    - environment (str): The environment to use. 'default' uses the environment specified in the configuration file.
    - seed (int): Optional seed for reproducibility.
    - header (str): Optional Python code to execute before initializing the model, such as importing libraries.
    """

  if checkpoint == 'none':
    # Use no checkpoint, the agent is freshly created.
    checkpoint_path = None
    tonic.logger.log('Not loading any weights')
  else:
    checkpoint_path = os.path.join(path, 'checkpoints')
    if not os.path.isdir(checkpoint_path):
      tonic.logger.error(f'{checkpoint_path} is not a directory')
      checkpoint_path = None

    # List all the checkpoints.
    checkpoint_ids = []
    for file in os.listdir(checkpoint_path):
      if file[:5] == 'step_':
        checkpoint_id = file.split('.')[0]
        checkpoint_ids.append(int(checkpoint_id[5:]))

    if checkpoint_ids:
      if checkpoint == 'last':
        # Use the last checkpoint.
        checkpoint_id = max(checkpoint_ids)
        checkpoint_path = os.path.join(checkpoint_path, f'step_{checkpoint_id}')
      elif checkpoint == 'first':
        # Use the first checkpoint.
        checkpoint_id = min(checkpoint_ids)
        checkpoint_path = os.path.join(checkpoint_path, f'step_{checkpoint_id}')
      else:
        # Use the specified checkpoint.
        checkpoint_id = int(checkpoint)
        if checkpoint_id in checkpoint_ids:
          checkpoint_path = os.path.join(checkpoint_path, f'step_{checkpoint_id}')
        else:
          tonic.logger.error(f'Checkpoint {checkpoint_id} not found in {checkpoint_path}')
          checkpoint_path = None
    else:
      tonic.logger.error(f'No checkpoint found in {checkpoint_path}')
      checkpoint_path = None

  # Load the experiment configuration.
  arguments_path = os.path.join(path, 'config.yaml')
  with open(arguments_path, 'r') as config_file:
    config = yaml.load(config_file, Loader=yaml.FullLoader)
  config = argparse.Namespace(**config)

  # Run the header first, e.g. to load an ML framework.
  try:
    if config.header:
      exec(config.header)
    if header:
      exec(header)
  except:
    pass

  # Build the agent.
  agent = eval(config.agent)

  # Build the environment.
  if environment == 'default':
    environment  = tonic.environments.distribute(lambda: eval(config.environment))
  else:
    environment  = tonic.environments.distribute(lambda: eval(environment))
  if seed is not None:
    environment.seed(seed)

  # Initialize the agent.
  agent.initialize(
    observation_space=environment.observation_space,
    action_space=environment.action_space,
    seed=seed,
  )

  # Load the weights of the agent form a checkpoint.
  if checkpoint_path:
    agent.load(checkpoint_path)

  steps = 0
  test_observations = environment.start()
  frames = [environment.render('rgb_array',camera_id=0, width=640, height=480)[0]]
  score, length = 0, 0

  while True:
      # Select an action.
      actions = agent.test_step(test_observations, steps)
      assert not np.isnan(actions.sum())

      # Take a step in the environment.
      test_observations, infos = environment.step(actions)
      frames.append(environment.render('rgb_array',camera_id=0, width=640, height=480)[0])
      agent.test_update(**infos, steps=steps)

      score += infos['rewards'][0]
      length += 1

      if infos['resets'][0]:
          break
  video_path = os.path.join(path, 'video.mp4')
  print('Reward for the run: ', score)
  return display_video(frames,video_path)

Let’s visualize the agent with a pretrained MLP model. Once you have your pretrained model, you can replace the experiment path to visualize the agent with your model.

# play_model('data/local/experiments/tonic/swimmer-swim/mlp_256')
play_model('data/local/experiments/tonic/swimmer-swim/pretrained_mlp_ppo')
/opt/hostedtoolcache/Python/3.10.20/x64/lib/python3.10/site-packages/gym/spaces/box.py:127: UserWarning: WARN: Box bound precision lowered by casting to float32
  logger.warn(f"Box bound precision lowered by casting to {self.dtype}")

Loading weights from data/local/experiments/tonic/swimmer-swim/pretrained_mlp_ppo/checkpoints/step_10000000.pt
Reward for the run:  979.3922030646354
/tmp/ipykernel_8472/4073379236.py:65: MatplotlibDeprecationWarning: Auto-close()ing of figures upon backend switching is deprecated since 3.8 and will be removed two minor releases later.  To suppress this warning, explicitly call plt.close('all') first.
  matplotlib.use(orig_backend)  # Switch back to the original backend.
Loading...

Try testing the model on a modification of the enviroment it was trained on.

  • Train on basic swim task and test on a environment with higher viscosity.

  • Can we train on the basic 6 link swimmer and test on a larger 12 link swimmer?

  • Train the model for a bit on the modified environment and see how quickly the model can adapt to the new environment.

# add your code

Section 3: NCAP

Now that we are familiar with how to train standard models on the swimmer agent let’s take a look at NCAP a model that was inspired from the C. elegans motor circuit. Our hope with using such models is that they would already have really good priors which should lead to much better transfer, faster learning curves and possibly really good innate performance (zero shot performance).

Check out NCAP paper to learn more about the model.

Paper Illustration

Source
#@title Paper Illustration

from IPython.display import Image, display
import os
from pathlib import Path

url = "https://github.com/neuromatch/NeuroAI_Course/blob/main/projects/project-notebooks/static/NCAPPaper.png?raw=true"

display(Image(url=url))
Loading...

3.1 NCAP classes

Now we are going to define SwimmerModule (NCAP model) and SwimmerActor (wrapper around NCAP model to make it compatible with tonic) classes.

Section 3.1.1 Defining the constraints

# ==================================================================================================
# Weight constraints.


def excitatory(w, upper=None):
    return w.clamp(min=0, max=upper)


def inhibitory(w, lower=None):
    return w.clamp(min=lower, max=0)


def unsigned(w, lower=None, upper=None):
    return w if lower is None and upper is None else w.clamp(min=lower, max=upper)


# ==================================================================================================
# Activation constraints.


def graded(x):
    return x.clamp(min=0, max=1)


# ==================================================================================================
# Weight initialization.


def excitatory_uniform(shape=(1,), lower=0., upper=1.):
    assert lower >= 0
    return nn.init.uniform_(nn.Parameter(torch.empty(shape)), a=lower, b=upper)


def inhibitory_uniform(shape=(1,), lower=-1., upper=0.):
    assert upper <= 0
    return nn.init.uniform_(nn.Parameter(torch.empty(shape)), a=lower, b=upper)


def unsigned_uniform(shape=(1,), lower=-1., upper=1.):
    return nn.init.uniform_(nn.Parameter(torch.empty(shape)), a=lower, b=upper)


def excitatory_constant(shape=(1,), value=1.):
    return nn.Parameter(torch.full(shape, value))


def inhibitory_constant(shape=(1,), value=-1.):
    return nn.Parameter(torch.full(shape, value))


def unsigned_constant(shape=(1,), lower=-1., upper=1., p=0.5):
    with torch.no_grad():
        weight = torch.empty(shape).uniform_(0, 1)
        mask = weight < p
        weight[mask] = upper
        weight[~mask] = lower
        return nn.Parameter(weight)

Can you think of more kinds of weight initializations and constraints that might be useful for the swimmer agent?

Section 3.1.2: Defining the SwimmerModule

The SwimmerModule class represents the neural network module inspired by the C. elegans neural circuitry, designed for controlling a robotic swimmer with specific architectural priors, such as proprioception and oscillatory movement patterns.

class SwimmerModule(nn.Module):
    """C.-elegans-inspired neural circuit architectural prior."""

    def __init__(
            self,
            n_joints: int,
            n_turn_joints: int = 1,
            oscillator_period: int = 60,
            use_weight_sharing: bool = True,
            use_weight_constraints: bool = True,
            use_weight_constant_init: bool = True,
            include_proprioception: bool = True,
            include_head_oscillators: bool = True,
            include_speed_control: bool = False,
            include_turn_control: bool = False,
    ):
        super().__init__()
        self.n_joints = n_joints
        self.n_turn_joints = n_turn_joints
        self.oscillator_period = oscillator_period
        self.include_proprioception = include_proprioception
        self.include_head_oscillators = include_head_oscillators
        self.include_speed_control = include_speed_control
        self.include_turn_control = include_turn_control

        # Log activity
        self.connections_log = []

        # Timestep counter (for oscillations).
        self.timestep = 0

        # Weight sharing switch function.
        self.ws = lambda nonshared, shared: shared if use_weight_sharing else nonshared

        # Weight constraint and init functions.
        if use_weight_constraints:
            self.exc = excitatory
            self.inh = inhibitory
            if use_weight_constant_init:
                exc_param = excitatory_constant
                inh_param = inhibitory_constant
            else:
                exc_param = excitatory_uniform
                inh_param = inhibitory_uniform
        else:
            self.exc = unsigned
            self.inh = unsigned
            if use_weight_constant_init:
                exc_param = inh_param = unsigned_constant
            else:
                exc_param = inh_param = unsigned_uniform

        # Learnable parameters.
        self.params = nn.ParameterDict()
        if use_weight_sharing:
            if self.include_proprioception:
                self.params['bneuron_prop'] = exc_param()
            if self.include_speed_control:
                self.params['bneuron_speed'] = inh_param()
            if self.include_turn_control:
                self.params['bneuron_turn'] = exc_param()
            if self.include_head_oscillators:
                self.params['bneuron_osc'] = exc_param()
            self.params['muscle_ipsi'] = exc_param()
            self.params['muscle_contra'] = inh_param()
        else:
            for i in range(self.n_joints):
                if self.include_proprioception and i > 0:
                    self.params[f'bneuron_d_prop_{i}'] = exc_param()
                    self.params[f'bneuron_v_prop_{i}'] = exc_param()

                if self.include_speed_control:
                    self.params[f'bneuron_d_speed_{i}'] = inh_param()
                    self.params[f'bneuron_v_speed_{i}'] = inh_param()

                if self.include_turn_control and i < self.n_turn_joints:
                    self.params[f'bneuron_d_turn_{i}'] = exc_param()
                    self.params[f'bneuron_v_turn_{i}'] = exc_param()

                if self.include_head_oscillators and i == 0:
                    self.params[f'bneuron_d_osc_{i}'] = exc_param()
                    self.params[f'bneuron_v_osc_{i}'] = exc_param()

                self.params[f'muscle_d_d_{i}'] = exc_param()
                self.params[f'muscle_d_v_{i}'] = inh_param()
                self.params[f'muscle_v_v_{i}'] = exc_param()
                self.params[f'muscle_v_d_{i}'] = inh_param()

    def reset(self):
        self.timestep = 0

    def log_activity(self, activity_type, neuron):
        """Logs an active connection between neurons."""
        self.connections_log.append((self.timestep, activity_type, neuron))

    def forward(
            self,
            joint_pos,
            right_control=None,
            left_control=None,
            speed_control=None,
            timesteps=None,
            log_activity=True,
            log_file='log.txt'
    ):
        """Forward pass.

    Args:
      joint_pos (torch.Tensor): Joint positions in [-1, 1], shape (..., n_joints).
      right_control (torch.Tensor): Right turn control in [0, 1], shape (..., 1).
      left_control (torch.Tensor): Left turn control in [0, 1], shape (..., 1).
      speed_control (torch.Tensor): Speed control in [0, 1], 0 stopped, 1 fastest, shape (..., 1).
      timesteps (torch.Tensor): Timesteps in [0, max_env_steps], shape (..., 1).

    Returns:
      (torch.Tensor): Joint torques in [-1, 1], shape (..., n_joints).
    """

        exc = self.exc
        inh = self.inh
        ws = self.ws

        # Separate into dorsal and ventral sensor values in [0, 1], shape (..., n_joints).
        joint_pos_d = joint_pos.clamp(min=0, max=1)
        joint_pos_v = joint_pos.clamp(min=-1, max=0).neg()

        # Convert speed signal from acceleration into brake.
        if self.include_speed_control:
            assert speed_control is not None
            speed_control = 1 - speed_control.clamp(min=0, max=1)

        joint_torques = []  # [shape (..., 1)]
        for i in range(self.n_joints):
            bneuron_d = bneuron_v = torch.zeros_like(joint_pos[..., 0, None])  # shape (..., 1)

            # B-neurons recieve proprioceptive input from previous joint to propagate waves down the body.
            if self.include_proprioception and i > 0:
                bneuron_d = bneuron_d + joint_pos_d[
                    ..., i - 1, None] * exc(self.params[ws(f'bneuron_d_prop_{i}', 'bneuron_prop')])
                bneuron_v = bneuron_v + joint_pos_v[
                    ..., i - 1, None] * exc(self.params[ws(f'bneuron_v_prop_{i}', 'bneuron_prop')])
                self.log_activity('exc', f'bneuron_d_prop_{i}')
                self.log_activity('exc', f'bneuron_v_prop_{i}')

            # Speed control unit modulates all B-neurons.
            if self.include_speed_control:
                bneuron_d = bneuron_d + speed_control * inh(
                    self.params[ws(f'bneuron_d_speed_{i}', 'bneuron_speed')]
                )
                bneuron_v = bneuron_v + speed_control * inh(
                    self.params[ws(f'bneuron_v_speed_{i}', 'bneuron_speed')]
                )
                self.log_activity('inh', f'bneuron_d_speed_{i}')
                self.log_activity('inh', f'bneuron_v_speed_{i}')

            # Turn control units modulate head B-neurons.
            if self.include_turn_control and i < self.n_turn_joints:
                assert right_control is not None
                assert left_control is not None
                turn_control_d = right_control.clamp(min=0, max=1)  # shape (..., 1)
                turn_control_v = left_control.clamp(min=0, max=1)
                bneuron_d = bneuron_d + turn_control_d * exc(
                    self.params[ws(f'bneuron_d_turn_{i}', 'bneuron_turn')]
                )
                bneuron_v = bneuron_v + turn_control_v * exc(
                    self.params[ws(f'bneuron_v_turn_{i}', 'bneuron_turn')]
                )
                self.log_activity('exc', f'bneuron_d_turn_{i}')
                self.log_activity('exc', f'bneuron_v_turn_{i}')

            # Oscillator units modulate first B-neurons.
            if self.include_head_oscillators and i == 0:
                if timesteps is not None:
                    phase = timesteps.round().remainder(self.oscillator_period)
                    mask = phase < self.oscillator_period // 2
                    oscillator_d = torch.zeros_like(timesteps)  # shape (..., 1)
                    oscillator_v = torch.zeros_like(timesteps)  # shape (..., 1)
                    oscillator_d[mask] = 1.
                    oscillator_v[~mask] = 1.
                else:
                    phase = self.timestep % self.oscillator_period  # in [0, oscillator_period)
                    if phase < self.oscillator_period // 2:
                        oscillator_d, oscillator_v = 1.0, 0.0
                    else:
                        oscillator_d, oscillator_v = 0.0, 1.0
                bneuron_d = bneuron_d + oscillator_d * exc(
                    self.params[ws(f'bneuron_d_osc_{i}', 'bneuron_osc')]
                )
                bneuron_v = bneuron_v + oscillator_v * exc(
                    self.params[ws(f'bneuron_v_osc_{i}', 'bneuron_osc')]
                )

                self.log_activity('exc', f'bneuron_d_osc_{i}')
                self.log_activity('exc', f'bneuron_v_osc_{i}')

            # B-neuron activation.
            bneuron_d = graded(bneuron_d)
            bneuron_v = graded(bneuron_v)

            # Muscles receive excitatory ipsilateral and inhibitory contralateral input.
            muscle_d = graded(
                bneuron_d * exc(self.params[ws(f'muscle_d_d_{i}', 'muscle_ipsi')]) +
                bneuron_v * inh(self.params[ws(f'muscle_d_v_{i}', 'muscle_contra')])
            )
            muscle_v = graded(
                bneuron_v * exc(self.params[ws(f'muscle_v_v_{i}', 'muscle_ipsi')]) +
                bneuron_d * inh(self.params[ws(f'muscle_v_d_{i}', 'muscle_contra')])
            )

            # Joint torque from antagonistic contraction of dorsal and ventral muscles.
            joint_torque = muscle_d - muscle_v
            joint_torques.append(joint_torque)

        self.timestep += 1

        out = torch.cat(joint_torques, -1)  # shape (..., n_joints)
        return out

Section 3.1.3: Defining the SwimmerActor wrapper

The SwimmerActor class acts as a wrapper around the SwimmerModule, managing high-level control signals and observations coming from the environment and passing them to the SwimmerModule in a suitable format. This class is basically responsible for making the SwimmerModule compatible with the tonic library. If you wish to use any other library to try a algorithm not present in tonic you have to write a new wrapper to make SwimmerModule compatible with that library.

class SwimmerActor(nn.Module):
    def __init__(
            self,
            swimmer,
            controller=None,
            distribution=None,
            timestep_transform=(-1, 1, 0, 1000),
    ):
        super().__init__()
        self.swimmer = swimmer
        self.controller = controller
        self.distribution = distribution
        self.timestep_transform = timestep_transform

    def initialize(
            self,
            observation_space,
            action_space,
            observation_normalizer=None,
    ):
        self.action_size = action_space.shape[0]

    def forward(self, observations):
        joint_pos = observations[..., :self.action_size]
        timesteps = observations[..., -1, None]

        # Normalize joint positions by max joint angle (in radians).
        joint_limit = 2 * np.pi / (self.action_size + 1)  # In dm_control, calculated with n_bodies.
        joint_pos = torch.clamp(joint_pos / joint_limit, min=-1, max=1)

        # Convert normalized time signal into timestep.
        if self.timestep_transform:
            low_in, high_in, low_out, high_out = self.timestep_transform
            timesteps = (timesteps - low_in) / (high_in - low_in) * (high_out - low_out) + low_out

        # Generate high-level control signals.
        if self.controller:
            right, left, speed = self.controller(observations)
        else:
            right, left, speed = None, None, None

        # Generate low-level action signals.
        actions = self.swimmer(
            joint_pos,
            timesteps=timesteps,
            right_control=right,
            left_control=left,
            speed_control=speed,
        )

        # Pass through distribution for stochastic policy.
        if self.distribution:
            actions = self.distribution(actions)

        return actions

3.2: Train NCAP

We will now define functions akin to those for MLP we defined in Section 3.2, but tailored for the SwimmerActor model.

from tonic.torch import models, normalizers
import torch


def ppo_swimmer_model(
        n_joints=5,
        action_noise=0.1,
        critic_sizes=(64, 64),
        critic_activation=nn.Tanh,
        **swimmer_kwargs,
):
    return models.ActorCritic(
        actor=SwimmerActor(
            swimmer=SwimmerModule(n_joints=n_joints, **swimmer_kwargs),
            distribution=lambda x: torch.distributions.normal.Normal(x, action_noise),
        ),
        critic=models.Critic(
            encoder=models.ObservationEncoder(),
            torso=models.MLP(critic_sizes, critic_activation),
            head=models.ValueHead(),
        ),
        observation_normalizer=normalizers.MeanStd(),
    )


def d4pg_swimmer_model(
  n_joints=5,
  critic_sizes=(256, 256),
  critic_activation=nn.ReLU,
  **swimmer_kwargs,
):
  return models.ActorCriticWithTargets(
    actor=SwimmerActor(swimmer=SwimmerModule(n_joints=n_joints, **swimmer_kwargs),),
    critic=models.Critic(
      encoder=models.ObservationActionEncoder(),
      torso=models.MLP(critic_sizes, critic_activation),
      # These values are for the control suite with 0.99 discount.
      head=models.DistributionalValueHead(-150., 150., 51),
    ),
    observation_normalizer=normalizers.MeanStd(),
  )

train('import tonic.torch',
      # 'tonic.torch.agents.D4PG(model=d4pg_swimmer_model(n_joints=5,critic_sizes=(128,128)))',
      'tonic.torch.agents.PPO(model=ppo_swimmer_model(n_joints=5,critic_sizes=(256,256)))',
  'tonic.environments.ControlSuite("swimmer-swim",time_feature=True)',
  name = 'ncap_ppo',
  trainer = 'tonic.Trainer(steps=int(1e5),save_steps=int(5e4))')
/opt/hostedtoolcache/Python/3.10.20/x64/lib/python3.10/site-packages/gym/spaces/box.py:127: UserWarning: WARN: Box bound precision lowered by casting to float32
  logger.warn(f"Box bound precision lowered by casting to {self.dtype}")
Config file saved to data/local/experiments/tonic/swimmer-swim/ncap_ppo/config.yaml
          Time left:  epoch 0:00:49  total 0:04:05          
          Time left:  epoch 0:00:16  total 0:01:24          
          Time left:  epoch 0:00:16  total 0:01:23          
          Time left:  epoch 0:00:15  total 0:01:23          
          Time left:  epoch 0:00:15  total 0:01:22          
          Time left:  epoch 0:00:15  total 0:01:22          
          Time left:  epoch 0:00:15  total 0:01:22          
          Time left:  epoch 0:00:14  total 0:01:22          
          Time left:  epoch 0:00:14  total 0:01:21          
          Time left:  epoch 0:00:14  total 0:01:21          
          Time left:  epoch 0:00:14  total 0:01:21          
          Time left:  epoch 0:00:13  total 0:01:21          
          Time left:  epoch 0:00:13  total 0:01:20          
          Time left:  epoch 0:00:18  total 0:01:52          
          Time left:  epoch 0:00:17  total 0:01:49          
          Time left:  epoch 0:00:16  total 0:01:47          
          Time left:  epoch 0:00:16  total 0:01:45          
          Time left:  epoch 0:00:15  total 0:01:43          
          Time left:  epoch 0:00:15  total 0:01:41          
          Time left:  epoch 0:00:14  total 0:01:40          
          Time left:  epoch 0:00:14  total 0:01:39          
          Time left:  epoch 0:00:13  total 0:01:37          
          Time left:  epoch 0:00:13  total 0:01:36          
          Time left:  epoch 0:00:12  total 0:01:35          
          Time left:  epoch 0:00:12  total 0:01:34          
          Time left:  epoch 0:00:13  total 0:01:48          
          Time left:  epoch 0:00:13  total 0:01:46          
          Time left:  epoch 0:00:12  total 0:01:45          
          Time left:  epoch 0:00:12  total 0:01:43          
          Time left:  epoch 0:00:11  total 0:01:42          
          Time left:  epoch 0:00:11  total 0:01:41          
          Time left:  epoch 0:00:10  total 0:01:40          
          Time left:  epoch 0:00:10  total 0:01:39          
          Time left:  epoch 0:00:09  total 0:01:38          
          Time left:  epoch 0:00:09  total 0:01:37          
          Time left:  epoch 0:00:09  total 0:01:36          
          Time left:  epoch 0:00:08  total 0:01:35          
          Time left:  epoch 0:00:09  total 0:01:44          
          Time left:  epoch 0:00:08  total 0:01:42          
          Time left:  epoch 0:00:08  total 0:01:41          
          Time left:  epoch 0:00:07  total 0:01:40          
          Time left:  epoch 0:00:07  total 0:01:39          
          Time left:  epoch 0:00:06  total 0:01:38          
          Time left:  epoch 0:00:06  total 0:01:37          
          Time left:  epoch 0:00:06  total 0:01:36          
          Time left:  epoch 0:00:05  total 0:01:36          
          Time left:  epoch 0:00:05  total 0:01:35          
          Time left:  epoch 0:00:04  total 0:01:34          
          Time left:  epoch 0:00:04  total 0:01:33          
          Time left:  epoch 0:00:04  total 0:01:32          
          Time left:  epoch 0:00:03  total 0:01:39          
          Time left:  epoch 0:00:03  total 0:01:38          
          Time left:  epoch 0:00:03  total 0:01:37          
          Time left:  epoch 0:00:02  total 0:01:36          
          Time left:  epoch 0:00:02  total 0:01:35          
          Time left:  epoch 0:00:01  total 0:01:34          
          Time left:  epoch 0:00:01  total 0:01:33          
          Time left:  epoch 0:00:01  total 0:01:33          
          Time left:  epoch 0:00:00  total 0:01:32          
          Time left:  epoch 0:00:00  total 0:01:31          
          Time left:  epoch 0:00:00  total 0:01:30          

actor                                                       
  clip fraction                                       0.0137
  entropy                                             -0.884
  iterations                                              80
  kl                                                 0.00277
  loss                                             -0.000817
  std                                                    0.1
  stop                                                     0
critic                                                      
  iterations                                              80
  loss                                                  51.1
  v                                                     33.3
test                                                        
  action                                                    
    max                                                 1.34
    mean                                              0.0289
    min                                                -1.31
    size                                               5,000
    std                                                0.802
  episode length                                            
    max                                                1,000
    mean                                               1e+03
    min                                                1,000
    size                                                   5
    std                                                    0
  episode score                                             
    max                                                  786
    mean                                                 782
    min                                                  776
    size                                                   5
    std                                                 3.45
train                                                       
  action                                                    
    max                                                 1.42
    mean                                              0.0297
    min                                                -1.38
    size                                              20,000
    std                                                0.804
  episode length                                            
    max                                                1,000
    mean                                               1e+03
    min                                                1,000
    size                                                  20
    std                                                    0
  episode score                                             
    max                                                  792
    mean                                                 780
    min                                                  759
    size                                                  20
    std                                                  7.2
  episodes                                                20
  epoch seconds                                         26.6
  epoch steps                                          2e+04
  epochs                                                   1
  seconds                                               26.6
  steps                                                2e+04
  steps per second                                       753
  worker steps                                         2e+04

Logging data to data/local/experiments/tonic/swimmer-swim/ncap_ppo/log.csv
          Time left:  epoch 0:00:26  total 0:01:46          
          Time left:  epoch 0:00:26  total 0:01:45          
          Time left:  epoch 0:00:26  total 0:01:50          
          Time left:  epoch 0:00:26  total 0:01:48          
          Time left:  epoch 0:00:25  total 0:01:47          
          Time left:  epoch 0:00:24  total 0:01:46          
          Time left:  epoch 0:00:24  total 0:01:45          
          Time left:  epoch 0:00:23  total 0:01:44          
          Time left:  epoch 0:00:23  total 0:01:43          
          Time left:  epoch 0:00:22  total 0:01:42          
          Time left:  epoch 0:00:22  total 0:01:41          
          Time left:  epoch 0:00:21  total 0:01:40          
          Time left:  epoch 0:00:21  total 0:01:39          
          Time left:  epoch 0:00:20  total 0:01:38          
          Time left:  epoch 0:00:20  total 0:01:42          
          Time left:  epoch 0:00:20  total 0:01:41          
          Time left:  epoch 0:00:19  total 0:01:40          
          Time left:  epoch 0:00:19  total 0:01:39          
          Time left:  epoch 0:00:18  total 0:01:38          
          Time left:  epoch 0:00:18  total 0:01:37          
          Time left:  epoch 0:00:17  total 0:01:36          
          Time left:  epoch 0:00:17  total 0:01:36          
          Time left:  epoch 0:00:16  total 0:01:35          
          Time left:  epoch 0:00:16  total 0:01:34          
          Time left:  epoch 0:00:15  total 0:01:33          
          Time left:  epoch 0:00:15  total 0:01:32          
          Time left:  epoch 0:00:14  total 0:01:31          
          Time left:  epoch 0:00:14  total 0:01:34          
          Time left:  epoch 0:00:14  total 0:01:33          
          Time left:  epoch 0:00:13  total 0:01:33          
          Time left:  epoch 0:00:13  total 0:01:32          
          Time left:  epoch 0:00:12  total 0:01:31          
          Time left:  epoch 0:00:12  total 0:01:30          
          Time left:  epoch 0:00:11  total 0:01:29          
          Time left:  epoch 0:00:11  total 0:01:29          
          Time left:  epoch 0:00:10  total 0:01:28          
          Time left:  epoch 0:00:10  total 0:01:27          
          Time left:  epoch 0:00:09  total 0:01:27          
          Time left:  epoch 0:00:09  total 0:01:26          
          Time left:  epoch 0:00:09  total 0:01:28          
          Time left:  epoch 0:00:08  total 0:01:27          
          Time left:  epoch 0:00:08  total 0:01:27          
          Time left:  epoch 0:00:07  total 0:01:26          
          Time left:  epoch 0:00:07  total 0:01:25          
          Time left:  epoch 0:00:06  total 0:01:24          
          Time left:  epoch 0:00:06  total 0:01:24          
          Time left:  epoch 0:00:06  total 0:01:23          
          Time left:  epoch 0:00:05  total 0:01:22          
          Time left:  epoch 0:00:05  total 0:01:22          
          Time left:  epoch 0:00:04  total 0:01:21          
          Time left:  epoch 0:00:04  total 0:01:20          
          Time left:  epoch 0:00:03  total 0:01:22          
          Time left:  epoch 0:00:03  total 0:01:21          
          Time left:  epoch 0:00:03  total 0:01:21          
          Time left:  epoch 0:00:02  total 0:01:20          
          Time left:  epoch 0:00:02  total 0:01:19          
          Time left:  epoch 0:00:01  total 0:01:19          
          Time left:  epoch 0:00:01  total 0:01:18          
          Time left:  epoch 0:00:00  total 0:01:17          
          Time left:  epoch 0:00:00  total 0:01:17          
          Time left:  epoch 0:00:00  total 0:01:16          

actor                                                       
  clip fraction                                       0.0202
  entropy                                             -0.884
  iterations                                              80
  kl                                                 0.00246
  loss                                              -0.00122
  std                                                    0.1
  stop                                                     0
critic                                                      
  iterations                                              80
  loss                                                  3.73
  v                                                     66.8
test                                                        
  action                                                    
    max                                                 1.38
    mean                                              0.0413
    min                                                -1.35
    size                                               5,000
    std                                                0.811
  episode length                                            
    max                                                1,000
    mean                                               1e+03
    min                                                1,000
    size                                                   5
    std                                                    0
  episode score                                             
    max                                                  797
    mean                                                 789
    min                                                  777
    size                                                   5
    std                                                 7.47
train                                                       
  action                                                    
    max                                                 1.37
    mean                                              0.0211
    min                                                -1.43
    size                                              20,000
    std                                                0.809
  episode length                                            
    max                                                1,000
    mean                                               1e+03
    min                                                1,000
    size                                                  20
    std                                                    0
  episode score                                             
    max                                                  796
    mean                                                 787
    min                                                  777
    size                                                  20
    std                                                 5.68
  episodes                                                40
  epoch seconds                                         28.3
  epoch steps                                          2e+04
  epochs                                                   2
  seconds                                                 55
  steps                                                4e+04
  steps per second                                       706
  worker steps                                         4e+04

          Time left:  epoch 0:00:27  total 0:01:22          
          Time left:  epoch 0:00:26  total 0:01:21          
          Time left:  epoch 0:00:26  total 0:01:21          
          Time left:  epoch 0:00:26  total 0:01:22          
          Time left:  epoch 0:00:25  total 0:01:21          
          Time left:  epoch 0:00:25  total 0:01:20          
          Time left:  epoch 0:00:24  total 0:01:20          
          Time left:  epoch 0:00:24  total 0:01:19          
          Time left:  epoch 0:00:23  total 0:01:18          
          Time left:  epoch 0:00:23  total 0:01:18          
          Time left:  epoch 0:00:22  total 0:01:17          
          Time left:  epoch 0:00:22  total 0:01:16          
          Time left:  epoch 0:00:21  total 0:01:16          
          Time left:  epoch 0:00:21  total 0:01:15          
          Time left:  epoch 0:00:20  total 0:01:14          
          Time left:  epoch 0:00:20  total 0:01:14          
          Time left:  epoch 0:00:20  total 0:01:15          
          Time left:  epoch 0:00:19  total 0:01:14          
          Time left:  epoch 0:00:19  total 0:01:13          
          Time left:  epoch 0:00:18  total 0:01:13          
          Time left:  epoch 0:00:18  total 0:01:12          
          Time left:  epoch 0:00:17  total 0:01:11          
          Time left:  epoch 0:00:17  total 0:01:11          
          Time left:  epoch 0:00:16  total 0:01:10          
          Time left:  epoch 0:00:16  total 0:01:10          
          Time left:  epoch 0:00:15  total 0:01:09          
          Time left:  epoch 0:00:15  total 0:01:08          
          Time left:  epoch 0:00:14  total 0:01:08          
          Time left:  epoch 0:00:14  total 0:01:09          
          Time left:  epoch 0:00:14  total 0:01:08          
          Time left:  epoch 0:00:13  total 0:01:07          
Saving weights to data/local/experiments/tonic/swimmer-swim/ncap_ppo/checkpoints/step_50000.pt
          Time left:  epoch 0:00:13  total 0:01:07          
          Time left:  epoch 0:00:12  total 0:01:06          
          Time left:  epoch 0:00:12  total 0:01:05          
          Time left:  epoch 0:00:11  total 0:01:05          
          Time left:  epoch 0:00:11  total 0:01:04          
          Time left:  epoch 0:00:10  total 0:01:04          
          Time left:  epoch 0:00:10  total 0:01:03          
          Time left:  epoch 0:00:09  total 0:01:02          
          Time left:  epoch 0:00:09  total 0:01:02          
          Time left:  epoch 0:00:09  total 0:01:03          
          Time left:  epoch 0:00:08  total 0:01:02          
          Time left:  epoch 0:00:08  total 0:01:01          
          Time left:  epoch 0:00:07  total 0:01:01          
          Time left:  epoch 0:00:07  total 0:01:00          
          Time left:  epoch 0:00:06  total 0:01:00          
          Time left:  epoch 0:00:06  total 0:00:59          
          Time left:  epoch 0:00:05  total 0:00:58          
          Time left:  epoch 0:00:05  total 0:00:58          
          Time left:  epoch 0:00:04  total 0:00:57          
          Time left:  epoch 0:00:04  total 0:00:57          
          Time left:  epoch 0:00:03  total 0:00:56          
          Time left:  epoch 0:00:03  total 0:00:56          
          Time left:  epoch 0:00:03  total 0:00:56          
          Time left:  epoch 0:00:02  total 0:00:56          
          Time left:  epoch 0:00:02  total 0:00:55          
          Time left:  epoch 0:00:01  total 0:00:54          
          Time left:  epoch 0:00:01  total 0:00:54          
          Time left:  epoch 0:00:00  total 0:00:53          
          Time left:  epoch 0:00:00  total 0:00:53          
          Time left:  epoch 0:00:00  total 0:00:52          

actor                                                       
  clip fraction                                       0.0179
  entropy                                             -0.884
  iterations                                              80
  kl                                                 0.00224
  loss                                              -0.00102
  std                                                    0.1
  stop                                                     0
critic                                                      
  iterations                                              80
  loss                                                 0.606
  v                                                     77.2
test                                                        
  action                                                    
    max                                                 1.41
    mean                                              0.0205
    min                                                -1.34
    size                                               5,000
    std                                                 0.82
  episode length                                            
    max                                                1,000
    mean                                               1e+03
    min                                                1,000
    size                                                   5
    std                                                    0
  episode score                                             
    max                                                  807
    mean                                                 799
    min                                                  796
    size                                                   5
    std                                                 4.48
train                                                       
  action                                                    
    max                                                 1.43
    mean                                              0.0207
    min                                                -1.41
    size                                              20,000
    std                                                0.816
  episode length                                            
    max                                                1,000
    mean                                               1e+03
    min                                                1,000
    size                                                  20
    std                                                    0
  episode score                                             
    max                                                  806
    mean                                                 794
    min                                                  787
    size                                                  20
    std                                                 5.17
  episodes                                                60
  epoch seconds                                         28.2
  epoch steps                                          2e+04
  epochs                                                   3
  seconds                                               83.2
  steps                                                6e+04
  steps per second                                       710
  worker steps                                         6e+04

          Time left:  epoch 0:00:27  total 0:00:55          
          Time left:  epoch 0:00:27  total 0:00:54          
          Time left:  epoch 0:00:26  total 0:00:54          
          Time left:  epoch 0:00:26  total 0:00:53          
          Time left:  epoch 0:00:25  total 0:00:53          
          Time left:  epoch 0:00:25  total 0:00:53          
          Time left:  epoch 0:00:25  total 0:00:52          
          Time left:  epoch 0:00:24  total 0:00:52          
          Time left:  epoch 0:00:24  total 0:00:51          
          Time left:  epoch 0:00:23  total 0:00:51          
          Time left:  epoch 0:00:23  total 0:00:50          
          Time left:  epoch 0:00:22  total 0:00:50          
          Time left:  epoch 0:00:22  total 0:00:49          
          Time left:  epoch 0:00:21  total 0:00:48          
          Time left:  epoch 0:00:21  total 0:00:48          
          Time left:  epoch 0:00:20  total 0:00:47          
          Time left:  epoch 0:00:20  total 0:00:47          
          Time left:  epoch 0:00:19  total 0:00:47          
          Time left:  epoch 0:00:19  total 0:00:46          
          Time left:  epoch 0:00:18  total 0:00:46          
          Time left:  epoch 0:00:18  total 0:00:45          
          Time left:  epoch 0:00:17  total 0:00:45          
          Time left:  epoch 0:00:17  total 0:00:44          
          Time left:  epoch 0:00:16  total 0:00:44          
          Time left:  epoch 0:00:16  total 0:00:43          
          Time left:  epoch 0:00:15  total 0:00:43          
          Time left:  epoch 0:00:15  total 0:00:42          
          Time left:  epoch 0:00:14  total 0:00:42          
          Time left:  epoch 0:00:14  total 0:00:41          
          Time left:  epoch 0:00:14  total 0:00:41          
          Time left:  epoch 0:00:13  total 0:00:41          
          Time left:  epoch 0:00:13  total 0:00:40          
          Time left:  epoch 0:00:12  total 0:00:40          
          Time left:  epoch 0:00:12  total 0:00:39          
          Time left:  epoch 0:00:11  total 0:00:38          
          Time left:  epoch 0:00:11  total 0:00:38          
          Time left:  epoch 0:00:10  total 0:00:37          
          Time left:  epoch 0:00:10  total 0:00:37          
          Time left:  epoch 0:00:09  total 0:00:36          
          Time left:  epoch 0:00:09  total 0:00:36          
          Time left:  epoch 0:00:08  total 0:00:35          
          Time left:  epoch 0:00:08  total 0:00:35          
          Time left:  epoch 0:00:08  total 0:00:35          
          Time left:  epoch 0:00:07  total 0:00:34          
          Time left:  epoch 0:00:07  total 0:00:34          
          Time left:  epoch 0:00:06  total 0:00:33          
          Time left:  epoch 0:00:06  total 0:00:33          
          Time left:  epoch 0:00:05  total 0:00:32          
          Time left:  epoch 0:00:05  total 0:00:32          
          Time left:  epoch 0:00:04  total 0:00:31          
          Time left:  epoch 0:00:04  total 0:00:31          
          Time left:  epoch 0:00:04  total 0:00:30          
          Time left:  epoch 0:00:03  total 0:00:30          
          Time left:  epoch 0:00:03  total 0:00:29          
          Time left:  epoch 0:00:02  total 0:00:29          
          Time left:  epoch 0:00:02  total 0:00:29          
          Time left:  epoch 0:00:01  total 0:00:28          
          Time left:  epoch 0:00:01  total 0:00:28          
          Time left:  epoch 0:00:00  total 0:00:27          
          Time left:  epoch 0:00:00  total 0:00:27          
          Time left:  epoch 0:00:00  total 0:00:26          

actor                                                       
  clip fraction                                       0.0348
  entropy                                             -0.884
  iterations                                              80
  kl                                                 0.00398
  loss                                              -0.00191
  std                                                    0.1
  stop                                                     0
critic                                                      
  iterations                                              80
  loss                                                 0.229
  v                                                     80.4
test                                                        
  action                                                    
    max                                                 1.34
    mean                                              0.0336
    min                                                -1.38
    size                                               5,000
    std                                                0.828
  episode length                                            
    max                                                1,000
    mean                                               1e+03
    min                                                1,000
    size                                                   5
    std                                                    0
  episode score                                             
    max                                                  815
    mean                                                 804
    min                                                  796
    size                                                   5
    std                                                 6.94
train                                                       
  action                                                    
    max                                                 1.39
    mean                                              0.0209
    min                                                -1.42
    size                                              20,000
    std                                                0.825
  episode length                                            
    max                                                1,000
    mean                                               1e+03
    min                                                1,000
    size                                                  20
    std                                                    0
  episode score                                             
    max                                                  808
    mean                                                 802
    min                                                  793
    size                                                  20
    std                                                 4.18
  episodes                                                80
  epoch seconds                                         27.7
  epoch steps                                          2e+04
  epochs                                                   4
  seconds                                                111
  steps                                                8e+04
  steps per second                                       721
  worker steps                                         8e+04

          Time left:  epoch 0:00:27  total 0:00:27          
          Time left:  epoch 0:00:27  total 0:00:27          
          Time left:  epoch 0:00:26  total 0:00:26          
          Time left:  epoch 0:00:26  total 0:00:26          
          Time left:  epoch 0:00:25  total 0:00:25          
          Time left:  epoch 0:00:25  total 0:00:25          
          Time left:  epoch 0:00:25  total 0:00:25          
          Time left:  epoch 0:00:24  total 0:00:24          
          Time left:  epoch 0:00:24  total 0:00:24          
          Time left:  epoch 0:00:23  total 0:00:23          
          Time left:  epoch 0:00:23  total 0:00:23          
          Time left:  epoch 0:00:22  total 0:00:22          
          Time left:  epoch 0:00:22  total 0:00:22          
          Time left:  epoch 0:00:21  total 0:00:21          
          Time left:  epoch 0:00:21  total 0:00:21          
          Time left:  epoch 0:00:20  total 0:00:20          
          Time left:  epoch 0:00:20  total 0:00:20          
          Time left:  epoch 0:00:19  total 0:00:19          
          Time left:  epoch 0:00:19  total 0:00:19          
          Time left:  epoch 0:00:18  total 0:00:18          
          Time left:  epoch 0:00:18  total 0:00:18          
          Time left:  epoch 0:00:17  total 0:00:17          
          Time left:  epoch 0:00:17  total 0:00:17          
          Time left:  epoch 0:00:16  total 0:00:16          
          Time left:  epoch 0:00:16  total 0:00:16          
          Time left:  epoch 0:00:15  total 0:00:15          
          Time left:  epoch 0:00:15  total 0:00:15          
          Time left:  epoch 0:00:15  total 0:00:15          
          Time left:  epoch 0:00:14  total 0:00:14          
          Time left:  epoch 0:00:14  total 0:00:14          
          Time left:  epoch 0:00:13  total 0:00:13          
          Time left:  epoch 0:00:13  total 0:00:13          
          Time left:  epoch 0:00:12  total 0:00:12          
          Time left:  epoch 0:00:12  total 0:00:12          
          Time left:  epoch 0:00:11  total 0:00:11          
          Time left:  epoch 0:00:11  total 0:00:11          
          Time left:  epoch 0:00:10  total 0:00:10          
          Time left:  epoch 0:00:10  total 0:00:10          
          Time left:  epoch 0:00:09  total 0:00:09          
          Time left:  epoch 0:00:09  total 0:00:09          
          Time left:  epoch 0:00:09  total 0:00:09          
          Time left:  epoch 0:00:08  total 0:00:08          
          Time left:  epoch 0:00:08  total 0:00:08          
          Time left:  epoch 0:00:07  total 0:00:07          
          Time left:  epoch 0:00:07  total 0:00:07          
          Time left:  epoch 0:00:06  total 0:00:06          
          Time left:  epoch 0:00:06  total 0:00:06          
          Time left:  epoch 0:00:05  total 0:00:05          
          Time left:  epoch 0:00:05  total 0:00:05          
          Time left:  epoch 0:00:04  total 0:00:04          
          Time left:  epoch 0:00:04  total 0:00:04          
          Time left:  epoch 0:00:04  total 0:00:04          
          Time left:  epoch 0:00:03  total 0:00:03          
          Time left:  epoch 0:00:03  total 0:00:03          
          Time left:  epoch 0:00:02  total 0:00:02          
          Time left:  epoch 0:00:02  total 0:00:02          
          Time left:  epoch 0:00:01  total 0:00:01          
          Time left:  epoch 0:00:01  total 0:00:01          
          Time left:  epoch 0:00:00  total 0:00:00          
          Time left:  epoch 0:00:00  total 0:00:00          
             Time left:  epoch 0:00:00  total               

actor                                                       
  clip fraction                                       0.0382
  entropy                                             -0.884
  iterations                                              80
  kl                                                 0.00348
  loss                                              -0.00213
  std                                                    0.1
  stop                                                     0
critic                                                      
  iterations                                              80
  loss                                                 0.121
  v                                                     81.9
test                                                        
  action                                                    
    max                                                 1.42
    mean                                              0.0176
    min                                                -1.33
    size                                               5,000
    std                                                0.839
  episode length                                            
    max                                                1,000
    mean                                               1e+03
    min                                                1,000
    size                                                   5
    std                                                    0
  episode score                                             
    max                                                  828
    mean                                                 817
    min                                                  811
    size                                                   5
    std                                                 6.32
train                                                       
  action                                                    
    max                                                 1.42
    mean                                              0.0292
    min                                                -1.38
    size                                              20,000
    std                                                0.833
  episode length                                            
    max                                                1,000
    mean                                               1e+03
    min                                                1,000
    size                                                  20
    std                                                    0
  episode score                                             
    max                                                  818
    mean                                                 809
    min                                                  802
    size                                                  20
    std                                                 3.99
  episodes                                               100
  epoch seconds                                         27.9
  epoch steps                                          2e+04
  epochs                                                   5
  seconds                                                139
  steps                                                1e+05
  steps per second                                       717
  worker steps                                         1e+05


Saving weights to data/local/experiments/tonic/swimmer-swim/ncap_ppo/checkpoints/step_100000.pt

Let’s visualize the trained NCAP agent in the environment.

# play_model('data/local/experiments/tonic/swimmer-swim/ncap_ppo')
play_model('data/local/experiments/tonic/swimmer-swim/pretrained_ncap_ppo')

Loading weights from data/local/experiments/tonic/swimmer-swim/pretrained_ncap_ppo/checkpoints/step_10000000.pt
Reward for the run:  877.4716531769373
/tmp/ipykernel_8472/4073379236.py:65: MatplotlibDeprecationWarning: Auto-close()ing of figures upon backend switching is deprecated since 3.8 and will be removed two minor releases later.  To suppress this warning, explicitly call plt.close('all') first.
  matplotlib.use(orig_backend)  # Switch back to the original backend.
Loading...

This architecture was designed using the C. elegans motor circuit that can swim right at birth i.e it should already have really good priors. Can you try visualizing an agent with an untrained NCAP model. Can it swim?

3.3 Plot perfomance

Now we are going to visualize performance of our model

def plot_performance(paths, ax=None,title='Model Performance'):
    """
    Plots the performance of multiple models on the same axes using Seaborn for styling.

    Reads CSV log files from specified paths and plots the mean episode scores
    achieved during testing against the cumulative time steps for each model.
    The plot uses a logarithmic scale for the x-axis to better display the progression
    over a wide range of steps. Each line's legend is set to the name of the last folder
    in the path, representing the model's name. Seaborn styles are applied for enhanced visualization.

    Parameters:
    - paths (list of str): Paths to the experiment directories.
    - ax (matplotlib.axes.Axes, optional): A matplotlib axis object to plot on. If None,
      a new figure and axis are created.
    """
    # Set the Seaborn style
    sns.set(style="whitegrid")
    colors = sns.color_palette("colorblind")  # Colorblind-friendly palette

    if ax is None:
        fig, ax = plt.subplots()

    for index, path in enumerate(paths):
        # Extract the model name from the path
        model_name = os.path.basename(path.rstrip('/'))

        # Load data
        df = pd.read_csv(os.path.join(path, 'log.csv'))
        scores = df['test/episode_score/mean']
        lengths = df['test/episode_length/mean']
        steps = np.cumsum(lengths)
        sns.lineplot(x=steps, y=scores, ax=ax, label=model_name, color=colors[index % len(colors)])

    ax.set_xscale('log')
    ax.set_xlabel('Cumulative Time Steps')
    ax.set_ylabel('Max Episode Score')
    ax.legend()
    ax.set_title(title)
fig, ax = plt.subplots()

#Replace the paths with the path to models you trained to plot their performance.
paths = [
    'data/local/experiments/tonic/swimmer-swim/pretrained_ncap_ppo',
    'data/local/experiments/tonic/swimmer-swim/pretrained_mlp_ppo'
]
plot_performance(paths, ax=ax, title='MLP v/s NCAP')
plt.tight_layout()
plt.show()
<Figure size 640x480 with 1 Axes>

  • Compare the performance and learning curve of NCAP to MLP for the basic swimmer agent.

  • Try testing the model on a modification of the environment (e.g., the 12-link swimmer) it was trained on.

  • What happens if we remove certain weight constraints (e.g., sign constraint) from the NCAP model?

# add your code

Section 4: Visualizing the sparse network

Given the importance of architectural choices in our project we have provided a function which can visualize the network architecture. This includes the ability to render the NCAP network, representing the C. Elegans connectome.

import networkx as nx
import matplotlib.pyplot as plt
import numpy as np


def draw_network(mode='NCAP', N=2, include_speed_control=False, include_turn_control=False):

    """
    Draws a network graph for a swimmer model based on either NCAP or MLP architecture.

    Parameters:
    - mode (str): Determines the architecture type ('NCAP' or 'MLP'). Defaults to 'NCAP'.
    - N (int): Number of joints in the swimmer model. Defaults to 2.
    - include_speed_control (bool): If True, includes nodes for speed control in the graph.
    - include_turn_control (bool): If True, includes nodes for turn control in the graph.
    """


    G = nx.DiGraph()

    n=2+N*4

    nodes =dict()

    if include_speed_control:
      nodes['1-s'] = n+7
    if include_turn_control:
      nodes['r'] = n+5
      nodes['l'] = n+3

    nodes['o'] = n-1
    nodes['$o^d$'] = n-0
    nodes['$o^v$']= n-2

    custom_node_positions = {}
    custom_node_positions['o'] = (1, nodes['o'])
    custom_node_positions['$o^d$'] = (1.5, nodes['$o^d$'])
    custom_node_positions['$o^v$'] = (1.5, nodes['$o^v$'])


    if include_speed_control:
      custom_node_positions['1-s'] = (1.5, nodes['1-s'])
    if include_turn_control:
      custom_node_positions['r'] = (1.5, nodes['r'])
      custom_node_positions['l'] = (1.5, nodes['l'])

    for i in range(1,N+1):
      nodes[f'$q_{i}$'] = 4*(N-i) + 1
      nodes[f'$q^d_{i}$'] = 4*(N-i) + 2
      nodes[f'$q^v_{i}$'] = 4*(N-i)
      nodes[f'$b^d_{i}$'] = 4*(N-i) + 2
      nodes[f'$b^v_{i}$'] = 4*(N-i)
      nodes[f'$m^d_{i}$'] = 4*(N-i) + 2
      nodes[f'$m^v_{i}$'] = 4*(N-i)
      nodes['$\overset{..}{q}$' + f'$_{i}$'] = 4*(N-i) + 1

      custom_node_positions[f'$q_{i}$'] = (1, nodes[f'$q_{i}$'])
      custom_node_positions[f'$q^d_{i}$'] = (1.5, nodes[f'$q^d_{i}$'])
      custom_node_positions[f'$q^v_{i}$'] = (1.5, nodes[f'$q^v_{i}$'])
      custom_node_positions[f'$b^d_{i}$'] = (2, nodes[f'$b^d_{i}$'])
      custom_node_positions[f'$b^v_{i}$'] = (2, nodes[f'$b^v_{i}$'])
      custom_node_positions[f'$m^d_{i}$'] = (2.5, nodes[f'$m^d_{i}$'])
      custom_node_positions[f'$m^v_{i}$'] = (2.5, nodes[f'$m^v_{i}$'])
      custom_node_positions['$\overset{..}{q}$' + f'$_{i}$'] = (3, nodes['$\overset{..}{q}$' + f'$_{i}$'])

    for node, layer in nodes.items():
        G.add_node(node, layer=layer)

    if mode=='NCAP':
        # Add edges between nodes
        edges_colors = ['green', 'orange', 'green', 'green']
        edge_labels = {
            ('o', '$o^d$'):'+1',
            ('o', '$o^v$'):'-1',
            ('$o^d$', '$b^d_1$'):'o',
            ('$o^v$', '$b^v_1$'):'o'
            }

        if include_speed_control:
          edges_colors += ['orange']
          edge_labels[('1-s', '$b^d_1$')] = 's, to all b'
        if include_turn_control:
          edges_colors += ['green', 'green']
          edge_labels[('r', '$b^d_1$')] = 't'
          edge_labels[('l', '$b^v_1$')] = 't'


        for i in range(1,N+1):
          if i < N:
            edges_colors += ['green', 'orange', 'green', 'green']

            edge_labels[((f'$q_{i}$', f'$q^d_{i}$'))] = '+1'
            edge_labels[((f'$q_{i}$', f'$q^v_{i}$'))] = '-1'
            edge_labels[((f'$q^d_{i}$', f'$b^d_{i+1}$'))] = 'p'
            edge_labels[((f'$q^v_{i}$', f'$b^v_{i+1}$'))] = 'p'

          edges_colors += ['green', 'orange', 'green', 'orange',
                          'orange', 'green']

          edge_labels[((f'$b^d_{i}$', f'$m^d_{i}$'))] = 'i'
          edge_labels[((f'$b^d_{i}$', f'$m^v_{i}$'))] = 'c'
          edge_labels[((f'$b^v_{i}$', f'$m^v_{i}$'))] = 'i'
          edge_labels[((f'$b^v_{i}$', f'$m^d_{i}$'))] = 'c'
          edge_labels[((f'$m^v_{i}$', '$\overset{..}{q}$' + f'$_{i}$'))] = '-1'
          edge_labels[((f'$m^d_{i}$', '$\overset{..}{q}$' + f'$_{i}$'))] = '+1'

        edges = edge_labels.keys()

    elif mode=='MLP':
      edges = []
      layers = [1, 1.5, 2, 2.5, 3]
      layers_nodes = [[], [], [], [], []]
      for key, value in custom_node_positions.items():
        ind = layers.index(value[0])
        layers_nodes[ind].append(key)
      for layer_ind in range(len(layers_nodes) - 1):
        for node1 in layers_nodes[layer_ind]:
          for node2 in layers_nodes[layer_ind+1]:
            edges.append((node1, node2))
      edges_colors = np.repeat('gray', len(edges))


    G.add_edges_from(edges)

    # Draw the graph using the custom node positions
    options = {"edge_color": edges_colors, "edgecolors": "tab:gray", "node_size": 500, 'node_color':'white'}
    nx.draw(G, pos=custom_node_positions, with_labels=True, arrowstyle="-", arrowsize=20, **options)
    if mode=='NCAP':
      nx.draw_networkx_edge_labels(G, pos=custom_node_positions, edge_labels=edge_labels)
draw_network('MLP', N=6)
draw_network('NCAP', N=6, include_speed_control=True, include_turn_control=True)

Notice that the NCAP architecture is highly sparse and interpretable as compared to the MLP. Moreover notice that the ncap architecture can be completely embedded within a fully connected MLP of 3 hidden layers and ReLU nonlinearities. This enables us to do a thorough investigation into how specific architectural elements influence both performance and the learning process. By leveraging this capability, we can systematically analyze the impact of the architectural preferences inherent to the model and make better design choices.

It might be useful to also visualize the network’s activity for NCAP. Given it only has 4 learnable parameters it becomes much easier to interpret the network.


Conclusion

Based on the concepts we’ve discussed in this tutorial, you should now be equipped to advance with the project. Within this project, there are multiple pathways you can explore. For each pathway, you will delve deeply into one of the main sections outlined in this notebook, allowing for a thorough investigation of different factors that can influence performance:

  • Exploring the effects of environment: Investigate how different environmental settings impact agent performance. This could involve altering parameters of the environment or the types of tasks and reward functions. Understanding these effects can help in making better architectural choices and learning algorithms that result in agents that are robust and adaptable

  • Exploring the effects of learning algorithms: Standard RL algorithms often struggle with sparse and constrained networks, which can lead to suboptimal performance. Explore where these algorithms fail and analyze potential reasons for their limitations. Experiment with modifications or alternative algorithms that might overcome these challenges.

  • Exploring the effects of model architecture: Investigate how various architectural decisions within the NCAP model influence its performance. Visualize the model and its activity and explore potential improvements by tweaking architectural elements, assessing how these changes affect learning outcomes and operational efficiency.