environment.py

# Imports from external libraries
import numpy as np
from perlin_noise import PerlinNoise

# Imports from this project
import constants
import config


# The environment class defines where the robot starts, where the goal is, and where the obstacle is.
class Environment:

    # Initialisation of a new random environment
    def __init__(self):
        # Set a random dynamics
        self.resistance = np.zeros([200, 100], dtype=np.float32)
        self.set_dynamics()
        # Set a random observation function
        self.observation_params = np.zeros(constants.OBSERVATION_DIMENSION, dtype=np.float32)
        self.set_observation_function()
        # Set the initial state
        self.state = np.array([0.0, 0.0])
        self.random_reset()

    # Reset the environment to a random initial state on the left of the environment
    def random_reset(self):
        # Set the current state to a random state
        self.state[0] = 0.05
        self.state[1] = np.random.uniform(0, constants.ENVIRONMENT_HEIGHT)
        # Return the observation of the state
        observation = self.observation_function(self.state)
        return observation

    # Reset the environment to a specific initial state
    def specific_reset(self, init_state):
        # Set the current state to a random state
        self.state = init_state
        # Return the observation of the state
        observation = self.observation_function(self.state)
        return observation

    # Step the environment (during training) by executing one action
    def step_training(self, action):
        # Use the dynamics to calculate the next state
        next_state = self.dynamics(self.state, action)
        self.state = next_state
        # Get the observations of this state
        next_observation = self.observation_function(next_state)
        # Get the reward
        distance_to_goal = np.abs(next_state[0] - 1.95)
        reward = -distance_to_goal
        # Return the next observation and the reward
        return next_observation, reward

    # Step the environment (during testing) by executing one action
    def step_testing(self, action):
        # Use the dynamics to calculate the next state
        next_state = self.dynamics(self.state, action)
        self.state = next_state
        # Get the observations of this state
        next_observation = self.observation_function(next_state)
        # Get the reward
        distance_to_goal = np.abs(next_state[0] - 1.95)
        reward = -distance_to_goal
        # Return the next observation and the reward
        return next_observation, reward

    # Set the random dynamics for this environment
    def set_dynamics(self):
        # Get some random perlin noise functions
        noise_1 = PerlinNoise(octaves=5, seed=config.SEED)
        noise_2 = PerlinNoise(octaves=10, seed=config.SEED)
        noise_3 = PerlinNoise(octaves=20, seed=config.SEED)
        # Get a random stretch factor
        num_cells_across = 200
        num_cells_down = 100
        cells = np.zeros([num_cells_across, num_cells_down], dtype=np.float32)
        # Loop over all the cells in the environment
        for col in range(num_cells_across):
            for row in range(num_cells_down):
                # Calculate the cell value as the weighted average of the perlin noise functions
                cell_value = 3 * noise_1([col / num_cells_across, row / num_cells_down])
                cell_value += 2 * noise_2([col / num_cells_across, row / num_cells_down])
                cell_value += 1 * noise_3([col / num_cells_across, row / num_cells_down])
                cells[col, row] = cell_value
        # Make sure the starting area does not have too much resistance
        for col in range(0, 20):
            for row in range(num_cells_down):
                cells[col, row] *= 1.0 * (col / 20.0)
        # Normalise the cells to between 0 and 1
        min_cell = np.min(cells)
        max_cell = np.max(cells)
        cells_normalised = (cells - min_cell) / (max_cell - min_cell)
        # Apply a sigmoid function to stretch out the values
        stretch_factor = 10
        cells_stretched = 1 / (1 + np.exp(-stretch_factor * (cells_normalised - 0.5)))
        self.resistance = cells_stretched

    # Set a random observation function
    def set_observation_function(self):
        for i in range(constants.OBSERVATION_DIMENSION):
            self.observation_params[i] = np.random.uniform(-1.0, 1.0)

    # The true environment dynamics
    def dynamics(self, state, action):
        # Check that the action is within the robot's maximum action limit
        action = np.clip(action, -constants.MAX_ACTION_MAGNITUDE, constants.MAX_ACTION_MAGNITUDE)
        # Get the cell that corresponds to this state
        cell_x = int(state[0] * 100)
        if cell_x < 0:
            cell_x = 0
        if cell_x > 199:
            cell_x = 199
        cell_y = int(state[1] * 100)
        if cell_y < 0:
            cell_y = 0
        if cell_y > 99:
            cell_y = 99
        # Get the environment's resistance for this state
        resistance = self.resistance[cell_x, cell_y]
        # Compute the next state
        next_state = state + 3 * np.power(1 - resistance, 2) * action
        # Clip the state to stay within the environment boundaries
        next_state[0] = np.clip(next_state[0], 0, 2)  # X-axis boundary [0, 2]
        next_state[1] = np.clip(next_state[1], 0, 1)  # Y-axis boundary [0, 1]
        # Return the next state
        return next_state

    # The observation function, which creates an observation of a state
    def observation_function(self, state):
        observation = np.zeros(constants.OBSERVATION_DIMENSION, dtype=np.float32)
        observation[0] = self.observation_params[0] + np.sin(self.observation_params[1] * state[0])
        observation[1] = np.cos(self.observation_params[2] * state[1])
        observation[2] = np.tan(self.observation_params[3] + state[0] + self.observation_params[4])
        observation[3] = np.exp(self.observation_params[3]) * state[1]
        observation[4] = self.observation_params[4]
        return observation