simulation.py

# a simple agent using dynamical generative model of FEP
# that can control it's own temperature
# directly
import numpy as np
from matplotlib import pyplot as plt

from utils import running_mean


def get_noise(scale=0.1, positive=False):
    """ Simulate noise with the mean at 0 and the standard deviation of 0.1 """
    noise = np.random.normal() * scale
    if positive:
        noise = np.abs(noise)

    return noise


class InteroceptiveAgent:
    """Interoceptive agent resembling homeostatic regulation """

    def __init__(self, i_s_z_0=0.1, i_s_z_1=0.1, i_s_w_0=0.1, i_s_w_1=0.1,
                 action_bound=6, temp_const_change_initial=0,
                 learn_r_a=None,
                 temp_viable_range=10, dt=0.1, learn_r=0.1,
                 simulate_current=False):
        # sigma (variances) of sensory noise (z) and model noise (w)
        self.i_s_z_0 = i_s_z_0
        self.i_s_z_1 = i_s_z_1
        self.i_s_w_0 = i_s_w_0
        self.i_s_w_1 = i_s_w_1

        # learning rate
        self.learn_r = learn_r
        self.learn_r_a = learn_r_a if learn_r_a else learn_r

        self.dt = dt

        # goal temperature
        self.T0 = 30
        self.temp_viable_mean = 30
        self.temp_viable_range = temp_viable_range

        # action bound -- agent can't set action more or less than this
        self.temp_change_action_bound = action_bound
        self.temp_change_environment_initial = temp_const_change_initial

        # if the current should be simulated
        self.simulate_current = simulate_current

        self.reset()

    def reset(self):
        # errors
        self.i_e_z_0 = []
        self.i_e_z_1 = []
        self.i_e_w_0 = []
        self.i_e_w_1 = []

        # senses
        self.sense_i = []
        self.sense_i_d1 = []

        # world
        # >> temperature params
        # start with temperature at T0
        self.temp = [self.T0]
        # overall change of temperature
        self.temp_change = [0]
        # change of agent's temperature, generated by the environment
        self.temp_change_environment = [self.temp_change_environment_initial]
        self.temp_change_instant_update = [0]
        # change of temperature, generated by the action
        self.temp_change_action = [0]
        # >> light params
        # overall light
        self.light_change = [0]
        # light change generated by the environment
        self.light_change_instant = [0]
        self.light_change_environment = [0]
        # light change generated by the movement
        self.light_change_movement = []
        # sense of light
        self.ex_sense = []
        # >> movement params
        self.velocity_action = [0]
        self.velocity = [0]

        # action (interoceptive)
        # temperature change set by action
        self.temp_relative_change = [0]
        self.temp_desire = [self.temp_viable_mean]

        # brain state mu
        self.mu_i = [0]
        self.mu_i_d1 = [0]
        self.mu_i_d2 = [0]

        # variational free energy
        # on the interoceptive layer
        self.vfe_i = []
        # total
        self.vfe = []

    def generate_senses(self):
        self.generate_sense()
        self.generate_sense_d1()

    def generate_sense(self):
        sense = self.temp[-1] + get_noise()
        self.sense_i.append(sense)

    def generate_sense_d1(self):
        """ change in sensation is felt change in temperature """
        sense_i_d1 = self.temp_change[-1] + get_noise()

        self.sense_i_d1.append(sense_i_d1)

    def update_world(self):
        """ update world parameters """

        # an agent lives in the water world where
        # a temperature changes at each timestep
        # at some time this rate of change can be adjusted

        if self.time == 50:
            # jump to 2
            new_temp_chage_desired = 3
        elif self.time == 100:
            # jump to 5
            new_temp_chage_desired = 5
        elif self.time == 150:
            # jump to -1
            new_temp_chage_desired = -1
        elif self.time == 200:
            # jump to -6
            new_temp_chage_desired = -6.2
        elif self.time == 250:
            # jump to 0
            new_temp_chage_desired = 0
        else:
            new_temp_chage_desired = self.temp_change_environment[-1]

        temp_change_instant_update = new_temp_chage_desired - \
            self.temp_change_environment[-1]

        self.temp_change_instant_update.append(temp_change_instant_update)

    def upd_velocity(self):
        velocity = self.velocity_action[-1]

        # simulate external force -- current after time step 300
        if self.simulate_current and self.time > 300:
            velocity += 0.1

        self.velocity.append(velocity)

    def upd_light_change(self):
        """ Light change is composed of the light change
        due to the environment and light change due to the movement of the
        agent in the environment (set by its velocity) """

        # light change due to the environment is updated by the
        light_change_environment = self.light_change_environment[-1]
        # instant change of light (Euler integrated)
        light_change_environment += self.light_change_instant[-1] * self.dt
        self.light_change_environment.append(light_change_environment)

        # the light change due to velocity (Euler integrated)
        # assuming 1:1 relationship
        self.light_change_movement.append(self.velocity[-1])

        # combined light change is light change due to environment
        light_change = self.light_change_environment[-1]
        # and light change due to the movement
        light_change += self.light_change_movement[-1]

        self.light_change.append(light_change)

    def upd_temp_change(self):
        """ Change in temperature or the organism is composed of 
        the change in temperature due to the environment and the change
        of temperature, generated by the agent.
        """

        # as agent moves though the world, the change of temperature
        # is affected by instant changes in change in temperature
        # and by movement of the agent
        # by adding velocity effect (Euler integrated)
        temp_change_environment = self.temp_change_environment[-1]
        # first update by instant change in temperature change (Euler integrating)
        temp_change_environment += self.temp_change_instant_update[-1]
        # update movement effect by adding Euler integrated velocity
        # assuming 1:1 relationship
        temp_change_environment += self.velocity[-1] * self.dt
        self.temp_change_environment.append(temp_change_environment)

        # total temperature change is change due to the environment
        # and change produced by the organism
        # add noise
        action = self.temp_change_action[-1] + get_noise() * self.dt
        temp_change = self.temp_change_environment[-1] + action

        self.temp_change.append(temp_change)

    def upd_temp(self):
        # update temperature with the current temperature update
        upd = self.temp_change[-1]
        # Euler integrating it
        upd *= self.dt

        self.temp.append(self.temp[-1] + upd)

    def upd_err_z_0(self):
        # error between sensation and generated sensations
        self.i_e_z_0.append(self.sense_i[-1] - self.mu_i[-1])

    def upd_err_z_1(self):
        # error between first derivatives of sensation and generated sensations
        self.i_e_z_1.append(self.sense_i_d1[-1] - self.mu_i_d1[-1])

    def upd_err_w_0(self):
        # error between model and generation of model
        # here: model of dynamics at 1st derivative
        # and it's generation for the 1st derivative
        self.i_e_w_0.append(self.mu_i_d1[-1] + self.mu_i[-1] - self.temp_desire[-1])

    def upd_err_w_1(self):
        # error between model and generation of model
        # here: model of dynamics at 2nd derivative)
        # and it's generation for the 2nd derivative
        self.i_e_w_1.append(self.mu_i_d2[-1] + self.mu_i_d1[-1])

    def upd_mu_i_d2(self):
        upd = -self.learn_r * (self.i_e_w_1[-1] / self.i_s_w_1)
        upd *= self.dt

        self.mu_i_d2.append(self.mu_i_d2[-1] + upd)

    def upd_mu_i_d1(self):
        upd = -self.learn_r * (-self.i_e_z_1[-1] / self.i_s_z_1 +
                               self.i_e_w_0[-1] / self.i_s_w_0 + self.i_e_w_1[-1] / self.i_s_w_1)
        upd += self.mu_i_d2[-2]
        upd *= self.dt

        self.mu_i_d1.append(self.mu_i_d1[-1] + upd)

    def upd_mu_i(self):
        upd = -self.learn_r * \
            (-self.i_e_z_0[-1] / self.i_s_z_0 + self.i_e_w_0[-1] / self.i_s_w_0)
        upd += self.mu_i_d1[-2]
        upd *= self.dt

        self.mu_i.append(self.mu_i[-1] + upd)

    def upd_vfe_i(self):
        def sqrd_err(err, sigma):
            return np.power(err, 2) / sigma

        vfe_i = 0.5 * (sqrd_err(self.i_e_z_0[-1], self.i_s_z_0) +
                       sqrd_err(self.i_e_z_1[-1], self.i_s_z_1) +
                       sqrd_err(self.i_e_w_0[-1], self.i_s_w_0) +
                       sqrd_err(self.i_e_w_1[-1], self.i_s_w_1))

        self.vfe_i.append(vfe_i)

    def upd_action(self):
        # TODO action is noisy! Add noise here but not forget about integration
        # sensation change over action is always 1
        upd = -self.learn_r_a * 1 * (self.i_e_z_1[-1] / self.i_s_z_1)
        upd *= self.dt
        action = self.temp_change_action[-1] + upd

        # action must be bound by some plausible constraints
        # e.g. temperature can't change more than action bound at each timestep
        if abs(action) > self.temp_change_action_bound:
            action = np.sign(action) * self.temp_change_action_bound

        # update action
        self.temp_change_action.append(action)

    def upd_no_action(self):
        # if agent is not acting update it's variables anyway
        self.temp_change_action.append(self.temp_change_action[-1])

    def interoception(self):
        # update interoception
        #   --> update errors
        self.upd_err_z_0()
        self.upd_err_z_1()
        self.upd_err_w_0()
        self.upd_err_w_1()

        #  --> update recognition dynamics
        self.upd_mu_i_d2()
        self.upd_mu_i_d1()
        self.upd_mu_i()

        #  update free energy
        self.upd_vfe_i()

    def active_inference(self):
        self.temp_desire.append(self.temp_viable_mean)
        self.interoception()

    def upd_vfe(self):
        self.vfe.append(self.vfe_i[-1])

    def plot_results(self):
        fig, ax = plt.subplots(3, 2, constrained_layout=True)

        timeline = [s * self.dt for s in range(self.steps)]

        # Temperature
        min_temp = self.temp_viable_mean - self.temp_viable_range
        max_temp = self.temp_viable_mean + self.temp_viable_range
        ax[0][0].plot(timeline, self.temp[1:])
        ax[0][0].plot(timeline, self.temp_desire[1:], ls='--', lw=0.75, c='green')
        ax[0][0].set_title('Temperature')
        ax[0][0].set_xlabel('time step')
        ax[0][0].set_ylabel('temperature')
        ax[0][0].plot(timeline, np.ones_like(timeline) *
                      min_temp, lw=0.75, ls='--', c='red')
        ax[0][0].plot(timeline, np.ones_like(timeline) *
                      max_temp, lw=0.75, ls='--', c='red')
        ax[0][0].legend(['agent, $T$', 'goal, $T_{goal}$', 'viability'],
                        loc='lower right')
        ax[0][0].set_ylim(-10, 50)
        ax[0][0].set_xlim(-10, self.time + 30)

        # Temperature change
        temp_change_organism = np.array(self.temp_change_environment[1:]) + \
            np.array(self.temp_change_action[1:])
        ax[1][0].plot(timeline, self.temp_change_environment[1:])
        ax[1][0].plot(timeline, self.temp_change_action[1:])
        ax[1][0].plot(timeline, temp_change_organism, lw=1.25, ls='--', c='g')
        ax[1][0].plot(timeline, np.ones_like(timeline) * 0, lw=0.65, c='gray')
        ax[1][0].legend(['env-t, $\\dot{T}_e$', 'action, $\\dot{T}_i$',
                         'agent, $\\dot{T}$'],
                        loc='lower right',
                        # loc='center right',
                        labelspacing=0.3)
        ax[1][0].set_title('Temperature change')
        ax[1][0].set_xlabel('time step')
        ax[1][0].set_ylabel('temperature change')
        ax[1][0].set_ylim(-8, 8)
        ax[1][0].set_xlim(-10, self.time + 30)

        # Light change
        ax[2][0].plot(timeline, running_mean(self.light_change[1:]), lw=2)
        ax[2][0].set_title('Light change')
        ax[2][0].set_xlabel('time step')
        ax[2][0].set_ylabel('light change')
        ax[2][0].legend(['total, $\\dot{L}$'], loc='lower right')
        ax[2][0].set_xlim(-10, self.time + 30)

        # mu
        ax[0][1].plot(timeline, self.mu_i[1:])
        ax[0][1].plot(timeline, self.mu_i_d1[1:])
        ax[0][1].plot(timeline, self.mu_i_d2[1:])
        ax[0][1].set_title('Environmental variable, $\\mu$')
        ax[0][1].set_xlabel('time step')
        ax[0][1].set_ylabel('$\\mu$')
        ax[0][1].legend(['$\\mu_i$', "$\\mu_i'$", "$\\mu_i''$"], loc='upper right')
        ax[0][1].set_ylim(-10, 50)

        # VFE
        ax[1][1].plot(timeline, running_mean(self.vfe), lw=3, c="#3f92d2")
        ax[1][1].set_title('Variational free energy (VFE), $F$')
        ax[1][1].set_xlabel('time step')
        ax[1][1].set_ylabel('$F$')
        ax[1][1].set_ylim(-5, 500)

        # Error terms
        ax[2][1].plot(timeline, running_mean(self.i_e_z_0), lw=0.75)
        ax[2][1].plot(timeline, running_mean(self.i_e_z_1), lw=0.75)
        ax[2][1].plot(timeline, running_mean(self.i_e_w_0), lw=0.75)
        ax[2][1].plot(timeline, running_mean(self.i_e_w_1), lw=0.75)
        ax[2][1].set_ylim(-10, 10)
        ax[2][1].set_title('Error terms, $\\epsilon$')
        ax[2][1].set_xlabel('time step')
        ax[2][1].set_ylabel('$\\epsilon$')
        ax[2][1].legend(['$\\epsilon_{z0}$', '$\\epsilon_{z1}$', '$\\epsilon_{w0}$',
                         '$\\epsilon_{w1}$'], loc='upper right')

        return ax

    def simulate(self, sim_time=300, act_time=50):
        self.reset()
        self.act_time = act_time

        plt.ion()

        self.steps = int(sim_time / self.dt)
        print(f'Simulating {self.steps} steps')

        for step in range(self.steps):
            self.time = step * self.dt
            # update world
            self.update_world()
            self.upd_velocity()
            self.upd_temp_change()
            self.upd_temp()
            self.upd_light_change()

            # generate sensations
            self.generate_senses()

            # perform active inference
            self.active_inference()

            # update variational free energy
            self.upd_vfe()

            #  act
            if step * self.dt > act_time:
                self.upd_action()

            else:
                self.upd_no_action()

        self.plot_results()
        plt.show()


