gym_utils.py

import multiprocessing.pool
import os

import gym
import numpy as np
from stable_baselines3.common.monitor import Monitor
from stable_baselines3.common.vec_env import DummyVecEnv, SubprocVecEnv
from stable_baselines3.common.vec_env.vec_normalize import VecNormalize


def make_env(env_id, env_kwargs, seed, allow_early_resets=True):
    def _init():
        env = gym.make(env_id, **env_kwargs)
        env.action_space = gym.spaces.Box(
            low=-1.0, high=1.0, shape=env.action_space.shape, dtype=np.float
        )
        env.seed(seed)
        env = Monitor(env, None, allow_early_resets=True)
        return env

    return _init


def make_vec_envs(
    env_id,
    env_kwargs,
    seed,
    num_processes,
    gamma=None,
    vecnormalize=False,
    subproc=True,
    allow_early_resets=True,
):
    envs = [
        make_env(env_id, env_kwargs, seed + i, allow_early_resets=allow_early_resets)
        for i in range(num_processes)
    ]

    if subproc and num_processes > 1:
        envs = SubprocVecEnv(envs)
    else:
        envs = DummyVecEnv(envs)

    if vecnormalize:
        if gamma is not None:
            envs = VecNormalize(envs, gamma=gamma)
        else:
            envs = VecNormalize(envs, norm_reward=False)

    return envs


from evogym import get_full_connectivity, has_actuator, is_connected


def load_robot(ROOT_DIR, robot_name, task=None):

    if robot_name == "default":
        robot_name = task
        robot_file = os.path.join(ROOT_DIR, "robot_files", f"{robot_name}.txt")
        assert os.path.exists(robot_file), f"defalt robot is not set on the task {task}"
    else:
        robot_file = os.path.join(ROOT_DIR, "robot_files", f"{robot_name}.txt")

    body = np.loadtxt(robot_file)
    assert is_connected(body), f"robot {robot_name} is not fully connected"
    assert has_actuator(body), f"robot {robot_name} have not actuator block"

    connections = get_full_connectivity(body)
    robot = {"body": body, "connections": connections}
    return robot