class ExteroceptiveAgent(InteroceptiveAgent):
    """Exteroceptive agent that infers the desired temperature
       and passes it down to the interoceptive layer"""

    def __init__(self, ex_s_z_0=0.01, simulate_current=True, **kwargs):
        # simulate current by default
        super().__init__(simulate_current=simulate_current, **kwargs)

        self.learn_r_ex = self.learn_r

        # sigma (variances)
        self.ex_s_z_0 = ex_s_z_0

    def reset(self):
        super().reset()

        # exteroceptive errors
        self.ex_e_z_0 = []
        self.ex_e_w_0 = []

        # brain state mu of the exteroception layer
        self.ex_mu = [0]

        # variational free energy of the exteroception layer
        self.vfe_ex = []

    def generate_senses(self):
        super().generate_senses()
        self.generate_ex_sense()

    def generate_ex_sense(self):
        ex_sense = self.light_change[-1] + get_noise()
        self.ex_sense.append(ex_sense)

    def upd_ex_err_z_0(self):
        # error between sensation and generated sensations
        self.ex_e_z_0.append(self.ex_sense[-1] - 0.1 * (-self.ex_mu[-1] + 30))

    def upd_ex_mu(self):
        upd = -self.learn_r_ex * \
            (0.1 * self.ex_e_z_0[-1] / self.ex_s_z_0)
        upd *= self.dt

        self.ex_mu.append(self.ex_mu[-1] + upd)

    def upd_vfe_ex(self):
        def sqrd_err(err, sigma):
            return np.power(err, 2) / sigma

        vfe_ex = 0.5 * (sqrd_err(self.ex_e_z_0[-1], self.ex_s_z_0))

        self.vfe_ex.append(vfe_ex)

    def upd_vfe(self):
        vfe = self.vfe_i[-1] + self.vfe_ex[-1]
        self.vfe.append(vfe)

    def exteroception(self):
        # the agent performs exteroception
        # that updates the desired temperature
        # setting new set point for the underlying
        # interoceptive inference

        #   --> update errors
        self.upd_ex_err_z_0()

        #  --> update recognition dynamics
        self.upd_ex_mu()

        #  update free energy
        self.upd_vfe_ex()

    def update_world(self):
        """ simulate light drop before temperature drop
            for a short period of time to show the proof of concept """
        # update world as in Interoceptive Agent
        super().update_world()

        # change in light starts before the temperature drop
        if int(self.time) == 175:
            light_change_instant = -0.7
        elif int(self.time) == 200:
            light_change_instant = 0.7
        else:
            light_change_instant = 0

        self.light_change_instant.append(light_change_instant)

    def active_inference(self):
        # exteroception first
        self.exteroception()
        # pass the prior about desired temperature to the interoception
        # generative model
        self.temp_desire.append(self.ex_mu[-1])
        # then perform interoception
        self.interoception()

    def simulate(self, sim_time=400, act_time=50):
        # simulate for 400 time steps by default
        super().simulate(sim_time=sim_time)

    def plot_results(self):
        ax = super().plot_results()

        timeline = [s * self.dt for s in range(self.steps)]

        ax[0][1].plot(timeline, self.ex_mu[1:], ls='--')
        ax[0][1].legend(['$\\mu_i$', "$\\mu_i'$", "$\\mu_i''$",
                         "$\\mu_e$"], loc='upper right')

        ax[1][1].plot(timeline, running_mean(self.vfe_i, 5), lw=1, c='tab:red')
        ax[1][1].plot(timeline, running_mean(self.vfe_ex), lw=1, c='tab:pink')
        ax[1][1].legend(['$F$', '$F_i$', '$F_e$'], loc='upper right')
        ax[1][1].set_ylim(-10, 200)

        ax[2][1].plot(timeline, running_mean(self.ex_e_z_0), lw=0.75)
        ax[2][1].legend(['$\\epsilon^{z0}_i$', '$\\epsilon^{z1}_i$', '$\\epsilon^{w0}_i$',
                         '$\\epsilon^{w1}_i$', '$\\epsilon^{z0}_e$'], loc='upper right')
        ax[2][1].set_ylim(-5, 5)

        return ax


class ActiveExteroceptiveAgent(ExteroceptiveAgent):
    """An agent now lives in the world where it can act
       In this world an agent can move up and down.
       Where moving up increases the amount of light 
       and, therefore, the temperature,
       while moving down has the contrary effect.

       The generative model of interoception is re-used for
       exteroceptive active inference. 

       As a result, an agent can both regulate it's temperature
       interoceptive and exteroceptively. At the same time,
       the generative dynamic, specifying a settling point resides on a level
       above. On a next level of the hierarchy the desired temperature is inferred."""

    def __init__(self, aex_s_z_0=0.1, aex_s_z_1=0.1, aex_s_w_0=0.1, aex_s_w_1=0.1,
                 aex_action_bound=0.2, supress_action=False, learn_r_aex=None,
                 supress_desired_temp_inference=False,
                 simulate_current=True,
                 **kwargs):
        super().__init__(simulate_current=simulate_current, **kwargs)

        self.learn_r_aex = self.learn_r if learn_r_aex is None else learn_r_aex

        # sigma (variances)
        self.aex_s_z_0 = aex_s_z_0
        self.aex_s_w_0 = aex_s_w_0

        # action bound of the action on the world
        self.aex_action_bound = aex_action_bound

        self.supress_action = supress_action
        self.supress_desired_temp_inference = supress_desired_temp_inference

    def reset(self):
        super().reset()

        # active exteroceptive errors
        self.aex_e_z_0 = [0]
        self.aex_e_w_0 = [0]

        # mu of the exteroceptive layer -- inferred change in light
        # produced by the agent
        self.aex_mu = [0]
        self.aex_mu_d1 = [0]

        # variational free energy of the active exteroception layer
        self.vfe_aex = []

        # action
        self.aex_action = [0]
        self.aex_action_pre_bound = [0]

    def update_world(self):
        super().update_world()

        if self.time < self.act_time:
            self.velocity_action.append(0)
            return

        if not self.supress_action:
            # extra_noise = get_noise() if np.random.rand() < 0.2 else 0
            noise = get_noise() if self.time < 300 else 0
            self.velocity_action.append(self.aex_action[-1] + noise)

    def upd_aex_mu_d1(self):
        upd = -self.learn_r_ex * (self.aex_e_w_0[-1] / self.aex_s_w_0)
        upd *= self.dt

        self.aex_mu_d1.append(self.aex_mu_d1[-1] + upd)

    def upd_aex_mu(self):
        upd = -self.learn_r_ex * \
            (+self.aex_e_z_0[-1] / self.aex_s_z_0 +
             +self.aex_e_w_0[-1] / self.aex_s_w_0
             )
        upd += self.aex_mu_d1[-1]
        upd *= self.dt

        self.aex_mu.append(self.aex_mu[-1] + upd)

    def upd_ex_err_z_0(self):
        # update of the error calculation at the exteroception layer
        # now we subtract the (prediction) of how much light change
        # is generated by the agent acting on the world
        prediction = self.ex_sense[-1] - self.aex_action[-1]

        self.ex_e_z_0.append(prediction - 0.1 * (-self.ex_mu[-1] + 30))

    def upd_aex_err_w_0(self):
        """ agent's expected dynamics is that environment temperature change
            needs to be explained
            set point is when the inferred change in light
            corresponds to the change in temperature """

        t_change_goal = self.sense_i_d1[-1] - self.temp_change_action[-1]
        self.aex_e_w_0.append(self.aex_mu_d1[-1] - t_change_goal + self.aex_mu[-1])

    def upd_aex_err_z_0(self):
        """ sensory error is the difference between the best guess
        about the change in light and the agent's prediction about it
        (through the copy of action)
        """

        error = self.aex_action[-1]
        self.aex_e_z_0.append(error - (-self.aex_mu[-1]))

    def upd_avfe_ex(self):
        """ active exteroception VFE """
        def sqrd_err(err, sigma):
            return np.power(err, 2) / sigma

        vfe_aex = 0.5 * (sqrd_err(self.aex_e_z_0[-1], self.aex_s_z_0)
                         + sqrd_err(self.aex_e_w_0[-1], self.aex_s_w_0))

        self.vfe_aex.append(vfe_aex)

    def upd_vfe(self):
        """ total VFE """
        vfe = self.vfe_i[-1] + self.vfe_ex[-1] + self.vfe_aex[-1]
        self.vfe.append(vfe)

    def upd_action(self):
        super().upd_action()

        # sensation change over action is always 1
        upd = -self.learn_r_aex * 1 * (self.aex_e_z_0[-1] / self.aex_s_z_0)
        upd *= self.dt
        aex_action = self.aex_action[-1] + upd
        aex_action += get_noise() * self.dt

        self.aex_action_pre_bound.append(aex_action)

        # action must be bound by some plausible constraints
        # e.g. can't move faster than some limit
        if abs(aex_action) > self.aex_action_bound:
            aex_action = np.sign(aex_action) * self.aex_action_bound

        # update action
        self.aex_action.append(aex_action)

    def upd_no_action(self):
        super().upd_no_action()
        self.aex_action.append(self.aex_action[-1])
        self.aex_action_pre_bound.append(self.aex_action_pre_bound[-1])

    def active_exteroception(self):
        """ an agents performs action
            based on it's exteroceptive inference
            about how temperature change causes light change """
        self.upd_aex_mu_d1()
        self.upd_aex_mu()

        # update errors
        self.upd_aex_err_z_0()
        self.upd_aex_err_w_0()

        # update free energy
        self.upd_avfe_ex()

    def active_inference(self):
        # exteroception first
        self.exteroception()

        # pass the prior about desired temperature to the generative model
        if not self.supress_desired_temp_inference:
            self.temp_desire.append(self.ex_mu[-1])
        else:
            self.temp_desire.append(self.temp_viable_mean)

        # then perform interoception
        self.interoception()
        # and active exteroception
        self.active_exteroception()

    def plot_results(self):
        ax = super().plot_results()

        timeline = [s * self.dt for s in range(self.steps)]

        # change in light
        ax[0][1].plot(timeline, self.aex_mu[1:], ls='--')
        ax[0][1].legend(['$\\mu_i$', "$\\mu_i'$", "$\\mu_i''$",
                         "$\\mu_e$", '$\\mu_a$'], loc='upper right')

        ax[2][1].plot(timeline, running_mean(self.aex_e_z_0[1:]), lw=0.75)
        ax[2][1].plot(timeline, running_mean(self.aex_e_w_0[1:]), lw=0.75)
        ax[2][1].legend(['$\\epsilon^{z0}_i$', '$\\epsilon^{z1}_i$', '$\\epsilon^{w0}_i$',
                         '$\\epsilon^{w1}_i$', '$\\epsilon^{z0}_e$',
                         '$\\epsilon^{z0}_a$', '$\\epsilon^{w0}_a$'],
                        loc='upper right',
                        labelspacing=0)
        ax[2][1].set_ylim(-3.5, 3.5)

        # change in light produced by the external force
        # is 0.1 after the time step 300
        light_change_external = np.zeros(self.steps + 1)
        light_change_external[3000:] = 0.1

        # environmental change in light (e.g. modelled sunset/sunrise)
        # is all change in light minus change in light generated by the agent's action
        # and generated by the external force dragging agent closer to surface
        light_change_env = np.array(self.light_change) - np.array(self.velocity_action) \
            - light_change_external

        # action  produced by the organism (including environmental noise)
        ax[2][0].plot(timeline, running_mean(self.velocity_action[1:]), lw=0.75)
        ax[2][0].plot(timeline, light_change_env[1:], lw=1, c='tab:red', ls='--')
        ax[2][0].plot(timeline[3000:], light_change_external[3001:], lw=1, c='tab:green')
        ax[2][0].legend(['total, $\\dot{L}$', 'action, $\\dot{L}_a$',
                         'env-t, $\\dot{L}_e$', 'external'],
                        loc='lower right')

        ax[1][1].plot(timeline, running_mean(self.vfe_aex), lw=1, c="tab:orange")
        ax[1][1].legend(['$F$', '$F_i$', '$F_e$', '$F_a$'])
        ax[1][1].set_ylim(-5, 175)

        plt.show()