phase-based Observations, Rewards, Coupling Ablation
# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice,
# this list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
#
# 3. Neither the name of the copyright holder nor the names of its
# contributors may be used to endorse or promote products derived from
# this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF
# THE POSSIBILITY OF SUCH DAMAGE.
#
# Copyright (c) 2021 ETH Zurich, Nikita Rudin

import os

LEGGED_GYM_ROOT_DIR = os.path.dirname(
    os.path.dirname(os.path.realpath(__file__)))
LEGGED_GYM_ENVS_DIR = os.path.join(LEGGED_GYM_ROOT_DIR, 'gym', 'envs')

gym/__init__.py
mit-biomimetics/ORCAgym
import importlib

from gym.utils.task_registry import task_registry

# * To add a new env:
# * 1. add the base env and env class name and location to the class dict
# * 2. add the config name and location to the config dict
# * 3. add the runner confg name and location to the runner config dict
# * 3. register the task experiment name to the env/config/ppo classes

# * from y import x where {y:x}
class_dict = {
    'LeggedRobot': '.base.legged_robot',
    'FixedRobot': '.base.fixed_robot',
    'MiniCheetah': '.mini_cheetah.mini_cheetah',
    'MiniCheetahRef': '.mini_cheetah.mini_cheetah_ref',
    'MiniCheetahOsc': '.mini_cheetah.mini_cheetah_osc',
    'MiniCheetahOscIK': '.mini_cheetah.mini_cheetah_osc_IK',
    'MiniCheetahFloquet': '.mini_cheetah.mini_cheetah_floquet',
}

config_dict = {
    'MiniCheetahCfg': '.mini_cheetah.mini_cheetah_config',
    'MiniCheetahRefCfg': '.mini_cheetah.mini_cheetah_ref_config',
    'MiniCheetahOscCfg': '.mini_cheetah.mini_cheetah_osc_config',
    'MiniCheetahOscIKCfg': '.mini_cheetah.mini_cheetah_osc_IK_config',
    'MiniCheetahFloquetCfg': '.mini_cheetah.mini_cheetah_floquet_config',
}

runner_config_dict = {
    'MiniCheetahRunnerCfg': '.mini_cheetah.mini_cheetah_config',
    'MiniCheetahRefRunnerCfg': '.mini_cheetah.mini_cheetah_ref_config',
    'MiniCheetahOscRunnerCfg': '.mini_cheetah.mini_cheetah_osc_config',
    'MiniCheetahOscIKRunnerCfg': '.mini_cheetah.mini_cheetah_osc_IK_config',
    'MiniCheetahFloquetRunnerCfg': '.mini_cheetah.mini_cheetah_floquet_config',
}

task_dict = {
    'mini_cheetah': ['MiniCheetah', 'MiniCheetahCfg', 'MiniCheetahRunnerCfg'],
    'mini_cheetah_ref':
        ['MiniCheetahRef', 'MiniCheetahRefCfg', 'MiniCheetahRefRunnerCfg'],
    'mini_cheetah_osc':
        ['MiniCheetahOsc', 'MiniCheetahOscCfg', 'MiniCheetahOscRunnerCfg'],
    'mini_cheetah_osc_IK':['MiniCheetahOscIK', 'MiniCheetahOscIKCfg',
                           'MiniCheetahOscIKRunnerCfg'],
    'mini_cheetah_floquet': ['MiniCheetahFloquet', 'MiniCheetahFloquetCfg',
                             'MiniCheetahFloquetRunnerCfg'],
}

for class_name, class_location in class_dict.items():
    locals()[class_name] = getattr(
        importlib.import_module(class_location, __name__), class_name)
for config_name, config_location in config_dict.items():
    locals()[config_name] = getattr(
        importlib.import_module(config_location, __name__), config_name)
for runner_config_name, runner_config_location in runner_config_dict.items():
    locals()[runner_config_name] = getattr(
        importlib.import_module(
            runner_config_location, __name__),
        runner_config_name)

for task_name, class_list in task_dict.items():
    task_registry.register(
        task_name,
        locals()[class_list[0]],
        locals()[class_list[1]](),
        locals()[class_list[2]]())

gym/envs/__init__.py
mit-biomimetics/ORCAgym
# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice,
# this list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
#
# 3. Neither the name of the copyright holder nor the names of its
# contributors may be used to endorse or promote products derived from
# this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF
# THE POSSIBILITY OF SUCH DAMAGE.
#
# Copyright (c) 2021 ETH Zurich, Nikita Rudin

import sys
from isaacgym import gymapi
from isaacgym import gymutil
import torch


# * Base class for RL tasks
class BaseTask():
    def __init__(self, gym, sim, cfg, sim_params, sim_device, headless):
        self.gym = gym
        self.sim = sim
        self.sim_params = sim_params
        self.sim_device = sim_device
        sim_device_type, self.sim_device_id = \
            gymutil.parse_device_str(self.sim_device)
        self.headless = headless

        # * env device is GPU only if sim is on GPU and use_gpu_pipeline=True,
        # * otherwise returned tensors are copied to CPU by physX.
        if sim_device_type == 'cuda' and sim_params.use_gpu_pipeline:
            self.device = self.sim_device
        else:
            self.device = 'cpu'

        # * graphics device for rendering, -1 for no rendering
        self.graphics_device_id = self.sim_device_id

        self.num_envs = cfg.env.num_envs
        self.num_actuators = cfg.env.num_actuators

        # * optimization flags for pytorch JIT
        torch._C._jit_set_profiling_mode(False)
        torch._C._jit_set_profiling_executor(False)

        # allocate buffers
        self.to_be_reset = torch.ones(self.num_envs,
                                      device=self.device,
                                      dtype=torch.bool)
        self.terminated = torch.ones(self.num_envs,
                                     device=self.device,
                                     dtype=torch.bool)
        self.episode_length_buf = torch.zeros(self.num_envs,
                                              device=self.device,
                                              dtype=torch.long)
        self.timed_out = torch.zeros(self.num_envs,
                                     device=self.device,
                                     dtype=torch.bool)

        self.extras = {}

        # todo: read from config
        self.enable_viewer_sync = True
        self.viewer = None

        # * if running with a viewer, set up keyboard shortcuts and camera
        if self.headless is False:
            # * subscribe to keyboard shortcuts
            self.viewer = self.gym.create_viewer(
                self.sim, gymapi.CameraProperties())
            self.gym.subscribe_viewer_keyboard_event(
                self.viewer, gymapi.KEY_ESCAPE, "QUIT")
            self.gym.subscribe_viewer_keyboard_event(
                self.viewer, gymapi.KEY_V, "toggle_viewer_sync")

    def get_states(self, obs_list):
        return torch.cat([self.get_state(obs) for obs in obs_list], dim=-1)

    def get_state(self, name):
        if name in self.scales.keys():
            return getattr(self, name)/self.scales[name]
        else:
            return getattr(self, name)

    def set_states(self, state_list, values):
        idx = 0
        for state in state_list:
            state_dim = getattr(self, state).shape[1]
            self.set_state(state, values[:, idx:idx+state_dim])
            idx += state_dim
        assert (idx == values.shape[1]), "Actions don't equal tensor shapes"

    def set_state(self, name, value):
        try:
            if name in self.scales.keys():
                setattr(self, name, value*self.scales[name])
            else:
                setattr(self, name, value)
        except AttributeError:
            print("Value for " + name + " does not match tensor shape")

    def _reset_idx(self, env_ids):
        """Reset selected robots"""
        raise NotImplementedError

    def reset(self):
        """ Reset all robots"""
        self._reset_idx(torch.arange(self.num_envs, device=self.device))
        self.step()

    def _reset_buffers(self):
        self.to_be_reset[:] = False
        self.terminated[:] = False
        self.timed_out[:] = False

    def compute_reward(self, reward_weights):
        ''' Compute and return a torch tensor of rewards
        reward_weights: dict with keys matching reward names, and values
            matching weights
        '''
        reward = torch.zeros(self.num_envs,
                             device=self.device, dtype=torch.float)
        for name, weight in reward_weights.items():
            reward += weight*self._eval_reward(name)
        return reward

    def _eval_reward(self, name):
        return eval('self._reward_'+name+'()')

    def _check_terminations_and_timeouts(self):
        """ Check if environments need to be reset
        """
        contact_forces = \
            self.contact_forces[:, self.termination_contact_indices, :]
        self.terminated = \
            torch.any(torch.norm(contact_forces, dim=-1) > 1., dim=1)
        self.timed_out = self.episode_length_buf > self.max_episode_length
        self.to_be_reset = self.timed_out | self.terminated

    def step(self, actions):
        raise NotImplementedError

    def _render(self, sync_frame_time=True):
        if self.viewer:
            # * check for window closed
            if self.gym.query_viewer_has_closed(self.viewer):
                sys.exit()

            # * check for keyboard events
            for evt in self.gym.query_viewer_action_events(self.viewer):
                if evt.action == "QUIT" and evt.value > 0:
                    sys.exit()
                elif evt.action == "toggle_viewer_sync" and evt.value > 0:
                    self.enable_viewer_sync = not self.enable_viewer_sync

            # * fetch results
            if self.device != 'cpu':
                self.gym.fetch_results(self.sim, True)

            # * step graphics
            if self.enable_viewer_sync:
                self.gym.step_graphics(self.sim)
                self.gym.draw_viewer(self.viewer, self.sim, True)
                if sync_frame_time:
                    self.gym.sync_frame_time(self.sim)
            else:
                self.gym.poll_viewer_events(self.viewer)
                # self.gym.draw_viewer(self.viewer, self.sim, True)

gym/envs/base/base_task.py
mit-biomimetics/ORCAgym
import os
import torch
import numpy as np
from isaacgym.torch_utils import to_torch, get_axis_params
from isaacgym import gymtorch, gymapi

from gym import LEGGED_GYM_ROOT_DIR
from gym.envs.base.base_task import BaseTask
from gym.utils.math import random_sample, exp_avg_filter
from gym.utils.helpers import class_to_dict


class FixedRobot(BaseTask):

    def __init__(self, gym, sim, cfg, sim_params, sim_device, headless):
        """ Parses the provided config file,
            calls create_sim() (which creates, simulation, terrain
            and environments), initilizes pytorch buffers used during training

        Args:
            cfg (Dict): Environment config file
            sim_params (gymapi.SimParams): simulation parameters
            physics_engine (gymapi.SimType): gymapi.SIM_PHYSX (must be PhysX)
            device_type (string): 'cuda' or 'cpu'
            device_id (int): 0, 1, ...
            headless (bool): Run without rendering if True
        """
        self.cfg = cfg
        self.sim_params = sim_params
        self.init_done = False
        self.device = sim_device
        self._parse_cfg(self.cfg)
        super().__init__(gym, sim, self.cfg, sim_params,
                         sim_device, headless)

        if not self.headless:
            self.set_camera(self.cfg.viewer.pos, self.cfg.viewer.lookat)
        self._initialize_sim()
        self.gym.prepare_sim(self.sim)
        self._init_buffers()
        self.init_done = True
        self.reset()

    def step(self):
        """ Apply actions, simulate, call self.post_physics_step()
            and pre_physics_step()

        Args:
            actions (torch.Tensor): Tensor of shape
                (num_envs, num_actions_per_env)
        """

        self._reset_buffers()
        self._pre_physics_step()
        # * step physics and render each frame
        self._render()
        for _ in range(self.cfg.control.decimation):
            self.torques = self._compute_torques()

            if self.cfg.asset.disable_motors:
                self.torques[:] = 0.
            torques_to_gym_tensor = torch.zeros(self.num_envs, self.num_dof,
                                                device=self.device)

            # todo encapsulate
            next_torques_idx = 0
            for dof_idx in range(self.num_dof):
                if self.cfg.control.actuated_joints_mask[dof_idx]:
                    torques_to_gym_tensor[:, dof_idx] = \
                        self.torques[:, next_torques_idx]
                    next_torques_idx += 1
                else:
                    torques_to_gym_tensor[:, dof_idx] = torch.zeros(
                        self.num_envs, device=self.device)

            self.gym.set_dof_actuation_force_tensor(
                self.sim, gymtorch.unwrap_tensor(torques_to_gym_tensor))
            self.gym.simulate(self.sim)
            if self.device == 'cpu':
                self.gym.fetch_results(self.sim, True)
            self.gym.refresh_dof_state_tensor(self.sim)

        self._post_physics_step()
        self._check_terminations_and_timeouts()

        env_ids = self.to_be_reset.nonzero().flatten()
        self._reset_idx(env_ids)

    def _pre_physics_step(self):
        pass

    def _post_physics_step(self):
        """
        check terminations, compute observations and rewards
        """
        self.gym.refresh_actor_root_state_tensor(self.sim)
        self.gym.refresh_net_contact_force_tensor(self.sim)

        self.episode_length_buf += 1
        self.common_step_counter += 1

        nact = self.num_actuators
        self.dof_pos_history[:, 2*nact:] = self.dof_pos_history[:, nact:2*nact]
        self.dof_pos_history[:, nact:2*nact] = self.dof_pos_history[:, :nact]
        self.dof_pos_history[:, :nact] = self.dof_pos_target

        self.dof_pos_obs = self.dof_pos - self.default_dof_pos

    def _reset_idx(self, env_ids):
        """ Reset some environments.
            Calls self._reset_dofs(env_ids), self._reset_root_states(env_ids),
            and self._resample_commands(env_ids)
            Logs episode info
            Resets some buffers

        Args:
            env_ids (list[int]): List of environment ids which must be reset
        """
        if len(env_ids) == 0:
            return

        # * reset robot states
        self._reset_system(env_ids)
        # * reset buffers
        self.dof_pos_history[env_ids] = 0.
        self.episode_length_buf[env_ids] = 0

    def _initialize_sim(self):
        """ Creates simulation, terrain and evironments
        """
        self.up_axis_idx = 2  # 2 for z, 1 for y -> adapt gravity accordingly
        self._create_ground_plane()
        self._create_envs()

    def set_camera(self, position, lookat):
        """
        Set the camera position and lookat.
        """
        cam_pos = gymapi.Vec3(position[0], position[1], position[2])
        cam_target = gymapi.Vec3(lookat[0], lookat[1], lookat[2])
        self.gym.viewer_camera_look_at(self.viewer, None, cam_pos, cam_target)

    # ------------- Callbacks --------------
    def _process_rigid_shape_props(self, props, env_id):
        """ Callback allowing to store/change/randomize the rigid shape
            properties of each environment. Called During environment creation.
            Base behavior: randomizes the friction of each environment

        Args:
            props (List[gymapi.RigidShapeProperties]): Properties of each
                shape of the asset
            env_id (int): Environment id

        Returns:
            [List[gymapi.RigidShapeProperties]]: Modified rigid shape
                properties
        """
        return props

    def _process_dof_props(self, props, env_id):
        """ Callback allowing to store/change/randomize the DOF properties of
            each environment. Called During environment creation.
            Base behavior: stores position, velocity and torques limits
                defined in the URDF

        Args:
            props (numpy.array): Properties of each DOF of the asset
            env_id (int): Environment id

        Returns:
            [numpy.array]: Modified DOF properties
        """
        if env_id == 0:  # ? why only for env_id == 0?
            self.dof_pos_limits = torch.zeros(self.num_dof, 2,
                                              dtype=torch.float,
                                              device=self.device)
            self.dof_vel_limits = torch.zeros(self.num_dof, dtype=torch.float,
                                              device=self.device)
            self.torque_limits = torch.zeros(self.num_actuators,
                                             dtype=torch.float,
                                             device=self.device)
            for i in range(len(props)):
                self.dof_pos_limits[i, 0] = props["lower"][i].item()
                self.dof_pos_limits[i, 1] = props["upper"][i].item()
                self.dof_vel_limits[i] = props["velocity"][i].item()
                try:  # ! this seems like a crime...
                    self.torque_limits[i] = props["effort"][i].item()
                except:
                    print("WARNING: passive joints need to be listed after "
                          + "active joints in the URDF.")
                # * soft limits
                m = (self.dof_pos_limits[i, 0] + self.dof_pos_limits[i, 1]) / 2
                r = self.dof_pos_limits[i, 1] - self.dof_pos_limits[i, 0]
                self.dof_pos_limits[i, 0] = \
                    m - 0.5 * r * self.cfg.reward_settings.soft_dof_pos_limit
                self.dof_pos_limits[i, 1] = \
                    m + 0.5 * r * self.cfg.reward_settings.soft_dof_pos_limit
        return props

    def _process_rigid_body_props(self, props, env_id):
        """
        Process rigid body properties.
        In `legged_robot` this is used to randomize the base mass.
        Implement as you see fit.
        """
        return props

    def _compute_torques(self):
        """ Compute torques from actions.
            Actions can be interpreted as position or velocity targets given
            to a PD controller, or directly as scaled torques.
            [NOTE]: torques must have the same dimension as the number of DOFs,
                even if some DOFs are not actuated.

        Args:
            actions (torch.Tensor): Actions

        Returns:
            [torch.Tensor]: Torques sent to the simulation
        """
        actuated_dof_pos = torch.zeros(self.num_envs, self.num_actuators,
                                       device=self.device)
        actuated_dof_vel = torch.zeros(self.num_envs, self.num_actuators,
                                       device=self.device)
        for dof_idx in range(self.num_dof):
            idx = 0
            if self.cfg.control.actuated_joints_mask[dof_idx]:
                actuated_dof_pos[:, idx] = self.dof_pos[:, dof_idx]
                actuated_dof_vel[:, idx] = self.dof_vel[:, dof_idx]
                idx += 1

        torques = (
            self.p_gains * (self.dof_pos_target + self.default_act_pos
                            - actuated_dof_pos)
            + self.d_gains * (self.dof_vel_target - actuated_dof_vel)
            + self.tau_ff)
        torques = torch.clip(torques, -self.torque_limits, self.torque_limits)
        return torques.view(self.torques.shape)

    def _reset_system(self, env_ids):
        """
        Reset the system.
        """

        if hasattr(self, self.cfg.init_state.reset_mode):
            eval(f"self.{self.cfg.init_state.reset_mode}(env_ids)")
        else:
            raise NameError(
                f"Unknown default setup: {self.cfg.init_state.reset_mode}")

        # self.root_states[env_ids] = self.base_init_state
        # self.root_states[env_ids, 7:13] = 0.0
        env_ids_int32 = env_ids.to(dtype=torch.int32)
        self.gym.set_dof_state_tensor_indexed(
            self.sim, gymtorch.unwrap_tensor(self.dof_state),
            gymtorch.unwrap_tensor(env_ids_int32), len(env_ids_int32))

    # * implement reset methods
    def reset_to_basic(self, env_ids):
        """
        Generate random samples for each entry of env_ids
        todo: pass in the actual number instead of the list env_ids
        """
        self.dof_pos[env_ids] = self.default_dof_pos
        self.dof_vel[env_ids] = 0

    def reset_to_range(self, env_ids):
        """
        Reset to a uniformly random distribution of states, sampled from a
        range for each state
        """
        # dof states
        self.dof_pos[env_ids] = random_sample(
            env_ids, self.dof_pos_range[:, 0],
            self.dof_pos_range[:, 1], device=self.device)
        self.dof_vel[env_ids] = random_sample(
            env_ids, self.dof_vel_range[:, 0],
            self.dof_vel_range[:, 1], device=self.device)

    # ----------------------------------------
    def _init_buffers(self):
        """ Initialize torch tensors which will contain simulation states and
        processed quantities
        """
        n_envs = self.num_envs
        # * get gym GPU state tensors
        actor_root_state = \
            self.gym.acquire_actor_root_state_tensor(self.sim)
        dof_state_tensor = \
            self.gym.acquire_dof_state_tensor(self.sim)
        net_contact_forces = \
            self.gym.acquire_net_contact_force_tensor(self.sim)
        self.gym.refresh_dof_state_tensor(self.sim)
        self.gym.refresh_actor_root_state_tensor(self.sim)
        self.gym.refresh_net_contact_force_tensor(self.sim)

        # * create some wrapper tensors for different slices
        # ! root_states probably not needed...
        self.root_states = gymtorch.wrap_tensor(actor_root_state)
        self.dof_state = gymtorch.wrap_tensor(dof_state_tensor)

        self.dof_pos = self.dof_state.view(n_envs, self.num_dof, 2)[..., 0]
        self.dof_vel = self.dof_state.view(n_envs, self.num_dof, 2)[..., 1]
        self.base_quat = self.root_states[:, 3:7]

        # * shape: num_envs, num_bodies, xyz axis
        self.contact_forces = \
            gymtorch.wrap_tensor(net_contact_forces).view(n_envs, -1, 3)

        # * initialize some data used later on
        self.common_step_counter = 0
        self.extras = {}
        self.gravity_vec = to_torch(get_axis_params(-1., self.up_axis_idx),
                                    device=self.device).repeat((n_envs, 1))
        self.forward_vec = to_torch([1., 0., 0.],
                                    device=self.device).repeat((n_envs, 1))
        self.torques = torch.zeros(n_envs, self.num_actuators,
                                   dtype=torch.float, device=self.device)
        self.p_gains = torch.zeros(self.num_actuators,
                                   dtype=torch.float, device=self.device)
        self.d_gains = torch.zeros(self.num_actuators,
                                   dtype=torch.float, device=self.device)
        self.actions = torch.zeros(n_envs, self.num_actuators,
                                   dtype=torch.float, device=self.device)
        self.dof_pos_target = torch.zeros(self.num_envs, self.num_actuators,
                                          dtype=torch.float,
                                          device=self.device)
        self.dof_vel_target = torch.zeros(self.num_envs,
                                          self.num_actuators,
                                          dtype=torch.float,
                                          device=self.device)
        self.tau_ff = torch.zeros(self.num_envs, self.num_actuators,
                                  dtype=torch.float, device=self.device)
        self.dof_pos_history = torch.zeros(self.num_envs,
                                           self.num_actuators*3,
                                           dtype=torch.float,
                                           device=self.device)
        # * joint positions offsets and PD gains
        # * added: default_act_pos, to differentiate from passive joints
        self.default_dof_pos = torch.zeros(self.num_dof,
                                           dtype=torch.float,
                                           device=self.device)
        self.default_act_pos = torch.zeros(self.num_actuators,
                                           dtype=torch.float,
                                           device=self.device)
        actuated_idx = []  # temp
        for i in range(self.num_dof):
            name = self.dof_names[i]
            # angle = self.cfg.init_state.default_joint_angles[name]
            angles = self.cfg.init_state.default_joint_angles
            # self.default_dof_pos[i] = angle
            found = False
            for dof_name in angles.keys():
                if dof_name in name:
                    self.default_dof_pos[i] = angles[dof_name]
                    found = True
            if not found:
                self.default_dof_pos[i] = 0.
                print(f"Default dof pos of joint {name} was not defined, "
                      + "setting to zero")

            found = False
            for dof_name in self.cfg.control.stiffness.keys():
                if dof_name in name:
                    self.p_gains[i] = self.cfg.control.stiffness[dof_name]
                    self.d_gains[i] = self.cfg.control.damping[dof_name]
                    self.default_act_pos[i] = angles[dof_name]
                    found = True
                    actuated_idx.append(i)
            if not found:
                try:
                    self.p_gains[i] = 0.
                    self.d_gains[i] = 0.
                    # todo remove if unnecessary
                    print("This should not happen anymore")
                    if self.cfg.control.control_type in ["P", "V"]:
                        print(f"PD gain of joint {name} not defined, "
                              + "set to zero")
                except:  # ! another crime - no bare excepts
                    pass

        self.default_dof_pos = self.default_dof_pos.unsqueeze(0)
        self.default_act_pos = self.default_act_pos.unsqueeze(0)
        # * store indices of actuated joints
        self.act_idx = to_torch(actuated_idx, dtype=torch.long,
                                device=self.device)
        # * check that init range highs and lows are consistent
        # * and repopulate to match
        if self.cfg.init_state.reset_mode == "reset_to_range":
            self.initialize_ranges_for_initial_conditions()

    def initialize_ranges_for_initial_conditions(self):
        self.dof_pos_range = torch.zeros(
            self.num_dof, 2,
            dtype=torch.float, device=self.device)
        self.dof_vel_range = torch.zeros(
            self.num_dof, 2,
            dtype=torch.float, device=self.device)

        for joint, vals in self.cfg.init_state.dof_pos_range.items():
            for i in range(self.num_dof):
                if joint in self.dof_names[i]:
                    self.dof_pos_range[i, :] = to_torch(vals,
                                                        device=self.device)

        for joint, vals in self.cfg.init_state.dof_vel_range.items():
            for i in range(self.num_dof):
                if joint in self.dof_names[i]:
                    self.dof_vel_range[i, :] = to_torch(vals,
                                                        device=self.device)

    def _create_ground_plane(self):
        """ Adds a ground plane to the simulation, sets friction and
            restitution based on the cfg.
        """
        plane_params = gymapi.PlaneParams()
        plane_params.normal = gymapi.Vec3(0.0, 0.0, 1.0)

        plane_params.static_friction = self.cfg.terrain.static_friction
        plane_params.dynamic_friction = self.cfg.terrain.dynamic_friction
        plane_params.restitution = self.cfg.terrain.restitution
        self.gym.add_ground(self.sim, plane_params)

    def _create_envs(self):
        """ Creates environments:
             1. loads the robot URDF/MJCF asset,
             2. For each environment
                2.1 creates the environment,
                2.2 calls DOF and Rigid shape properties callbacks,
                2.3 create actor with these properties and add them to the env
             3. Store indices of different bodies of the robot
        """
        asset_path = self.cfg.asset.file.format(
            LEGGED_GYM_ROOT_DIR=LEGGED_GYM_ROOT_DIR)
        asset_root = os.path.dirname(asset_path)
        asset_file = os.path.basename(asset_path)

        asset_options = gymapi.AssetOptions()
        asset_options.fix_base_link = self.cfg.asset.fix_base_link
        asset_options.density = self.cfg.asset.density
        asset_options.angular_damping = self.cfg.asset.angular_damping
        asset_options.linear_damping = self.cfg.asset.linear_damping
        asset_options.max_linear_velocity = self.cfg.asset.max_linear_velocity
        asset_options.armature = self.cfg.asset.armature
        asset_options.thickness = self.cfg.asset.thickness
        asset_options.disable_gravity = \
            self.cfg.asset.disable_gravity
        asset_options.default_dof_drive_mode = \
            self.cfg.asset.default_dof_drive_mode
        asset_options.collapse_fixed_joints = \
            self.cfg.asset.collapse_fixed_joints
        asset_options.replace_cylinder_with_capsule = \
            self.cfg.asset.replace_cylinder_with_capsule
        asset_options.flip_visual_attachments = \
            self.cfg.asset.flip_visual_attachments
        asset_options.max_angular_velocity = \
            self.cfg.asset.max_angular_velocity

        robot_asset = self.gym.load_asset(
            self.sim, asset_root, asset_file, asset_options)
        self.num_dof = self.gym.get_asset_dof_count(robot_asset)
        self.num_bodies = self.gym.get_asset_rigid_body_count(robot_asset)

        dof_props_asset = self.gym.get_asset_dof_properties(robot_asset)
        dof_props_asset['armature'] = self.cfg.asset.rotor_inertia
        dof_props_asset['damping'] = self.cfg.asset.joint_damping
        rigid_shape_props_asset = \
            self.gym.get_asset_rigid_shape_properties(robot_asset)

        # * save body names from the asset
        body_names = self.gym.get_asset_rigid_body_names(robot_asset)
        self.dof_names = self.gym.get_asset_dof_names(robot_asset)
        self.num_bodies = len(body_names)
        penalized_contact_names = []
        for name in self.cfg.asset.penalize_contacts_on:
            penalized_contact_names.extend(
                [s for s in body_names if name in s])
        termination_contact_names = []
        for name in self.cfg.asset.terminate_after_contacts_on:
            termination_contact_names.extend(
                [s for s in body_names if name in s])

        start_pose = gymapi.Transform()

        self._get_env_origins()
        env_lower = gymapi.Vec3(0., 0., 0.)  # ? what's this?
        env_upper = gymapi.Vec3(0., 0., 0.)
        self.actor_handles = []
        self.envs = []
        for i in range(self.num_envs):
            # create env instance
            env_handle = self.gym.create_env(
                self.sim, env_lower, env_upper, int(np.sqrt(self.num_envs)))
            pos = self.env_origins[i].clone()
            start_pose.p = gymapi.Vec3(*pos)

            rigid_shape_props = self._process_rigid_shape_props(
                rigid_shape_props_asset, i)
            self.gym.set_asset_rigid_shape_properties(
                robot_asset, rigid_shape_props)
            name = self.cfg.asset.file
            name = name.format(LEGGED_GYM_ROOT_DIR=LEGGED_GYM_ROOT_DIR)
            robot_handle = self.gym.create_actor(
                env_handle, robot_asset, start_pose, name, i,
                self.cfg.asset.self_collisions, 0)
            dof_props = self._process_dof_props(dof_props_asset, i)
            self.gym.set_actor_dof_properties(
                env_handle, robot_handle, dof_props)
            body_props = self.gym.get_actor_rigid_body_properties(
                env_handle, robot_handle)
            body_props = self._process_rigid_body_props(body_props, i)
            self.gym.set_actor_rigid_body_properties(
                env_handle, robot_handle, body_props, recomputeInertia=True)
            self.envs.append(env_handle)
            self.actor_handles.append(robot_handle)

        self.penalised_contact_indices = torch.zeros(
            len(penalized_contact_names),
            dtype=torch.long, device=self.device)
        for i in range(len(penalized_contact_names)):
            self.penalised_contact_indices[i] = \
                self.gym.find_actor_rigid_body_handle(
                    self.envs[0], self.actor_handles[0],
                    penalized_contact_names[i])

        self.termination_contact_indices = torch.zeros(
            len(termination_contact_names),
            dtype=torch.long, device=self.device)
        for i in range(len(termination_contact_names)):
            self.termination_contact_indices[i] = \
                self.gym.find_actor_rigid_body_handle(
                    self.envs[0], self.actor_handles[0],
                    termination_contact_names[i])

    def _get_env_origins(self):  # TODO: do without terrain
        """ Sets environment origins. On rough terrain the origins are defined
            by the terrain platforms. Otherwise create a grid.
        """
        # * removed terrain options
        self.custom_origins = False
        self.env_origins = torch.zeros(
            self.num_envs, 3,
            device=self.device)
        # * create a grid of robots
        num_cols = np.floor(np.sqrt(self.num_envs))
        num_rows = np.ceil(self.num_envs / num_cols)
        xx, yy = torch.meshgrid(torch.arange(num_rows), torch.arange(num_cols))
        spacing = self.cfg.env.env_spacing
        self.env_origins[:, 0] = spacing * xx.flatten()[:self.num_envs]
        self.env_origins[:, 1] = spacing * yy.flatten()[:self.num_envs]
        self.env_origins[:, 2] = self.cfg.env.root_height

    def _parse_cfg(self, cfg):
        self.dt = self.cfg.control.ctrl_dt
        self.scales = class_to_dict(self.cfg.scaling, self.device)
        # self.command_ranges = class_to_dict(self.cfg.commands.ranges)
        self.max_episode_length_s = self.cfg.env.episode_length_s
        self.max_episode_length = np.ceil(self.max_episode_length_s / self.dt)

    # ------------ reward functions----------------

    def _sqrdexp(self, x):
        """ shorthand helper for squared exponential
        """
        return torch.exp(
            -torch.square(x)/self.cfg.reward_settings.tracking_sigma)

    def _reward_torques(self):
        """Penalize torques"""
        return -torch.sum(torch.square(self.torques), dim=1)

    def _reward_dof_vel(self):
        """Penalize dof velocities"""
        return -torch.sum(torch.square(self.dof_vel), dim=1)

    def _reward_action_rate(self):
        """Penalize changes in actions"""
        nact = self.num_actuators
        dt2 = (self.dt*self.cfg.control.decimation)**2
        error = torch.square(
            self.dof_pos_history[:, :nact]
            - self.dof_pos_history[:, nact:2*nact])/dt2
        return -torch.sum(error, dim=1)

    def _reward_action_rate2(self):
        """Penalize changes in actions"""
        nact = self.num_actuators
        dt2 = (self.dt*self.cfg.control.decimation)**2
        error = torch.square(
            self.dof_pos_history[:, :nact]
            - 2*self.dof_pos_history[:, nact:2*nact]
            + self.dof_pos_history[:, 2*nact:])/dt2
        return -torch.sum(error, dim=1)

    def _reward_collision(self):
        """Penalize collisions on selected bodies"""
        return -torch.sum(
            1.*(torch.norm(
                self.contact_forces[:, self.penalised_contact_indices, :],
                dim=-1) > 0.1), dim=1)

    def _reward_termination(self):
        # Terminal reward / penalty
        return -self.terminated.float()

    def _reward_dof_pos_limits(self):
        """Penalize dof positions too close to the limit"""
        out_of_limits = (
            -(self.dof_pos - self.dof_pos_limits[:, 0])
            .clip(max=0.))  # lower limit
        out_of_limits += (
            self.dof_pos - self.dof_pos_limits[:, 1]).clip(min=0.)
        return -torch.sum(out_of_limits, dim=1)

    def _reward_dof_vel_limits(self):
        """Penalize dof velocities too close to the limit"""
        # * clip to max error = 1 rad/s per joint to avoid huge penalties
        return -torch.sum(
            (torch.abs(self.dof_vel)
             - self.dof_vel_limits*self.cfg.reward_settings.soft_dof_vel_limit)
            .clip(min=0., max=1.), dim=1)

gym/envs/base/fixed_robot.py
mit-biomimetics/ORCAgym

from .base_config import BaseConfig


class FixedRobotCfg(BaseConfig):
    class env:
        num_envs = 4096
        num_actuators = 1
        env_spacing = 4.  # not used with heightfields/trimeshes
        root_height = 2.
        episode_length_s = 4  # episode length in seconds

    class terrain:
        mesh_type = 'none'
        horizontal_scale = 0.1  # [m]
        vertical_scale = 0.005  # [m]
        border_size = 25  # [m]
        static_friction = 1.0
        dynamic_friction = 1.0
        restitution = 0.

    class init_state:
        # * reset_mode chooses how the initial conditions are chosen.
        # * "reset_to_basic" = a single position
        # * "reset_to_range" = uniformly random from a range defined below
        reset_mode = "reset_to_basic"

        # * target state when action = 0, also reset positions for basic mode
        default_joint_angles = {"joint_a": 0.,
                                "joint_b": 0.}

        # * initial conditiosn for reset_to_range
        dof_pos_range = {'joint_a': [-1., 1.],
                         'joint_b': [-1., 1.]}
        dof_vel_range = {'joint_a': [-1., 1.],
                         'joint_b': [-1., 1.]}

    class control:
        # PD Drive parameters:
        stiffness = {'joint_a': 10.0}  # [N*m/rad]
        damping = {'joint_a': 0.5}     # [N*m*s/rad]

        actuated_joints_mask = []  # for each dof: 1 if actuated, 0 if passive

        ctrl_frequency = 100
        desired_sim_frequency = 100

    class asset:
        file = ""
        penalize_contacts_on = []
        terminate_after_contacts_on = []
        disable_gravity = False
        disable_motors = False
        # * merge bodies connected by fixed joints.
        # * Specific fixed joints can be kept by adding
        # * " <... dont_collapse="true">
        collapse_fixed_joints = True
        fix_base_link = True  # fix the base of the robot
        # * see GymDofDriveModeFlags
        # * (0 is none, 1 is pos tgt, 2 is vel tgt, 3 effort)
        default_dof_drive_mode = 3
        self_collisions = 0  # 1 to disable, 0 to enable...bitwise filter
        # * replace collision cylinders with capsules,
        # * leads to faster/more stable simulation
        replace_cylinder_with_capsule = True
        # * Some .obj meshes must be flipped from y-up to z-up
        flip_visual_attachments = True

        density = 0.001
        angular_damping = 0.
        linear_damping = 0.
        max_angular_velocity = 1000.
        max_linear_velocity = 1000.
        armature = 0.
        thickness = 0.01
        rotor_inertia = 0.
        joint_damping = 0.

    class reward_settings:
        tracking_sigma = 0.25  # tracking reward = exp(-error^2/sigma)
        # * centage of urdf limits, values above this limit are penalized
        soft_dof_pos_limit = 1.
        soft_dof_vel_limit = 1.
        soft_torque_limit = 1.  # ! may want to turn this off

    class scaling:
        dof_pos = 1.
        dof_vel = 1.

        # Action scales
        tau_ff = 10
        dof_pos_target = 0.25

    # viewer camera:
    class viewer:
        ref_env = 0
        pos = [10, 0, 6]  # [m]
        lookat = [11., 5, 3.]  # [m]

    class sim:
        dt = 0.005
        substeps = 1
        gravity = [0., 0., -9.81]  # [m/s^2]
        up_axis = 1  # 0 is y, 1 is z

        class physx:
            num_threads = 10
            solver_type = 1  # 0: pgs, 1: tgs
            num_position_iterations = 4
            num_velocity_iterations = 0
            contact_offset = 0.01  # [m]
            rest_offset = 0.0   # [m]
            bounce_threshold_velocity = 0.5  # 0.5 [m/s]
            max_depenetration_velocity = 10.0
            # * 2**24 -> needed for 8000 envs and more
            max_gpu_contact_pairs = 2**23
            default_buffer_size_multiplier = 5
            # * 0: never, 1: last sub-step, 2: all sub-steps (default=2)
            contact_collection = 2


class FixedRobotCfgPPO(BaseConfig):
    seed = -1
    runner_class_name = 'OnPolicyRunner'

    class logging():
        enable_local_saving = True

    class policy:
        init_noise_std = 1.0
        actor_hidden_dims = [512, 256, 128]
        critic_hidden_dims = [512, 256, 128]
        # * can be elu, relu, selu, crelu, lrelu, tanh, sigmoid
        activation = 'elu'
        # only for 'ActorCriticRecurrent':
        # rnn_type = 'lstm'
        # rnn_hidden_size = 512
        # rnn_num_layers = 1

        actor_obs = [
            "observation_a",
            "observation_b",
            "these_need_to_be_atributes_(states)_of_the_robot_env"]

        critic_obs = [
            "observation_x",
            "observation_y",
            "critic_obs_can_be_the_same_or_different_than_actor_obs"]

        actions = ["tau_ff"]
        disable_actions = False

        class noise:
            noise = 0.1  # implement as needed, also in your robot class

        class rewards:
            class weights:
                torques = 0.
                dof_vel = 0.
                collision = 0.
                action_rate = 0.
                action_rate2 = 0.
                dof_pos_limits = 1.

            class termination_weight:
                termination = 0.0

    class algorithm:
        # * training params
        value_loss_coef = 1.0
        use_clipped_value_loss = True
        clip_param = 0.2
        entropy_coef = 0.01
        num_learning_epochs = 5
        # * mini batch size = num_envs*nsteps / nminibatches
        num_mini_batches = 4
        learning_rate = 1.e-3
        schedule = 'adaptive'  # could be adaptive, fixed
        gamma = 0.99
        lam = 0.95
        desired_kl = 0.01
        max_grad_norm = 1.

    class runner:
        policy_class_name = 'ActorCritic'
        algorithm_class_name = 'PPO'
        num_steps_per_env = 24  # per iteration
        max_iterations = 500  # number of policy updates

        # * logging
        # * check for potential saves every this many iterations
        save_interval = 50
        run_name = ''
        experiment_name = 'fixed_robot'

        # load and resume
        resume = False
        load_run = -1  # -1 = last run
        checkpoint = -1  # -1 = last saved model
        resume_path = None  # updated from load_run and chkpt
        device = 'cpu'

gym/envs/base/fixed_robot_config.py
mit-biomimetics/ORCAgym
# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice,
# this list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
#
# 3. Neither the name of the copyright holder nor the names of its
# contributors may be used to endorse or promote products derived from
# this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF
# THE POSSIBILITY OF SUCH DAMAGE.
#
# Copyright (c) 2021 ETH Zurich, Nikita Rudin

import os
import torch
import numpy as np
from isaacgym.torch_utils import (
    get_axis_params, torch_rand_float,
    quat_rotate_inverse, to_torch, quat_from_euler_xyz)
from isaacgym import gymtorch, gymapi

from gym import LEGGED_GYM_ROOT_DIR
from gym.envs.base.base_task import BaseTask
from gym.utils.terrain import Terrain
from gym.utils.math import random_sample, quat_apply_yaw
from gym.utils.helpers import class_to_dict


class LeggedRobot(BaseTask):
    def __init__(self, gym, sim, cfg, sim_params, sim_device, headless):
        """ Parses the provided config file,
            calls create_sim() (which creates, simulation, terrain and
            environments),
            initilizes pytorch buffers used during training

        Args:
            cfg (Dict): Environment config file
            sim_params (gymapi.SimParams): simulation parameters
            physics_engine (gymapi.SimType): gymapi.SIM_PHYSX (must be PhysX)
            device_type (string): 'cuda' or 'cpu'
            device_id (int): 0, 1, ...
            headless (bool): Run without rendering if True
        """
        self.cfg = cfg
        self.sim_params = sim_params
        self.height_samples = None
        self.init_done = False
        self.device = sim_device  # todo CRIME: remove this from __init__ then
        self._parse_cfg(self.cfg)
        super().__init__(gym, sim, self.cfg, sim_params, sim_device, headless)

        if not self.headless:
            self._set_camera(self.cfg.viewer.pos, self.cfg.viewer.lookat)

        self._initialize_sim()
        self.gym.prepare_sim(self.sim)
        self._init_buffers()
        self.init_done = True
        self.reset()

    def step(self):
        """ Apply actions, simulate, call self._post_physics_step()

        Args:
            actions (torch.Tensor): Tensor of shape (num_envs,
            num_actions_per_env)
        """
        self._reset_buffers()
        self._pre_physics_step()
        # * step physics and render each frame
        self._render()
        for _ in range(self.cfg.control.decimation):
            self._pre_torque_step()
            self.torques = self._compute_torques()
            self._post_torque_step()
            self._step_physx_sim()

        self._post_physics_step()
        self._check_terminations_and_timeouts()

        env_ids = self.to_be_reset.nonzero().flatten()
        self._reset_idx(env_ids)

    def _pre_physics_step(self):
        return None

    def _pre_torque_step(self):
        return None

    def _post_torque_step(self):
        return None

    def _step_physx_sim(self):
        self.gym.set_dof_actuation_force_tensor(
                self.sim, gymtorch.unwrap_tensor(self.torques))
        self.gym.simulate(self.sim)
        if self.device == 'cpu':
            self.gym.fetch_results(self.sim, True)
        self.gym.refresh_dof_state_tensor(self.sim)

    def _post_physics_step(self):
        self.gym.refresh_actor_root_state_tensor(self.sim)
        self.gym.refresh_net_contact_force_tensor(self.sim)

        self.episode_length_buf += 1
        self.common_step_counter += 1

        if self.cfg.terrain.measure_heights:
            self.measured_heights = self._get_heights()

        # prepare quantities
        self.base_quat[:] = self.root_states[:, 3:7]
        self.base_lin_vel[:] = quat_rotate_inverse(
            self.base_quat, self.root_states[:, 7:10])
        self.base_ang_vel[:] = quat_rotate_inverse(
            self.base_quat, self.root_states[:, 10:13])
        self.projected_gravity[:] = quat_rotate_inverse(
            self.base_quat, self.gravity_vec)

        self.base_height = self.root_states[:, 2:3]

        nact = self.num_actuators
        self.dof_pos_history[:, 2*nact:] = self.dof_pos_history[:, nact:2*nact]
        self.dof_pos_history[:, nact:2*nact] = self.dof_pos_history[:, :nact]
        self.dof_pos_history[:, :nact] = self.dof_pos_target
        self.dof_pos_obs = (self.dof_pos - self.default_dof_pos)

        env_ids = (self.episode_length_buf
                   % int(self.cfg.commands.resampling_time / self.dt) == 0)
        self._resample_commands(env_ids.nonzero().flatten())
        if self.cfg.push_robots.toggle:
            if (self.common_step_counter % self.cfg.push_interval == 0):
                self._push_robots()

    def _reset_idx(self, env_ids):
        if len(env_ids) == 0:
            return

        # * reset robot states
        self._reset_system(env_ids)
        self._resample_commands(env_ids)
        # * reset buffers
        self.dof_pos_history[env_ids] = 0.
        self.episode_length_buf[env_ids] = 0

    def _initialize_sim(self):
        """ Creates simulation, terrain and evironments
        """
        self.up_axis_idx = 2  # 2 for z, 1 for y -> adapt gravity accordingly
        mesh_type = self.cfg.terrain.mesh_type
        if mesh_type in ['heightfield', 'trimesh']:
            self.terrain = Terrain(self.cfg.terrain, self.num_envs)
        if mesh_type == 'plane':
            self._create_ground_plane()
        elif mesh_type == 'heightfield':
            self._create_heightfield()
        elif mesh_type == 'trimesh':
            self._create_trimesh()
        elif mesh_type is not None:
            raise ValueError(
                "Terrain mesh type not recognised. Allowed types are "
                + "[None, plane, heightfield, trimesh]")
        self._create_envs()

    def _resample_commands(self, env_ids):
        """ Randommly select commands of some environments

        Args:
            env_ids (List[int]):
            Environments ids for which new commands are needed
        """
        if len(env_ids) == 0:
            return

        self.commands[env_ids, 0] = torch_rand_float(
            self.command_ranges["lin_vel_x"][0],
            self.command_ranges["lin_vel_x"][1],
            (len(env_ids), 1),
            device=self.device).squeeze(1)
        self.commands[env_ids, 1] = torch_rand_float(
            -self.command_ranges["lin_vel_y"],
            self.command_ranges["lin_vel_y"],
            (len(env_ids), 1),
            device=self.device).squeeze(1)
        max_yaw_vel = self.command_ranges["yaw_vel"]
        self.commands[env_ids, 2] = torch_rand_float(
            -max_yaw_vel, max_yaw_vel, (len(env_ids), 1),
            device=self.device).squeeze(1)

        # set small commands to zero
        self.commands[env_ids, :2] *= (
            torch.norm(self.commands[env_ids, :3], dim=1) > 0.2).unsqueeze(1)

    def _set_camera(self, position, lookat):
        """ Set camera position and direction
        """
        cam_pos = gymapi.Vec3(position[0], position[1], position[2])
        cam_target = gymapi.Vec3(lookat[0], lookat[1], lookat[2])
        self.gym.viewer_camera_look_at(self.viewer, None, cam_pos, cam_target)

    def _process_rigid_shape_props(self, props, env_id):
        """ Callback allowing to store/change/randomize the rigid shape
            properties of each environment. Called During environment creation.
            Base behavior: randomizes the friction of each environment

        Args:
            props (List[gymapi.RigidShapeProperties]):
                Properties of each shape of the asset
            env_id (int): Environment id

        Returns:
            [List[gymapi.RigidShapeProperties]]:
                Modified rigid shape properties
        """
        if self.cfg.domain_rand.randomize_friction:
            if env_id == 0:
                # prepare friction randomization
                friction_range = self.cfg.domain_rand.friction_range
                num_buckets = 64
                bucket_ids = torch.randint(0, num_buckets, (self.num_envs, 1))
                friction_buckets = torch_rand_float(friction_range[0],
                                                    friction_range[1],
                                                    (num_buckets, 1),
                                                    device='cpu')
                self.friction_coeffs = friction_buckets[bucket_ids]

            for s in range(len(props)):
                props[s].friction = self.friction_coeffs[env_id]
        return props

    def _process_dof_props(self, props, env_id):
        """ Callback allowing to store/change/randomize the DOF properties of
            each environment. Called During environment creation.
            Base behavior: stores position, velocity and torques limits
                defined in the URDF

        Args:
            props (numpy.array): Properties of each DOF of the asset
            env_id (int): Environment id

        Returns:
            [numpy.array]: Modified DOF properties
        """
        if env_id == 0:
            self.dof_pos_limits = torch.zeros(self.num_dof, 2,
                                              dtype=torch.float,
                                              device=self.device)
            self.dof_vel_limits = torch.zeros(self.num_dof,
                                              dtype=torch.float,
                                              device=self.device)
            self.torque_limits = torch.zeros(self.num_dof,
                                             dtype=torch.float,
                                             device=self.device)

            for i in range(len(props)):
                self.dof_pos_limits[i, 0] = props["lower"][i].item()
                self.dof_pos_limits[i, 1] = props["upper"][i].item()
                self.dof_vel_limits[i] = props["velocity"][i].item()
                self.torque_limits[i] = props["effort"][i].item()
                # * soft limits
                m = (self.dof_pos_limits[i, 0] + self.dof_pos_limits[i, 1]) / 2
                r = self.dof_pos_limits[i, 1] - self.dof_pos_limits[i, 0]
                self.dof_pos_limits[i, 0] = \
                    m - 0.5 * r * self.cfg.reward_settings.soft_dof_pos_limit
                self.dof_pos_limits[i, 1] = \
                    m + 0.5 * r * self.cfg.reward_settings.soft_dof_pos_limit
        return props

    def _process_rigid_body_props(self, props, env_id):
        # if env_id==0:
        #     sum = 0
        #     for i, p in enumerate(props):
        #         sum += p.mass
        #         print(f"Mass of body {i}: {p.mass} (before randomization)")
        #     print(f"Total mass {sum} (before randomization)")
        # randomize base mass
        if self.cfg.domain_rand.randomize_base_mass:
            rng = self.cfg.domain_rand.added_mass_range
            props[0].mass += np.random.uniform(rng[0], rng[1])
        return props

    def _compute_torques(self):
        if self.cfg.asset.disable_motors:
            self.torques[:] = 0.
            return
        torques = (self.p_gains*(self.dof_pos_target
                                 + self.default_dof_pos
                                 - self.dof_pos)
                   + self.d_gains*(self.dof_vel_target - self.dof_vel)
                   + self.tau_ff)
        torques = torch.clip(torques, -self.torque_limits, self.torque_limits)
        return torques.view(self.torques.shape)

    def _reset_system(self, env_ids):
        """ Resets selected environmments
        Args:
            env_ids (List[int]): Environemnt ids
        """
        if hasattr(self, self.cfg.init_state.reset_mode):
            eval(f"self.{self.cfg.init_state.reset_mode}(env_ids)")
        else:
            raise NameError(
                f"Unknown default setup: {self.cfg.init_state.reset_mode}")

        env_ids_int32 = env_ids.to(dtype=torch.int32)
        self.gym.set_dof_state_tensor_indexed(
            self.sim,
            gymtorch.unwrap_tensor(self.dof_state),
            gymtorch.unwrap_tensor(env_ids_int32),
            len(env_ids_int32))

        # * start base position shifted in X-Y plane
        if self.custom_origins:
            self.root_states[env_ids, :3] += self.env_origins[env_ids]
            # * xy position within 1m of the center
            self.root_states[env_ids, :2] += \
                torch_rand_float(
                    -1., 1., (len(env_ids), 2), device=self.device)
        else:
            self.root_states[env_ids, :3] += self.env_origins[env_ids]

        self.gym.set_actor_root_state_tensor_indexed(
            self.sim,
            gymtorch.unwrap_tensor(self.root_states),
            gymtorch.unwrap_tensor(env_ids_int32),
            len(env_ids_int32))

    # * implement reset methods
    def reset_to_basic(self, env_ids):
        """
        Reset to a single initial state
        """
        self.dof_pos[env_ids] = self.default_dof_pos
        self.dof_vel[env_ids] = 0
        self.root_states[env_ids] = self.base_init_state

    def reset_to_range(self, env_ids):
        """
        Reset to a uniformly random distribution of states, sampled from a
        range for each state
        """
        # * dof states
        self.dof_pos[env_ids] = random_sample(
            env_ids, self.dof_pos_range[:, 0],
            self.dof_pos_range[:, 1], device=self.device)
        self.dof_vel[env_ids] = random_sample(
            env_ids, self.dof_vel_range[:, 0],
            self.dof_vel_range[:, 1], device=self.device)

        # * base states
        random_com_pos = random_sample(
            env_ids, self.root_pos_range[:, 0],
            self.root_pos_range[:, 1], device=self.device)

        self.root_states[env_ids, 0:7] = torch.cat(
            (random_com_pos[:, 0:3],
             quat_from_euler_xyz(
                random_com_pos[:, 3],
                random_com_pos[:, 4],
                random_com_pos[:, 5])), 1)
        self.root_states[env_ids, 7:13] = random_sample(
            env_ids, self.root_vel_range[:, 0],
            self.root_vel_range[:, 1], device=self.device)

    def _push_robots(self):
        """ Random pushes the robots. Emulates an impulse by setting a
            randomized base velocity.
        """
        max_vel = self.cfg.push_robots.max_push_vel_xy
        self.root_states[:, 7:9] += torch_rand_float(-max_vel, max_vel,
                                                     (self.num_envs, 2),
                                                     device=self.device)
        self.gym.set_actor_root_state_tensor(
            self.sim, gymtorch.unwrap_tensor(self.root_states))

    # ----------------------------------------
    def _init_buffers(self):
        """ Initialize torch tensors which will contain simulation states and
            processed quantities
        """
        # * get gym GPU state tensors
        actor_root_state = \
            self.gym.acquire_actor_root_state_tensor(self.sim)
        dof_state_tensor = \
            self.gym.acquire_dof_state_tensor(self.sim)
        net_contact_forces = \
            self.gym.acquire_net_contact_force_tensor(self.sim)
        rigid_body_state = \
            self.gym.acquire_rigid_body_state_tensor(self.sim)

        self.gym.refresh_dof_state_tensor(self.sim)
        self.gym.refresh_actor_root_state_tensor(self.sim)
        self.gym.refresh_net_contact_force_tensor(self.sim)
        self.gym.refresh_rigid_body_state_tensor(self.sim)

        # * create some wrapper tensors for different slices

        self.root_states = gymtorch.wrap_tensor(actor_root_state)
        self.dof_state = gymtorch.wrap_tensor(dof_state_tensor)
        self._rigid_body_state = gymtorch.wrap_tensor(rigid_body_state)

        self._rigid_body_pos = self._rigid_body_state.view(
            self.num_envs, self.num_bodies, 13)[..., 0:3]
        self.dof_pos = self.dof_state.view(
            self.num_envs, self.num_dof, 2)[..., 0]
        self.dof_vel = self.dof_state.view(
            self.num_envs, self.num_dof, 2)[..., 1]
        self.base_quat = self.root_states[:, 3:7]

        # * shape: num_envs, num_bodies, xyz axis
        self.contact_forces = (
            gymtorch.wrap_tensor(net_contact_forces)
            .view(self.num_envs, -1, 3))

        self._rigid_body_pos = self._rigid_body_state.view(
            self.num_envs, self.num_bodies, 13)[..., 0:3]
        self._rigid_body_quat = self._rigid_body_state.view(
            self.num_envs, self.num_bodies, 13)[..., 3:7]
        self._rigid_body_lin_vel = self._rigid_body_state.view(
            self.num_envs, self.num_bodies, 13)[..., 7:10]
        self._rigid_body_ang_vel = self._rigid_body_state.view(
            self.num_envs, self.num_bodies, 13)[..., 10:13]

        # * initialize some data used later on
        self.common_step_counter = 0
        self.extras = {}

        self.gravity_vec = to_torch(
            get_axis_params(-1., self.up_axis_idx),
            device=self.device).repeat((self.num_envs, 1))
        self.torques = torch.zeros(self.num_envs, self.num_actuators,
                                   dtype=torch.float, device=self.device)
        self.p_gains = torch.zeros(self.num_actuators,
                                   dtype=torch.float, device=self.device)
        self.d_gains = torch.zeros(self.num_actuators,
                                   dtype=torch.float, device=self.device)
        self.dof_pos_target = torch.zeros(self.num_envs, self.num_actuators,
                                          dtype=torch.float,
                                          device=self.device)
        self.dof_vel_target = torch.zeros(self.num_envs, self.num_actuators,
                                          dtype=torch.float,
                                          device=self.device)
        self.tau_ff = torch.zeros(self.num_envs, self.num_actuators,
                                  dtype=torch.float, device=self.device)

        self.dof_pos_history = torch.zeros(
            self.num_envs, self.num_actuators*3,
            dtype=torch.float, device=self.device)
        self.commands = torch.zeros(
            self.num_envs, 3,
            dtype=torch.float, device=self.device)
        self.base_lin_vel = quat_rotate_inverse(self.base_quat,
                                                self.root_states[:, 7:10])
        self.base_ang_vel = quat_rotate_inverse(self.base_quat,
                                                self.root_states[:, 10:13])
        self.projected_gravity = quat_rotate_inverse(self.base_quat,
                                                     self.gravity_vec)
        self.dof_pos_obs = torch.zeros_like(self.dof_pos)
        self.base_height = torch.zeros(self.num_envs, 1,
                                       dtype=torch.float, device=self.device)

        if self.cfg.terrain.measure_heights:
            self.height_points = self._init_height_points()
        self.measured_heights = 0

        # * joint positions offsets and PD gains
        self.default_dof_pos = torch.zeros(
            self.num_dof,
            dtype=torch.float, device=self.device)
        for i in range(self.num_dof):
            name = self.dof_names[i]
            angles = self.cfg.init_state.default_joint_angles

            found = False
            for dof_name in angles.keys():
                if dof_name in name:
                    self.default_dof_pos[i] = angles[dof_name]
                    found = True
            if not found:
                self.default_dof_pos[i] = 0.0
                print(f"Default dof pos of joint {name} was not defined, "
                      + "setting to zero")

            found = False
            for dof_name in self.cfg.control.stiffness.keys():
                if dof_name in name:
                    self.p_gains[i] = self.cfg.control.stiffness[dof_name]
                    self.d_gains[i] = self.cfg.control.damping[dof_name]
                    found = True
            if not found:
                self.p_gains[i] = 0.
                self.d_gains[i] = 0.
                if self.cfg.control.control_type in ["P", "V"]:
                    print(f"PD gain of joint {name} were not defined, "
                          + "setting them to zero")
        self.default_dof_pos = self.default_dof_pos.unsqueeze(0)

        # * check that init range highs and lows are consistent
        # * and repopulate to match
        if self.cfg.init_state.reset_mode == "reset_to_range":
            self.initialize_ranges_for_initial_conditions()

    def initialize_ranges_for_initial_conditions(self):
        self.dof_pos_range = torch.zeros(
            self.num_dof, 2,
            dtype=torch.float, device=self.device)
        self.dof_vel_range = torch.zeros(
            self.num_dof, 2,
            dtype=torch.float, device=self.device)

        for joint, vals in self.cfg.init_state.dof_pos_range.items():
            for i in range(self.num_dof):
                if joint in self.dof_names[i]:
                    self.dof_pos_range[i, :] = to_torch(
                        vals, device=self.device)

        for joint, vals in self.cfg.init_state.dof_vel_range.items():
            for i in range(self.num_dof):
                if joint in self.dof_names[i]:
                    self.dof_vel_range[i, :] = to_torch(
                        vals, device=self.device)

        self.root_pos_range = torch.tensor(
            self.cfg.init_state.root_pos_range,
            dtype=torch.float, device=self.device)
        self.root_vel_range = torch.tensor(
            self.cfg.init_state.root_vel_range,
            dtype=torch.float, device=self.device)

    def _create_ground_plane(self):
        """ Adds a ground plane to the simulation, sets friction and
            restitution based on the cfg.
        """
        plane_params = gymapi.PlaneParams()
        plane_params.normal = gymapi.Vec3(0.0, 0.0, 1.0)
        plane_params.static_friction = self.cfg.terrain.static_friction
        plane_params.dynamic_friction = self.cfg.terrain.dynamic_friction
        plane_params.restitution = self.cfg.terrain.restitution
        self.gym.add_ground(self.sim, plane_params)

    def _create_heightfield(self):
        """ Adds a heightfield terrain to the simulation, sets parameters
            based on the cfg.
        """
        hf_params = gymapi.HeightFieldProperties()
        hf_params.column_scale = self.terrain.horizontal_scale
        hf_params.row_scale = self.terrain.horizontal_scale
        hf_params.vertical_scale = self.terrain.vertical_scale
        hf_params.nbRows = self.terrain.tot_cols
        hf_params.nbColumns = self.terrain.tot_rows
        hf_params.transform.p.x = -self.terrain.border_size
        hf_params.transform.p.y = -self.terrain.border_size
        hf_params.transform.p.z = 0.0
        hf_params.static_friction = self.cfg.terrain.static_friction
        hf_params.dynamic_friction = self.cfg.terrain.dynamic_friction
        hf_params.restitution = self.cfg.terrain.restitution

        self.gym.add_heightfield(
            self.sim, self.terrain.heightsamples, hf_params)
        self.height_samples = torch.tensor(
            self.terrain.heightsamples).view(
                self.terrain.tot_rows, self.terrain.tot_cols).to(self.device)

    def _create_trimesh(self):
        """ Adds a triangle mesh terrain to the simulation, sets parameters
            based on the cfg.
        # """
        tm_params = gymapi.TriangleMeshParams()
        tm_params.nb_vertices = self.terrain.vertices.shape[0]
        tm_params.nb_triangles = self.terrain.triangles.shape[0]

        tm_params.transform.p.x = -self.terrain.cfg.border_size
        tm_params.transform.p.y = -self.terrain.cfg.border_size
        tm_params.transform.p.z = 0.0
        tm_params.static_friction = self.cfg.terrain.static_friction
        tm_params.dynamic_friction = self.cfg.terrain.dynamic_friction
        tm_params.restitution = self.cfg.terrain.restitution
        self.gym.add_triangle_mesh(
            self.sim, self.terrain.vertices.flatten(order='C'),
            self.terrain.triangles.flatten(order='C'), tm_params)
        self.height_samples = torch.tensor(
            self.terrain.heightsamples).view(
                self.terrain.tot_rows, self.terrain.tot_cols).to(self.device)

    def _create_envs(self):
        """ Creates environments:
             1. loads the robot URDF/MJCF asset,
             2. For each environment
                2.1 creates the environment,
                2.2 calls DOF and Rigid shape properties callbacks,
                2.3 create actor with these properties and add them to the env
             3. Store indices of different bodies of the robot
        """
        asset_path = self.cfg.asset.file.format(
            LEGGED_GYM_ROOT_DIR=LEGGED_GYM_ROOT_DIR)
        asset_root = os.path.dirname(asset_path)
        asset_file = os.path.basename(asset_path)

        asset_options = gymapi.AssetOptions()
        asset_options.fix_base_link = self.cfg.asset.fix_base_link
        asset_options.density = self.cfg.asset.density
        asset_options.angular_damping = self.cfg.asset.angular_damping
        asset_options.linear_damping = self.cfg.asset.linear_damping
        asset_options.max_linear_velocity = self.cfg.asset.max_linear_velocity
        asset_options.armature = self.cfg.asset.armature
        asset_options.thickness = self.cfg.asset.thickness
        asset_options.disable_gravity = self.cfg.asset.disable_gravity
        asset_options.default_dof_drive_mode = \
            self.cfg.asset.default_dof_drive_mode
        asset_options.collapse_fixed_joints = \
            self.cfg.asset.collapse_fixed_joints
        asset_options.replace_cylinder_with_capsule = \
            self.cfg.asset.replace_cylinder_with_capsule
        asset_options.flip_visual_attachments = \
            self.cfg.asset.flip_visual_attachments
        asset_options.max_angular_velocity = \
            self.cfg.asset.max_angular_velocity

        robot_asset = self.gym.load_asset(
            self.sim, asset_root, asset_file, asset_options)
        self.num_dof = self.gym.get_asset_dof_count(robot_asset)
        self.num_bodies = self.gym.get_asset_rigid_body_count(robot_asset)

        dof_props_asset = self.gym.get_asset_dof_properties(robot_asset)
        dof_props_asset['armature'] = self.cfg.asset.rotor_inertia
        dof_props_asset['damping'] = self.cfg.asset.joint_damping
        rigid_shape_props_asset = \
            self.gym.get_asset_rigid_shape_properties(robot_asset)

        body_names = self.gym.get_asset_rigid_body_names(robot_asset)
        self.dof_names = self.gym.get_asset_dof_names(robot_asset)
        self.num_bodies = len(body_names)
        feet_names = [s for s in body_names if self.cfg.asset.foot_name in s]
        penalized_contact_names = []
        for name in self.cfg.asset.penalize_contacts_on:
            penalized_contact_names.extend(
                [s for s in body_names if name in s])
        termination_contact_names = []
        for name in self.cfg.asset.terminate_after_contacts_on:
            termination_contact_names.extend(
                [s for s in body_names if name in s])

        base_init_state_list = (
            self.cfg.init_state.pos + self.cfg.init_state.rot
            + self.cfg.init_state.lin_vel + self.cfg.init_state.ang_vel)
        self.base_init_state = (
            to_torch(base_init_state_list,
                     device=self.device))
        start_pose = gymapi.Transform()
        start_pose.p = gymapi.Vec3(*self.base_init_state[:3])

        self._get_env_origins()
        env_lower = gymapi.Vec3(0., 0., 0.)
        env_upper = gymapi.Vec3(0., 0., 0.)
        self.actor_handles = []
        self.envs = []

        for i in range(self.num_envs):
            # * create env instance
            env_handle = self.gym.create_env(self.sim, env_lower, env_upper,
                                             int(np.sqrt(self.num_envs)))
            pos = self.env_origins[i].clone()
            pos[:2] += torch_rand_float(-1., 1., (2, 1),
                                        device=self.device).squeeze(1)
            start_pose.p = gymapi.Vec3(*pos)

            rigid_shape_props = \
                self._process_rigid_shape_props(rigid_shape_props_asset, i)
            self.gym.set_asset_rigid_shape_properties(robot_asset,
                                                      rigid_shape_props)

            name = self.cfg.asset.file
            name = name.format(LEGGED_GYM_ROOT_DIR=LEGGED_GYM_ROOT_DIR)
            robot_handle = self.gym.create_actor(
                env_handle, robot_asset, start_pose, name, i,
                self.cfg.asset.self_collisions, 0)
            dof_props = self._process_dof_props(dof_props_asset, i)
            self.gym.set_actor_dof_properties(
                env_handle, robot_handle, dof_props)
            body_props = self.gym.get_actor_rigid_body_properties(
                env_handle, robot_handle)
            body_props = self._process_rigid_body_props(body_props, i)
            self.gym.set_actor_rigid_body_properties(
                env_handle, robot_handle, body_props, recomputeInertia=True)
            self.envs.append(env_handle)
            self.actor_handles.append(robot_handle)
        self.feet_indices = torch.zeros(
            len(feet_names),
            dtype=torch.long, device=self.device)
        for i in range(len(feet_names)):
            self.feet_indices[i] = \
                self.gym.find_actor_rigid_body_handle(
                    self.envs[0], self.actor_handles[0], feet_names[i])

        self.penalised_contact_indices = torch.zeros(
            len(penalized_contact_names),
            dtype=torch.long, device=self.device)
        for i in range(len(penalized_contact_names)):
            self.penalised_contact_indices[i] = \
                self.gym.find_actor_rigid_body_handle(
                    self.envs[0], self.actor_handles[0],
                    penalized_contact_names[i])

        self.termination_contact_indices = torch.zeros(
            len(termination_contact_names),
            dtype=torch.long, device=self.device)
        for i in range(len(termination_contact_names)):
            self.termination_contact_indices[i] = \
                self.gym.find_actor_rigid_body_handle(
                    self.envs[0], self.actor_handles[0],
                    termination_contact_names[i])

    def _get_env_origins(self):
        """ Sets environment origins. On rough terrain the origins are defined
            by the terrain platforms. Otherwise create a grid.
        """
        if self.cfg.terrain.mesh_type in ["heightfield", "trimesh"]:
            self.custom_origins = True
            self.env_origins = torch.zeros(
                self.num_envs, 3,
                device=self.device)
            # * put robots at the origins defined by the terrain
            max_init_level = self.cfg.terrain.num_rows - 1
            self.terrain_levels = torch.randint(
                0, max_init_level+1, (self.num_envs,), device=self.device)
            self.terrain_types = torch.div(
                torch.arange(self.num_envs, device=self.device),
                (self.num_envs/self.cfg.terrain.num_cols),
                rounding_mode='floor').to(torch.long)
            self.max_terrain_level = self.cfg.terrain.num_rows
            self.terrain_origins = torch.from_numpy(
                self.terrain.env_origins).to(self.device).to(torch.float)
            self.env_origins[:] = \
                self.terrain_origins[self.terrain_levels, self.terrain_types]
        else:
            self.custom_origins = False
            self.env_origins = torch.zeros(
                self.num_envs, 3,
                device=self.device)
            # * create a grid of robots
            num_cols = np.floor(np.sqrt(self.num_envs))
            num_rows = np.ceil(self.num_envs / num_cols)
            xx, yy = torch.meshgrid(
                torch.arange(num_rows), torch.arange(num_cols))
            spacing = self.cfg.env.env_spacing
            self.env_origins[:, 0] = spacing * xx.flatten()[:self.num_envs]
            self.env_origins[:, 1] = spacing * yy.flatten()[:self.num_envs]
            self.env_origins[:, 2] = 0.

    def _parse_cfg(self, cfg):
        self.dt = self.cfg.control.ctrl_dt
        self.scales = class_to_dict(self.cfg.scaling, self.device)
        self.command_ranges = class_to_dict(self.cfg.commands.ranges)
        self.max_episode_length_s = self.cfg.env.episode_length_s
        self.max_episode_length = int(np.ceil(
            self.max_episode_length_s / self.dt))
        self.cfg.push_interval = np.ceil(
            self.cfg.push_robots.interval_s / self.dt)

    def _init_height_points(self):
        """ Returns points at which the height measurments are sampled
            (in base frame)

        Returns:
            [torch.Tensor]: Tensor of shape
                (num_envs, self.num_height_points, 3)
        """
        y = torch.tensor(
            self.cfg.terrain.measured_points_y,
            device=self.device)
        x = torch.tensor(
            self.cfg.terrain.measured_points_x,
            device=self.device)
        grid_x, grid_y = torch.meshgrid(x, y)

        self.num_height_points = grid_x.numel()
        points = torch.zeros(
            self.num_envs, self.num_height_points, 3,
            device=self.device)
        points[:, :, 0] = grid_x.flatten()
        points[:, :, 1] = grid_y.flatten()
        return points

    def _get_heights(self, env_ids=None):
        """ Samples heights of the terrain at required points around each
            robot. The points are offset by the base's position and rotated
            by the base's yaw

        Args:
            env_ids (List[int], optional): Subset of environments for which to
                return the heights. Defaults to None.

        Raises:
            NameError: [description]

        Returns:
            [type]: [description]
        """
        if self.cfg.terrain.mesh_type == 'plane':
            return torch.zeros(
                self.num_envs, self.num_height_points,
                device=self.device)
        elif self.cfg.terrain.mesh_type == 'none':
            raise NameError(
                "Can't measure height with terrain mesh type 'none'")

        if env_ids:
            points = quat_apply_yaw(
                self.base_quat[env_ids].repeat(1, self.num_height_points),
                (self.height_points[env_ids])
                + (self.root_states[env_ids, :3]).unsqueeze(1))
        else:
            points = quat_apply_yaw(
                self.base_quat.repeat(1, self.num_height_points),
                self.height_points) + (self.root_states[:, :3]).unsqueeze(1)

        points += self.terrain.cfg.border_size
        points = (points/self.terrain.cfg.horizontal_scale).long()
        px = points[:, :, 0].view(-1)
        py = points[:, :, 1].view(-1)
        px = torch.clip(px, 0, self.height_samples.shape[0]-2)
        py = torch.clip(py, 0, self.height_samples.shape[1]-2)

        heights1 = self.height_samples[px, py]
        heights2 = self.height_samples[px+1, py]
        heights3 = self.height_samples[px, py+1]
        heights = torch.min(heights1, heights2)
        heights = torch.min(heights, heights3)

        return heights.view(self.num_envs, -1) \
            * self.terrain.cfg.vertical_scale

    def _sqrdexp(self, x, scale=None):
        """ shorthand helper for squared exponential
        """
        if scale is None:
            scale = 1.
        return torch.exp(
            -torch.square(x/scale)/self.cfg.reward_settings.tracking_sigma)

    # ------------ reward functions----------------

    def _reward_lin_vel_z(self):
        """Penalize z axis base linear velocity"""
        return -torch.square(self.base_lin_vel[:, 2])

    def _reward_ang_vel_xy(self):
        """Penalize xy axes base angular velocity"""
        return -torch.sum(torch.square(self.base_ang_vel[:, :2]), dim=1)

    def _reward_orientation(self):
        """Penalize non flat base orientation"""
        return -torch.sum(torch.square(self.projected_gravity[:, :2]), dim=1)

    def _reward_base_height(self):
        """Penalize base height away from target"""
        base_height = torch.mean(
            self.root_states[:, 2].unsqueeze(1) - self.measured_heights, dim=1)
        return -torch.square(
            base_height - self.cfg.reward_settings.base_height_target)

    def _reward_torques(self):
        """Penalize torques"""
        return -torch.sum(torch.square(self.torques), dim=1)

    def _reward_dof_vel(self):
        """Penalize dof velocities"""
        return -torch.sum(torch.square(self.dof_vel), dim=1)

    def _reward_action_rate(self):
        """Penalize changes in actions"""
        nact = self.num_actuators
        dt2 = (self.dt*self.cfg.control.decimation)**2
        error = torch.square(self.dof_pos_history[:, :nact]
                             - self.dof_pos_history[:, nact:2*nact])/dt2
        return -torch.sum(error, dim=1)

    def _reward_action_rate2(self):
        """Penalize changes in actions"""
        nact = self.num_actuators
        dt2 = (self.dt*self.cfg.control.decimation)**2
        error = torch.square(self.dof_pos_history[:, :nact]
                             - 2*self.dof_pos_history[:, nact:2*nact]
                             + self.dof_pos_history[:, 2*nact:])/dt2
        return -torch.sum(error, dim=1)

    def _reward_collision(self):
        """Penalize collisions on selected bodies"""
        return -torch.sum(1.*(torch.norm(
            self.contact_forces[:, self.penalised_contact_indices, :],
            dim=-1) > 0.1), dim=1)

    def _reward_termination(self):
        # Terminal reward / penalty
        return -self.terminated.float()

    def _reward_dof_pos_limits(self):
        """Penalize dof positions too close to the limit"""
        # * lower limit
        out_of_limits = -(
            self.dof_pos - self.dof_pos_limits[:, 0]).clip(max=0.)
        out_of_limits += (
            self.dof_pos - self.dof_pos_limits[:, 1]).clip(min=0.)
        return -torch.sum(out_of_limits, dim=1)

    def _reward_dof_vel_limits(self):
        """Penalize dof velocities too close to the limit"""
        # * clip to max error = 1 rad/s per joint to avoid huge penalties
        soft_dof_vel_limit = self.cfg.reward_settings.soft_dof_vel_limit
        return -torch.sum((
            torch.abs(self.dof_vel)
            - (self.dof_vel_limits * soft_dof_vel_limit))
                         .clip(min=0., max=1.), dim=1)

    def _reward_torque_limits(self):
        """penalize torques too close to the limit"""
        return -torch.sum((
            torch.abs(self.torques)
            - self.torque_limits*self.cfg.reward_settings.soft_torque_limit)
                         .clip(min=0.), dim=1)

    def _reward_tracking_lin_vel(self):
        """Tracking of linear velocity commands (xy axes)"""
        error = torch.square(self.commands[:, :2] - self.base_lin_vel[:, :2])
        error = torch.exp(-error/self.cfg.reward_settings.tracking_sigma)
        return torch.sum(error, dim=1)

    def _reward_tracking_ang_vel(self):
        """Tracking of angular velocity commands (yaw)"""
        ang_vel_error = torch.square(
            self.commands[:, 2] - self.base_ang_vel[:, 2])
        return self._sqrdexp(ang_vel_error/self.scales['base_ang_vel'][2])

    def _reward_feet_contact_forces(self):
        """penalize high contact forces"""
        return -torch.sum((
            torch.norm(self.contact_forces[:, self.feet_indices, :], dim=-1)
            - self.cfg.reward_settings.max_contact_force).clip(min=0.), dim=1)

gym/envs/base/legged_robot.py
mit-biomimetics/ORCAgym
# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice,
# this list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
#
# 3. Neither the name of the copyright holder nor the names of its
# contributors may be used to endorse or promote products derived from
# this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF
# THE POSSIBILITY OF SUCH DAMAGE.
#
# Copyright (c) 2021 ETH Zurich, Nikita Rudin

from .base_config import BaseConfig


class LeggedRobotCfg(BaseConfig):
    class env:
        num_envs = 4096
        num_actuators = 12
        env_spacing = 0.  # not used with heightfields/trimeshes
        episode_length_s = 20  # episode length in seconds

    class terrain:
        mesh_type = 'plane'  # none, plane, heightfield or trimesh
        horizontal_scale = 0.1  # [m]
        vertical_scale = 0.005  # [m]
        border_size = 25  # [m]
        curriculum = False
        static_friction = 1.0
        dynamic_friction = 1.0
        restitution = 0.
        # * rough terrain only:
        measure_heights = False
        # * 1mx1.6m rectangle (without center line)
        measured_points_x = [
            -0.8, -0.7, -0.6, -0.5, -0.4, -0.3, -0.2, -0.1,
            0., 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8]
        measured_points_y = [
            -0.5, -0.4, -0.3, -0.2, -0.1, 0., 0.1, 0.2, 0.3, 0.4, 0.5]
        selected = False  # select a unique terrain type and pass all arguments
        terrain_kwargs = None  # Dict of arguments for selected terrain
        max_init_terrain_level = 2  # starting curriculum state
        terrain_length = 8.
        terrain_width = 8.
        num_rows = 10  # number of terrain rows (levels)
        num_cols = 20  # number of terrain cols (types)
        # * terrain types:
        # * [smooth slope, rough slope, stairs up, stairs down, discrete]
        terrain_proportions = [0.6, 0.4, 0., 0., 0.]
        # * trimesh only:
        slope_treshold = 0.75
        # * slopes above this threshold will be corrected to vertical surfaces

    class commands:
        resampling_time = 10.  # time before command are changed[s]

        class ranges:
            lin_vel_x = [-1.0, 1.0]  # min max [m/s]
            lin_vel_y = 1.  # max [m/s]
            yaw_vel = 1    # max [rad/s]

    class push_robots:
        toggle = True
        interval_s = 15
        max_push_vel_xy = 0.05

    class init_state:

        # * target state when action = 0, also reset positions for basic mode
        default_joint_angles = {"joint_a": 0.,
                                "joint_b": 0.}

        # * reset setup chooses how the initial conditions are chosen.
        # * "reset_to_basic" = a single position
        # * "reset_to_range" = uniformly random from a range defined below
        reset_mode = "reset_to_basic"

        # * default COM for basic initialization
        pos = [0.0, 0.0, 1.]        # x,y,z [m]
        rot = [0.0, 0.0, 0.0, 1.0]  # x,y,z,w [quat]
        lin_vel = [0.0, 0.0, 0.0]   # x,y,z [m/s]
        ang_vel = [0.0, 0.0, 0.0]   # x,y,z [rad/s]

        # * initial conditiosn for reset_to_range
        dof_pos_range = {'joint_a': [-1., 1.],
                         'joint_b': [-1., 1.]}
        dof_vel_range = {'joint_a': [-1., 1.],
                         'joint_b': [-1., 1.]}
        root_pos_range = [[0., 0.],     # x
                          [0., 0.],     # y
                          [0.5, 0.75],  # z
                          [0., 0.],     # roll
                          [0., 0.],     # pitch
                          [0., 0.]]     # yaw
        root_vel_range = [[-0.1, 0.1],  # x
                          [-0.1, 0.1],  # y
                          [-0.1, 0.1],  # z
                          [-0.1, 0.1],  # roll
                          [-0.1, 0.1],  # pitch
                          [-0.1, 0.1]]  # yaw

    class control:
        # * PD Drive parameters:
        stiffness = {'joint_a': 10.0, 'joint_b': 15.}  # [N*m/rad]
        damping = {'joint_a': 1.0, 'joint_b': 1.5}     # [N*m*s/rad]
        ctrl_frequency = 100
        desired_sim_frequency = 200

    class asset:
        file = ""
        # * name of the feet bodies,
        # * used to index body state and contact force tensors
        foot_name = "None"
        penalize_contacts_on = []
        terminate_after_contacts_on = []
        end_effector_names = []
        disable_gravity = False
        disable_motors = False
        # * merge bodies connected by fixed joints.
        # * Specific fixed joints can be kept by adding
        # * " <... dont_collapse="true">
        collapse_fixed_joints = True
        # * fix the base of the robot
        fix_base_link = False
        # * see GymDofDriveModeFlags
        # * (0 is none, 1 is pos tgt, 2 is vel tgt, 3 effort)
        default_dof_drive_mode = 3
        self_collisions = 0  # 1 to disable, 0 to enable...bitwise filter
        replace_cylinder_with_capsule = True
        # * Some .obj meshes must be flipped from y-up to z-up
        flip_visual_attachments = True

        density = 0.001
        angular_damping = 0.
        linear_damping = 0.
        max_angular_velocity = 1000.
        max_linear_velocity = 1000.
        armature = 0.
        thickness = 0.01
        rotor_inertia = 0.
        joint_damping = 0.

    class domain_rand:
        randomize_friction = False
        friction_range = [0.5, 1.25]
        randomize_base_mass = False
        added_mass_range = [-1., 1.]  # added to each link!

    class reward_settings:
        # * tracking reward = exp(-error^2/sigma)
        tracking_sigma = 0.25
        # * percentage of urdf limits, values above this limit are penalized
        soft_dof_pos_limit = 1.
        soft_dof_vel_limit = 1.
        soft_torque_limit = 1.
        base_height_target = 1.
        # * forces above this value are penalized
        max_contact_force = 100.

    class scaling:
        commands = 1
        # * Action scales
        tau_ff = 1  # scale by torque limits
        dof_pos = 1
        dof_pos_obs = dof_pos
        dof_pos_target = dof_pos  # scale by range of motion

    # * viewer camera:
    class viewer:
        ref_env = 0
        pos = [10, 0, 6]  # [m]
        lookat = [11., 5, 3.]  # [m]


class LeggedRobotRunnerCfg(BaseConfig):
    seed = -1
    runner_class_name = 'OnPolicyRunner'

    class logging():
        enable_local_saving = True

    class policy:
        init_noise_std = 1.0
        actor_hidden_dims = [512, 256, 128]
        critic_hidden_dims = [512, 256, 128]
        # can be elu, relu, selu, crelu, lrelu, tanh, sigmoid
        activation = 'elu'

        actor_obs = [
            "observation_a",
            "observation_b",
            "these_need_to_be_atributes_(states)_of_the_robot_env"]
        critic_obs = [
            "observation_x",
            "observation_y",
            "critic_obs_can_be_the_same_or_different_than_actor_obs"]

        actions = ["q_des"]
        disable_actions = False

        class noise:
            dof_pos_obs = 0.01
            dof_vel = 1.5
            lin_vel = 0.1
            ang_vel = 0.2
            projected_gravity = 0.05
            height_measurements = 0.1

        class reward:
            class weights:
                tracking_lin_vel = .0
                tracking_ang_vel = 0.
                lin_vel_z = 0
                ang_vel_xy = 0.
                orientation = 0.
                torques = 0.
                dof_vel = 0.
                base_height = 0.
                collision = 0.
                action_rate = 0.
                action_rate2 = 0.
                stand_still = 0.
                dof_pos_limits = 0.

            class termination_weight:
                termination = 0.01

    class algorithm:
        # * training params
        value_loss_coef = 1.0
        use_clipped_value_loss = True
        clip_param = 0.2
        entropy_coef = 0.01
        num_learning_epochs = 5
        # * mini batch size = num_envs*nsteps / nminibatches
        num_mini_batches = 4
        learning_rate = 1.e-3
        schedule = 'adaptive'  # could be adaptive, fixed
        gamma = 0.99
        lam = 0.95
        desired_kl = 0.01
        max_grad_norm = 1.

    class runner:
        policy_class_name = 'ActorCritic'
        algorithm_class_name = 'PPO'
        num_steps_per_env = 24
        max_iterations = 1500
        save_interval = 50
        run_name = ''
        experiment_name = 'legged_robot'

        # load and resume
        resume = False
        load_run = -1  # -1 = last run
        checkpoint = -1  # -1 = last saved model
        resume_path = None  # updated from load_run and chkpt
        device = 'cpu'

gym/envs/base/legged_robot_config.py
mit-biomimetics/ORCAgym
class SimCfg:
    dt = 0.001  # calculated based on env.config.control.desired_sim_frequency
    substeps = 1
    gravity = [0., 0., -9.81]  # [m/s^2]
    up_axis = 1  # 0 is y, 1 is z

    class physx:
        num_threads = 10
        solver_type = 1  # 0: pgs, 1: tgs (more robust, more expensive)
        num_position_iterations = 4
        num_velocity_iterations = 0
        contact_offset = 0.01
        rest_offset = 0.0
        bounce_threshold_velocity = 0.5  # 0.5 [m/s]
        # * 0: never, 1: last sub-step, 2: all sub-steps (default=2)
        contact_collection = 2
        max_depenetration_velocity = 10.0
        # * 2**24 needed for 8000 envs or more
        max_gpu_contact_pairs = 2**23
        default_buffer_size_multiplier = 5

gym/envs/base/sim_config.py
mit-biomimetics/ORCAgym
import torch

from gym.envs.base.legged_robot import LeggedRobot


class MiniCheetah(LeggedRobot):
    def __init__(self, gym, sim, cfg, sim_params, sim_device, headless):
        super().__init__(gym, sim, cfg, sim_params, sim_device, headless)

    def _init_buffers(self):
        super()._init_buffers()

    def _reward_lin_vel_z(self):
        """Penalize z axis base linear velocity with squared exp"""
        return self._sqrdexp(self.base_lin_vel[:, 2]
                             / self.scales["base_lin_vel"])

    def _reward_ang_vel_xy(self):
        """Penalize xy axes base angular velocity"""
        error = self._sqrdexp(self.base_ang_vel[:, :2]
                              / self.scales["base_ang_vel"])
        return torch.sum(error, dim=1)

    def _reward_orientation(self):
        """Penalize non-flat base orientation"""
        error = (torch.square(self.projected_gravity[:, :2])
                 / self.cfg.reward_settings.tracking_sigma)
        return torch.sum(torch.exp(-error), dim=1)

    def _reward_min_base_height(self):
        """Squared exponential saturating at base_height target"""
        error = (self.base_height-self.cfg.reward_settings.base_height_target)
        error /= self.scales["base_height"]
        error = torch.clamp(error, max=0, min=None).flatten()
        return self._sqrdexp(error)

    def _reward_tracking_lin_vel(self):
        """Tracking of linear velocity commands (xy axes)"""
        # just use lin_vel?
        error = self.commands[:, :2] - self.base_lin_vel[:, :2]
        # * scale by (1+|cmd|): if cmd=0, no scaling.
        error *= 1./(1. + torch.abs(self.commands[:, :2]))
        error = torch.sum(torch.square(error), dim=1)
        return torch.exp(-error/self.cfg.reward_settings.tracking_sigma)

    def _reward_tracking_ang_vel(self):
        """Tracking of angular velocity commands (yaw)"""
        ang_vel_error = torch.square((self.commands[:, 2]
                                      - self.base_ang_vel[:, 2])/5.)
        return self._sqrdexp(ang_vel_error)

    def _reward_dof_vel(self):
        """Penalize dof velocities"""
        return torch.sum(
            self._sqrdexp(self.dof_vel / self.scales["dof_vel"]), dim=1)

    def _reward_dof_near_home(self):
        return torch.sum(
            self._sqrdexp(
                (self.dof_pos - self.default_dof_pos)
                / self.scales["dof_pos_obs"]), dim=1)

gym/envs/mini_cheetah/mini_cheetah.py
mit-biomimetics/ORCAgym
from gym.envs.base.legged_robot_config \
    import LeggedRobotCfg, LeggedRobotRunnerCfg

BASE_HEIGHT_REF = 0.33


class MiniCheetahCfg(LeggedRobotCfg):
    class env(LeggedRobotCfg.env):
        num_envs = 2**12
        num_actuators = 12
        episode_length_s = 10

    class terrain(LeggedRobotCfg.terrain):
        mesh_type = 'plane'

    class init_state(LeggedRobotCfg.init_state):
        default_joint_angles = {
            "haa": 0.0,
            "hfe": -0.785398,
            "kfe": 1.596976,
        }

        # * reset setup chooses how the initial conditions are chosen.
        # * "reset_to_basic" = a single position
        # * "reset_to_range" = uniformly random from a range defined below
        reset_mode = "reset_to_basic"

        # * default COM for basic initialization
        pos = [0.0, 0.0, 0.33]  # x,y,z [m]
        rot = [0.0, 0.0, 0.0, 1.0]  # x,y,z,w [quat]
        lin_vel = [0.0, 0.0, 0.0]  # x,y,z [m/s]
        ang_vel = [0.0, 0.0, 0.0]  # x,y,z [rad/s]

        # * initialization for random range setup
        dof_pos_range = {'haa': [-0.05, 0.05],
                         'hfe': [-0.85, -0.6],
                         'kfe': [-1.45, 1.72]}
        dof_vel_range = {'haa': [0., 0.],
                         'hfe': [0., 0.],
                         'kfe': [0., 0.]}
        root_pos_range = [[0., 0.],       # x
                          [0., 0.],       # y
                          [0.37, 0.4],    # z
                          [0., 0.],       # roll
                          [0., 0.],       # pitch
                          [0., 0.]]       # yaw
        root_vel_range = [[-0.05, 0.05],  # x
                          [0., 0.],       # y
                          [-0.05, 0.05],  # z
                          [0., 0.],       # roll
                          [0., 0.],       # pitch
                          [0., 0.]]       # yaw

    class control(LeggedRobotCfg.control):
        # * PD Drive parameters:
        stiffness = {'haa': 20., 'hfe': 20., 'kfe': 20.}
        damping = {'haa': 0.5, 'hfe': 0.5, 'kfe': 0.5}
        ctrl_frequency = 100
        desired_sim_frequency = 1000

    class commands:
        # * time before command are changed[s]
        resampling_time = 10.

        class ranges:
            lin_vel_x = [-1.0, 1.0]  # min max [m/s]
            lin_vel_y = 1.           # max [m/s]
            yaw_vel = 1              # max [rad/s]

    class push_robots:
        toggle = False
        interval_s = 1
        max_push_vel_xy = 0.5

    class domain_rand:
        randomize_friction = True
        friction_range = [0.5, 1.25]
        randomize_base_mass = True
        added_mass_range = [-1., 1.]

    class asset(LeggedRobotCfg.asset):
        file = "{LEGGED_GYM_ROOT_DIR}/resources/robots/" \
               + "mini_cheetah/urdf/mini_cheetah_simple.urdf"
        foot_name = "foot"
        penalize_contacts_on = []
        terminate_after_contacts_on = ["base", "thigh"]
        end_effector_names = ['foot']
        # * merge bodies connected by fixed joints.
        collapse_fixed_joints = False
        # * 1 to disable, 0 to enable...bitwise filter
        self_collisions = 1
        flip_visual_attachments = False
        disable_gravity = False
        # * set all torques set to 0
        disable_motors = False
        joint_damping = 0.1
        rotor_inertia = [0.002268, 0.002268, 0.005484]*4

    class reward_settings(LeggedRobotCfg.reward_settings):
        soft_dof_pos_limit = 0.9
        soft_dof_vel_limit = 0.9
        soft_torque_limit = 0.9
        max_contact_force = 600.
        base_height_target = BASE_HEIGHT_REF
        tracking_sigma = 0.25

    class scaling(LeggedRobotCfg.scaling):
        base_ang_vel = 1.
        base_lin_vel = 1.
        commands = 1
        dof_vel = 100.  # ought to be roughly max expected speed.
        base_height = 1
        dof_pos = 1
        dof_pos_obs = dof_pos
        # * Action scales
        dof_pos_target = dof_pos
        # tau_ff = 4*[18, 18, 28]  # hip-abad, hip-pitch, knee
        clip_actions = 1000.


class MiniCheetahRunnerCfg(LeggedRobotRunnerCfg):
    seed = -1

    class policy(LeggedRobotRunnerCfg.policy):
        actor_hidden_dims = [256, 256, 256]
        critic_hidden_dims = [256, 256, 256]
        # * can be elu, relu, selu, crelu, lrelu, tanh, sigmoid
        activation = 'elu'

        actor_obs = [
            "base_height",
            "base_lin_vel",
            "base_ang_vel",
            "projected_gravity",
            "commands",
            "dof_pos_obs",
            "dof_vel"
            ]
        critic_obs = [
            "base_height",
            "base_lin_vel",
            "base_ang_vel",
            "projected_gravity",
            "commands",
            "dof_pos_obs",
            "dof_vel"
            ]

        actions = ["dof_pos_target"]
        add_noise = False

        class noise:
            dof_pos_obs = 0.005  # can be made very low
            dof_vel = 0.005
            base_ang_vel = 0.05
            projected_gravity = 0.02

        class reward(LeggedRobotRunnerCfg.policy.reward):
            class weights(LeggedRobotRunnerCfg.policy.reward.weights):
                tracking_lin_vel = 5.0
                tracking_ang_vel = 5.0
                lin_vel_z = 0.
                ang_vel_xy = 0.0
                orientation = 1.0
                torques = 5.e-7
                dof_vel = 1.
                base_height = 1.
                action_rate = 0.001
                action_rate2 = 0.0001
                stand_still = 0.
                dof_pos_limits = 0.
                feet_contact_forces = 0.
                dof_near_home = 1.

            class termination_weight:
                termination = 0.01

    class algorithm(LeggedRobotRunnerCfg.algorithm):
        # * training params
        value_loss_coef = 1.0
        use_clipped_value_loss = True
        clip_param = 0.2
        entropy_coef = 0.01
        num_learning_epochs = 6
        # * mini batch size = num_envs*nsteps / nminibatches
        num_mini_batches = 6
        learning_rate = 1.e-4
        schedule = 'adaptive'  # can be adaptive or fixed
        discount_horizon = 1.  # [s]
        GAE_bootstrap_horizon = 1.  # [s]
        desired_kl = 0.01
        max_grad_norm = 1.

    class runner(LeggedRobotRunnerCfg.runner):
        run_name = ''
        experiment_name = 'mini_cheetah'
        # * number of policy updates
        max_iterations = 1000
        algorithm_class_name = 'PPO'
        # * per iteration
        # * (n_steps in Rudin 2021 paper - batch_size = n_steps * n_robots)
        num_steps_per_env = 24

gym/envs/mini_cheetah/mini_cheetah_config.py
mit-biomimetics/ORCAgym
import torch
import pandas as pd
from isaacgym.torch_utils import torch_rand_float, to_torch

from gym.utils.math import exp_avg_filter

from gym import LEGGED_GYM_ROOT_DIR
from gym.envs.mini_cheetah.mini_cheetah_osc import MiniCheetahOsc

MINI_CHEETAH_MASS = 8.292 * 9.81  # Weight of mini cheetah in Newtons


class MiniCheetahFloquet(MiniCheetahOsc):

    def _init_buffers(self):
        super()._init_buffers()

        self.last_phase_1 = torch.zeros(self.num_envs, 4, 2,
                                        device=self.device)
        self.last_phase_2 = torch.zeros(self.num_envs, 4, 2,
                                        device=self.device)
        self.last_phase_3 = torch.zeros(self.num_envs, 4, 2,
                                        device=self.device)
        self.last_phase_4 = torch.zeros(self.num_envs, 4, 2,
                                        device=self.device)
        self.crossings = torch.zeros_like(self.oscillators, dtype=torch.bool)
        self.floquet_reward = torch.zeros_like(self.crossings,
                                               dtype=torch.float32)

        self.average_frequency = (torch.ones_like(self.oscillators)
                                  * self.osc_omega)
        self.max_freq_diff = torch.zeros(self.num_envs, device=self.device)
        self.last_cross = torch.zeros_like(self.oscillators,
                                                   dtype=torch.long)

    def _pre_physics_step(self):
        super()._pre_physics_step()

    def _post_physics_step(self):
        super()._post_physics_step()
        self.max_freq_diff = self.average_frequency.max(dim=1)[0] \
                             - self.average_frequency.min(dim=1)[0]

    def _step_oscillators(self, dt=None):
        if dt is None:
            dt = self.dt

        local_feedback = self.osc_coupling * (torch.cos(self.oscillators)
                                              + self.osc_offset)
        grf = self._compute_grf()
        self.oscillators_vel = self.osc_omega - grf * local_feedback
        # self.oscillators_vel *= torch_rand_float(0.9,
        #                                          1.1,
        #                                          self.oscillators_vel.shape,
        #                                          self.device)
        self.oscillators_vel += (torch.randn(self.oscillators_vel.shape,
                                             device=self.device)
                                 * self.cfg.osc.process_noise_std)

        self.oscillators_vel *= 2*torch.pi
        self.oscillators += self.oscillators_vel*dt
        self.crossings = self.oscillators >= 2 * torch.pi
        self.oscillators = torch.remainder(self.oscillators, 2*torch.pi)

        for osc_id in range(4):
            env_id = self.crossings[:, osc_id].nonzero().flatten()
            if len(env_id) == 0:
                continue
            if osc_id == 0:
                self.last_phase_1.roll(1, dims=2)
                self.last_phase_1[env_id, :, 0] =\
                    self.oscillators[env_id, :]
            if osc_id == 1:
                self.last_phase_2.roll(1, dims=2)
                self.last_phase_2[env_id, :, 0] =\
                    self.oscillators[env_id, :]
            if osc_id == 2:
                self.last_phase_3.roll(1, dims=2)
                self.last_phase_3[env_id, :, 0] =\
                    self.oscillators[env_id, :]
            if osc_id == 3:
                self.last_phase_4.roll(1, dims=2)
                self.last_phase_4[env_id, :, 0] =\
                    self.oscillators[env_id, :]
            # * update average frequency of oscillators that reset
            self._update_avg_frequency(env_id, osc_id)

        self.oscillator_obs = torch.cat((torch.cos(self.oscillators),
                                         torch.sin(self.oscillators)), dim=1)

    def _update_avg_frequency(self, env_id, osc_id):
        time_steps = (self.common_step_counter
                      - self.last_cross[env_id, osc_id])
        self.average_frequency[env_id, osc_id] =\
            exp_avg_filter(1.0/(time_steps*self.dt),
                            self.average_frequency[env_id, osc_id])
        self.last_cross[env_id, osc_id] = self.common_step_counter

    def _reward_floquet(self):
        phase_diff = \
            torch.cat(((self.last_phase_1[:, :, 0] - self.last_phase_1[:, :, 1]).norm(dim=1).unsqueeze(1),
                       (self.last_phase_2[:, :, 0] - self.last_phase_2[:, :, 1]).norm(dim=1).unsqueeze(1),
                       (self.last_phase_3[:, :, 0] - self.last_phase_3[:, :, 1]).norm(dim=1).unsqueeze(1),
                       (self.last_phase_4[:, :, 0] - self.last_phase_4[:, :, 1]).norm(dim=1).unsqueeze(1)), dim=1)
        phase_diff /= 2*torch.pi
        self.floquet_reward = torch.where(self.crossings,
                                          phase_diff, self.floquet_reward)
        vel_error = (self.commands[:, :2] - self.base_lin_vel[:, :2]).norm(dim=1)
        vel_error = self._sqrdexp(vel_error/(self.scales["base_lin_vel"]/2.))
        # * when tracking is decent, penalize large deviations of the floquet multipliers
        reward = -self.floquet_reward.sum(dim=1)*vel_error
        return reward

    def _reward_locked_frequency(self):
        return -self.max_freq_diff


gym/envs/mini_cheetah/mini_cheetah_floquet.py
mit-biomimetics/ORCAgym
from gym.envs.mini_cheetah.mini_cheetah_osc_config \
    import MiniCheetahOscCfg, MiniCheetahOscRunnerCfg

BASE_HEIGHT_REF = 0.33


class MiniCheetahFloquetCfg(MiniCheetahOscCfg):
    class env(MiniCheetahOscCfg.env):
        num_envs = 4096
        num_actuators = 12
        episode_length_s = 1.

    class terrain(MiniCheetahOscCfg.terrain):
        pass

    class init_state(MiniCheetahOscCfg.init_state):
        pass

    class control(MiniCheetahOscCfg.control):
        pass

    class osc(MiniCheetahOscCfg.osc):
        pass

    class commands(MiniCheetahOscCfg.commands):
        pass
        # resampling_time = 4.  # time before command are changed[s]

        class ranges(MiniCheetahOscCfg.commands.ranges):
            pass
            # lin_vel_x = [-1., 4.]  # min max [m/s]
            # lin_vel_y = 1.  # max [m/s]
            # yaw_vel = 6    # max [rad/s]

    class push_robots:
        toggle = True
        interval_s = 10
        max_push_vel_xy = 0.2

    class domain_rand(MiniCheetahOscCfg.domain_rand):
        pass

    class asset(MiniCheetahOscCfg.asset):
        pass

    class reward_settings(MiniCheetahOscCfg.reward_settings):
        pass

    class scaling(MiniCheetahOscCfg.scaling):
        pass

class MiniCheetahFloquetRunnerCfg(MiniCheetahOscRunnerCfg):
    seed = -1

    class policy(MiniCheetahOscRunnerCfg.policy):
        actor_hidden_dims = [256, 256, 128]
        critic_hidden_dims = [256, 256, 128]
        # * can be elu, relu, selu, crelu, lrelu, tanh, sigmoid
        activation = 'elu'

        actor_obs = [
                     "base_ang_vel",
                     "projected_gravity",
                     "commands",
                     "dof_pos_obs",
                     "dof_vel",
                     "oscillator_obs",
                     "dof_pos_target"
                     ]

        critic_obs = ["base_height",
                      "base_lin_vel",
                      "base_ang_vel",
                      "projected_gravity",
                      "commands",
                      "dof_pos_obs",
                      "dof_vel",
                      "oscillator_obs",
                      "oscillators_vel",
                      "dof_pos_target"
                      ]

        actions = ["dof_pos_target"]

        class noise(MiniCheetahOscRunnerCfg.policy.noise):
            pass

        class reward(MiniCheetahOscRunnerCfg.policy.reward):
            class weights(MiniCheetahOscRunnerCfg.policy.reward.weights):
                floquet = 0.
                locked_frequency = 0.

            class termination_weight:
                pass
                # termination = 15./100.

    class algorithm(MiniCheetahOscRunnerCfg.algorithm):
        pass

    class runner(MiniCheetahOscRunnerCfg.runner):
        run_name = 'floquet'
        experiment_name = 'osc'
        max_iterations = 3000  # number of policy updates
        algorithm_class_name = 'PPO'
        num_steps_per_env = 32

gym/envs/mini_cheetah/mini_cheetah_floquet_config.py
mit-biomimetics/ORCAgym
import torch
import pandas as pd
import numpy as np
from isaacgym.torch_utils import torch_rand_float, to_torch

from gym import LEGGED_GYM_ROOT_DIR
from gym.envs.mini_cheetah.mini_cheetah import MiniCheetah

from isaacgym import gymtorch, gymapi

MINI_CHEETAH_MASS = 8.292 * 9.81  # Weight of mini cheetah in Newtons

class MiniCheetahOsc(MiniCheetah):
    def __init__(self, gym, sim, cfg, sim_params, sim_device, headless):
        super().__init__(gym, sim, cfg, sim_params, sim_device, headless)
        self.process_noise_std = self.cfg.osc.process_noise_std

    def _init_buffers(self):
        super()._init_buffers()
        self.oscillators = torch.zeros(self.num_envs, 4,
                                       device=self.device)
        self.oscillator_obs = torch.zeros(self.num_envs, 8,
                                          device=self.device)

        self.oscillators_vel = torch.zeros_like(self.oscillators)
        self.grf = torch.zeros(self.num_envs, 4, device=self.device)
        self.osc_omega = self.cfg.osc.omega \
            * torch.ones(self.num_envs, 1, device=self.device)
        self.osc_coupling = self.cfg.osc.coupling \
            * torch.ones(self.num_envs, 1, device=self.device)
        self.osc_offset = self.cfg.osc.offset \
            * torch.ones(self.num_envs, 1, device=self.device)

    def _reset_oscillators(self, env_ids):
        if len(env_ids) == 0:
            return
                # * random
        if self.cfg.osc.init_to == 'random':
            self.oscillators[env_ids] = torch_rand_float(
                0, 2*torch.pi, shape=self.oscillators[env_ids].shape,
                device=self.device)
        elif self.cfg.osc.init_to == 'standing':
            self.oscillators[env_ids] = 3*torch.pi/2
        elif self.cfg.osc.init_to == 'trot':
            self.oscillators[env_ids] = torch.tensor([0., torch.pi, torch.pi, 0.],
                                                     device=self.device)
        elif self.cfg.osc.init_to == 'pace':
            self.oscillators[env_ids] = torch.tensor([0., torch.pi, 0., torch.pi],
                                                     device=self.device)
            if self.cfg.osc.init_w_offset:
                self.oscillators[env_ids, :] += \
                    torch.rand_like(self.oscillators[env_ids, 0]).unsqueeze(1) * 2 * torch.pi
        elif self.cfg.osc.init_to == 'pronk':
            self.oscillators[env_ids, :] *= 0.
        elif self.cfg.osc.init_to == 'bound':
            self.oscillators[env_ids, :] = torch.tensor([torch.pi, torch.pi, 0., 0.], device=self.device)
        else:
            raise NotImplementedError

        if self.cfg.osc.init_w_offset:
                self.oscillators[env_ids, :] += \
                    torch.rand_like(self.oscillators[env_ids, 0]).unsqueeze(1) * 2 * torch.pi
        self.oscillators = torch.remainder(self.oscillators, 2*torch.pi)


    def _reset_system(self, env_ids):
        if len(env_ids) == 0:
            return
        self._reset_oscillators(env_ids)

        self.oscillator_obs = torch.cat((torch.cos(self.oscillators),
                                         torch.sin(self.oscillators)), dim=1)

        # * keep some robots in the same starting state as they ended
        timed_out_subset = (self.timed_out & ~self.terminated) * \
            (torch.rand(self.num_envs, device=self.device)
             < self.cfg.init_state.timeout_reset_ratio)

        env_ids = (self.terminated | timed_out_subset).nonzero().flatten()
        if len(env_ids) == 0:
            return
        super()._reset_system(env_ids)

    def _pre_physics_step(self):
        super()._pre_physics_step()
        # self.grf = self._compute_grf()
        if not self.cfg.osc.randomize_osc_params:
            self.compute_osc_slope()

    def compute_osc_slope(self):
        cmd_x = torch.abs(self.commands[:, 0:1]) - self.cfg.osc.stop_threshold
        stop = (cmd_x < 0)

        self.osc_offset = stop * self.cfg.osc.offset
        self.osc_omega = stop * self.cfg.osc.omega_stop \
            + torch.randn_like(self.osc_omega) * self.cfg.osc.omega_var
        self.osc_coupling = stop * self.cfg.osc.coupling_stop \
            + torch.randn_like(self.osc_coupling) * self.cfg.osc.coupling_var

        self.osc_omega += (~stop) * torch.clamp(cmd_x*self.cfg.osc.omega_slope
                                                + self.cfg.osc.omega_step,
                                                min=0.,
                                                max=self.cfg.osc.omega_max)
        self.osc_coupling += \
            (~stop) * torch.clamp(cmd_x*self.cfg.osc.coupling_slope
                                  + self.cfg.osc.coupling_step,
                                  min=0.,
                                  max=self.cfg.osc.coupling_max)

        self.osc_omega = torch.clamp_min(self.osc_omega, 0.1)
        self.osc_coupling = torch.clamp_min(self.osc_coupling, 0)

    def _process_rigid_body_props(self, props, env_id):
        if env_id == 0:
            # * init buffers for the domain rand changes
            self.mass = torch.zeros(self.num_envs, 1, device=self.device)
            self.com = torch.zeros(self.num_envs, 3, device=self.device)

        # * randomize mass
        if self.cfg.domain_rand.randomize_base_mass:
            lower = self.cfg.domain_rand.lower_mass_offset
            upper = self.cfg.domain_rand.upper_mass_offset
            # self.mass_
            props[0].mass += np.random.uniform(lower, upper)
            self.mass[env_id] = props[0].mass
            # * randomize com position
            lower = self.cfg.domain_rand.lower_z_offset
            upper = self.cfg.domain_rand.upper_z_offset
            props[0].com.z += np.random.uniform(lower, upper)
            self.com[env_id, 2] = props[0].com.z

            lower = self.cfg.domain_rand.lower_x_offset
            upper = self.cfg.domain_rand.upper_x_offset
            props[0].com.x += np.random.uniform(lower, upper)
            self.com[env_id, 0] = props[0].com.x
        return props


    def _post_physics_step(self):
        """ Update all states that are not handled in PhysX """
        super()._post_physics_step()
        self.grf = self._compute_grf()
        # self._step_oscillators()

    def _post_torque_step(self):
        super()._post_torque_step()
        self._step_oscillators(self.dt/self.cfg.control.decimation)
        return None

    def _step_oscillators(self, dt=None):
        if dt is None:
            dt = self.dt

        local_feedback = self.osc_coupling * (torch.cos(self.oscillators)
                                              + self.osc_offset)
        grf = self._compute_grf()
        self.oscillators_vel = self.osc_omega - grf * local_feedback
        # self.oscillators_vel *= torch_rand_float(0.9,
        #                                          1.1,
        #                                          self.oscillators_vel.shape,
        #                                          self.device)
        self.oscillators_vel += (torch.randn(self.oscillators_vel.shape,
                                             device=self.device)
                                 * self.cfg.osc.process_noise_std)

        self.oscillators_vel *= 2*torch.pi
        self.oscillators += self.oscillators_vel * dt  # torch.clamp(self.oscillators_vel * dt, min=0)
        self.oscillators = torch.remainder(self.oscillators, 2*torch.pi)
        self.oscillator_obs = torch.cat((torch.cos(self.oscillators),
                                         torch.sin(self.oscillators)), dim=1)

    def _resample_commands(self, env_ids):
        """ Randommly select commands of some environments

        Args:
            env_ids (List[int]): Environments ids for which new commands are needed
        """
        if len(env_ids) == 0:
            return
        super()._resample_commands(env_ids)
        possible_commands = torch.tensor(self.command_ranges["lin_vel_x"],
                                         device=self.device)
        self.commands[env_ids, 0:1] = possible_commands[torch.randint(
            0, len(possible_commands), (len(env_ids), 1),
            device=self.device)]
        # add some gaussian noise to the commands
        self.commands[env_ids, 0:1] += torch.randn((len(env_ids), 1),
                                                   device=self.device) \
                                        * self.cfg.commands.var

        # possible_commands = torch.tensor(self.command_ranges["lin_vel_y"],
        #                                  device=self.device)
        # self.commands[env_ids, 1:2] = possible_commands[torch.randint(
        #     0, len(possible_commands), (len(env_ids), 1),
        #     device=self.device)]
        # possible_commands = torch.tensor(self.command_ranges["yaw_vel"],
        #                                  device=self.device)
        # self.commands[env_ids, 0:1] = possible_commands[torch.randint(
        #     0, len(possible_commands), (len(env_ids), 1),
        #     device=self.device)]

        if (0 in self.cfg.commands.ranges.lin_vel_x):
            # * with 20% chance, reset to 0 commands except for forward
            self.commands[env_ids, 1:] *= (torch_rand_float(0, 1, (len(env_ids), 1),
                device=self.device).squeeze(1) < 0.8).unsqueeze(1)
            # * with 20% chance, reset to 0 commands except for rotation
            self.commands[env_ids, :2] *= (torch_rand_float(0, 1, (len(env_ids), 1),
                device=self.device).squeeze(1) < 0.8).unsqueeze(1)
            # * with 10% chance, reset to 0
            self.commands[env_ids, :] *= (torch_rand_float(0, 1, (len(env_ids), 1),
                device=self.device).squeeze(1) < 0.9).unsqueeze(1)

        if self.cfg.osc.randomize_osc_params:
            self._resample_osc_params(env_ids)

    def _resample_osc_params(self, env_ids):
        if (len(env_ids) > 0):
            self.osc_omega[env_ids, 0] = torch_rand_float(self.cfg.osc.omega_range[0],
                                                        self.cfg.osc.omega_range[1],
                                                        (len(env_ids), 1),
                                                        device=self.device).squeeze(1)
            self.osc_coupling[env_ids, 0] = torch_rand_float(self.cfg.osc.coupling_range[0],
                                                            self.cfg.osc.coupling_range[1],
                                                            (len(env_ids), 1),
                                                            device=self.device).squeeze(1)
            self.osc_offset[env_ids, 0] = torch_rand_float(self.cfg.osc.offset_range[0],
                                                        self.cfg.osc.offset_range[1],
                                                        (len(env_ids), 1),
                                                        device=self.device).squeeze(1)

    def perturb_base_velocity(self, velocity_delta, env_ids=None):
        if env_ids is None:
            env_ids = [range(self.num_envs)]
        self.root_states[env_ids, 7:10] += velocity_delta
        self.gym.set_actor_root_state_tensor(self.sim,
                                    gymtorch.unwrap_tensor(self.root_states))


    def _compute_grf(self, grf_norm=True):
        grf = torch.norm(self.contact_forces[:, self.feet_indices, :], dim=-1)
        if grf_norm:
            return torch.clamp_max(grf / MINI_CHEETAH_MASS, 1.0)
        else:
            return grf

    def _switch(self):
        c_vel = torch.linalg.norm(self.commands, dim=1)
        return torch.exp(-torch.square(torch.max(torch.zeros_like(c_vel),
                                                 c_vel-0.1))
                         / self.cfg.reward_settings.switch_scale)

    def _reward_cursorial(self):
        # penalize the abad joints being away from 0
        return -torch.mean(torch.square(self.dof_pos[:, 0:12:3]
                                        /self.scales["dof_pos"][0]), dim=1)

    def _reward_swing_grf(self):
        # Reward non-zero grf during swing (0 to pi)
        rew = self.get_swing_grf(self.cfg.osc.osc_bool, self.cfg.osc.grf_bool)
        return -torch.sum(rew, dim=1)

    def _reward_stance_grf(self):
        # Reward non-zero grf during stance (pi to 2pi)
        rew = self.get_stance_grf(self.cfg.osc.osc_bool, self.cfg.osc.grf_bool)
        return torch.sum(rew, dim=1)

    def get_swing_grf(self, osc_bool=False, contact_bool=False):
        if osc_bool:
            phase = torch.lt(self.oscillators, torch.pi).int()
        else:
            phase = torch.maximum(torch.zeros_like(self.oscillators),
                               torch.sin(self.oscillators))
        if contact_bool:
            return phase*torch.gt(self._compute_grf(),
                                  self.cfg.osc.grf_threshold)
        else:
            return phase*self._compute_grf()

    def get_stance_grf(self, osc_bool=False, contact_bool=False):
        if osc_bool:
            phase = torch.gt(self.oscillators, torch.pi).int()
        else:
            phase = torch.maximum(torch.zeros_like(self.oscillators),
                               - torch.sin(self.oscillators))
        if contact_bool:
            return phase*torch.gt(self._compute_grf(),
                                  self.cfg.osc.grf_threshold)
        else:
            return phase*self._compute_grf()

    def _reward_coupled_grf(self):
        """
        Multiply rewards for stance/swing grf, discount when undesirable
        behavior (grf during swing, no grf during stance)
        """
        swing_rew = self.get_swing_grf()
        stance_rew = self.get_stance_grf()
        combined_rew = self._sqrdexp(swing_rew*2) + stance_rew
        prod = torch.prod(torch.clip(combined_rew, 0, 1), dim=1)
        return prod - torch.ones_like(prod)

    def _reward_swing_velocity(self):
        """ Reward non-zero end effector velocity during swing (0 to pi) """
        # velocity = torch.tanh(torch.norm(self.end_effector_lin_vel, dim=-1))
        velocity = torch.zeros_like(self.oscillators)  # TODO: Grab velocity from AJ V2 code
        phase_bool = torch.lt(self.oscillators, torch.pi).int()
        phase_sin = torch.maximum(torch.zeros_like(self.oscillators),
                               torch.sin(self.oscillators))
        if self.cfg.osc.osc_bool:
            rew = phase_bool*velocity
        else:
            rew = phase_sin*velocity
        return torch.mean(rew, dim=1)

    def _reward_stance_velocity(self):
        """ Reward zero end effector velocity during swing (pi to 2pi) """
        # velocity = torch.tanh(torch.norm(self.end_effector_lin_vel, dim=-1))
        velocity = torch.zeros_like(self.oscillators)  # TODO: Grab velocity from AJ V2 code
        ph_bool = torch.gt(self.oscillators, torch.pi).int()
        ph_sin = torch.maximum(torch.zeros_like(self.oscillators),
                               - torch.sin(self.oscillators))
        if self.cfg.osc.osc_bool:
            rew = ph_bool*velocity
        else:
            rew = ph_sin*velocity
        # return torch.mean(self._sqrdexp(rew), dim=1)  # TODO: Higher reward if rew close to zero (how to decouple so that ignore if not in stance??)
        return - torch.mean(rew, dim=1)

    def _reward_dof_vel(self):
        return super()._reward_dof_vel()*self._switch()

    def _reward_dof_near_home(self):
        return super()._reward_dof_near_home()*self._switch()

    def _reward_stand_still(self):
        """Penalize motion at zero commands"""
        # * normalize angles so we care about being within 5 deg
        rew_pos = torch.mean(self._sqrdexp(
            (self.dof_pos - self.default_dof_pos)/torch.pi*36), dim=1)
        rew_vel = torch.mean(self._sqrdexp(self.dof_vel), dim=1)
        rew_base_vel = torch.mean(torch.square(self.base_lin_vel), dim=1)
        rew_base_vel += torch.mean(torch.square(self.base_ang_vel), dim=1)
        return (rew_vel+rew_pos-rew_base_vel)*self._switch()

    def _reward_standing_torques(self):
        """Penalize torques at zero commands"""
        return super()._reward_torques()*self._switch()

    # * gait similarity scores
    def angle_difference(self, theta1, theta2):
        diff = torch.abs(theta1 - theta2) % (2 * torch.pi)
        return torch.min(diff, 2*torch.pi - diff)

    def _reward_trot(self):
        # ! diagonal difference, front right and hind left
        angle = self.angle_difference(self.oscillators[:, 0],
                                      self.oscillators[:, 3])
        similarity = self._sqrdexp(angle, torch.pi)
        # ! diagonal difference, front left and hind right
        angle = self.angle_difference(self.oscillators[:, 1],
                                      self.oscillators[:, 2])
        similarity *= self._sqrdexp(angle, torch.pi)
        # ! diagonal difference, front left and hind right
        angle = self.angle_difference(self.oscillators[:, 0],
                                      self.oscillators[:, 1])
        similarity *= self._sqrdexp(angle - torch.pi, torch.pi)
        # ! diagonal difference, front left and hind right
        angle = self.angle_difference(self.oscillators[:, 2],
                                      self.oscillators[:, 3])
        similarity *= self._sqrdexp(angle - torch.pi, torch.pi)
        return similarity

    def _reward_pronk(self):
        # ! diagonal difference, front right and hind left
        angle = self.angle_difference(self.oscillators[:, 0],
                                      self.oscillators[:, 3])
        similarity = self._sqrdexp(angle, torch.pi)
        # ! diagonal difference, front left and hind right
        angle = self.angle_difference(self.oscillators[:, 1],
                                      self.oscillators[:, 2])
        similarity *= self._sqrdexp(angle, torch.pi)
        # ! difference, front right and front left
        angle = self.angle_difference(self.oscillators[:, 0],
                                      self.oscillators[:, 1])
        similarity *= self._sqrdexp(angle, torch.pi)
        # ! difference front right and hind right
        angle = self.angle_difference(self.oscillators[:, 0],
                                      self.oscillators[:, 2])
        similarity *= self._sqrdexp(angle, torch.pi)
        return similarity

    def _reward_pace(self):
        angle = self.angle_difference(self.oscillators[:, 0],
                                      self.oscillators[:, 2])
        similarity = self._sqrdexp(angle, torch.pi)
        # ! difference front left and hind left
        angle = self.angle_difference(self.oscillators[:, 1],
                                      self.oscillators[:, 3])
        similarity *= self._sqrdexp(angle, torch.pi)
        # ! difference front left and hind left
        angle = self.angle_difference(self.oscillators[:, 0],
                                      self.oscillators[:, 1])
        similarity *= self._sqrdexp(angle - torch.pi, torch.pi)
        # ! difference front left and hind left
        angle = self.angle_difference(self.oscillators[:, 2],
                                      self.oscillators[:, 3])
        similarity *= self._sqrdexp(angle - torch.pi, torch.pi)

        return similarity

    def _reward_any_symm_gait(self):
        rew_trot = self._reward_trot()
        rew_pace = self._reward_pace()
        rew_bound = self._reward_bound()
        return torch.max(torch.max(rew_trot, rew_pace), rew_bound)

    def _reward_enc_pace(self):
        return self._reward_pace()

    def _reward_bound(self):
        # ! difference, front right and front left
        angle = self.angle_difference(self.oscillators[:, 0],
                                      self.oscillators[:, 1])
        similarity = self._sqrdexp(angle, torch.pi)
        # ! difference hind right and hind left
        angle = self.angle_difference(self.oscillators[:, 2],
                                      self.oscillators[:, 3])
        similarity *= self._sqrdexp(angle, torch.pi)
        # ! difference right side
        angle = self.angle_difference(self.oscillators[:, 0],
                                      self.oscillators[:, 2])
        similarity *= self._sqrdexp(angle - torch.pi, torch.pi)
        # ! difference right side
        angle = self.angle_difference(self.oscillators[:, 1],
                                      self.oscillators[:, 3])
        similarity *= self._sqrdexp(angle - torch.pi, torch.pi)
        return similarity

    def _reward_asymettric(self):
        # ! hind legs
        angle = self.angle_difference(self.oscillators[:, 2],
                                      self.oscillators[:, 3])
        similarity = (1 - self._sqrdexp(angle, torch.pi))
        similarity *= (1 - self._sqrdexp((torch.pi - angle), torch.pi))
        # ! difference, left legs
        angle = self.angle_difference(self.oscillators[:, 1],
                                      self.oscillators[:, 3])
        similarity *= (1 - self._sqrdexp(angle, torch.pi))
        similarity *= (1 - self._sqrdexp((torch.pi - angle), torch.pi))
        # ! difference right legs
        angle = self.angle_difference(self.oscillators[:, 0],
                                      self.oscillators[:, 2])
        similarity *= (1 - self._sqrdexp(angle, torch.pi))
        similarity *= (1 - self._sqrdexp((torch.pi - angle), torch.pi))
        # ! front legs
        angle = self.angle_difference(self.oscillators[:, 0],
                                      self.oscillators[:, 1])
        similarity *= (1 - self._sqrdexp(angle, torch.pi))
        similarity *= (1 - self._sqrdexp((torch.pi - angle), torch.pi))
        # ! diagonal FR
        angle = self.angle_difference(self.oscillators[:, 0],
                                        self.oscillators[:, 3])
        similarity *= (1 - self._sqrdexp(angle, torch.pi))
        similarity *= (1 - self._sqrdexp((torch.pi - angle), torch.pi))
        # ! diagonal FL
        angle = self.angle_difference(self.oscillators[:, 1],
                                        self.oscillators[:, 2])
        similarity *= (1 - self._sqrdexp(angle, torch.pi))
        similarity *= (1 - self._sqrdexp((torch.pi - angle), torch.pi))
        return similarity
gym/envs/mini_cheetah/mini_cheetah_osc.py
mit-biomimetics/ORCAgym
import torch
import pandas as pd
from isaacgym.torch_utils import torch_rand_float, to_torch

from gym import LEGGED_GYM_ROOT_DIR
from gym.envs.mini_cheetah.mini_cheetah_osc import MiniCheetahOsc

MINI_CHEETAH_MASS = 8.292 * 9.81  # Weight of mini cheetah in Newtons

class MiniCheetahOscIK(MiniCheetahOsc):

    def _init_buffers(self):
        super()._init_buffers()
        self.ik_pos_target = torch.zeros(self.num_envs, 12,
                                         dtype=torch.float, device=self.device)

    def _pre_physics_step(self):
        super()._pre_physics_step()
        ik_defaults_L = torch.tensor([0., -0.019, -0.5]).to(self.device)
        ik_defaults_R = torch.tensor([0., 0.019, -0.5]).to(self.device)
        joints_ik_fr = self.IK_leg_3DOF(p_hr2ft_hr=self.ik_pos_target[:, 0:3]
                                        + ik_defaults_R)
        joints_ik_fl = self.IK_leg_3DOF(p_hr2ft_hr=self.ik_pos_target[:, 3:6]
                                        + ik_defaults_L)
        joints_ik_br = self.IK_leg_3DOF(p_hr2ft_hr=self.ik_pos_target[:, 6:9]
                                        + ik_defaults_R)
        joints_ik_bl = self.IK_leg_3DOF(p_hr2ft_hr=self.ik_pos_target[:, 9:12]
                                        + ik_defaults_L)
        self.dof_pos_target = torch.cat(
            (joints_ik_fr, joints_ik_fl, joints_ik_br, joints_ik_bl), dim=1)

    def IK_leg_3DOF(self, p_hr2ft_hr):
        # todo: clean the variable names here to make more sense
        r_hr2hp = torch.tensor([0.0, -0.019, 0.0], device=self.device)\
            .unsqueeze(0).repeat(self.num_envs, 1)
        r_hp2kp = torch.tensor([0.0, 0.0, -0.2085], device=self.device)\
            .unsqueeze(0).repeat(self.num_envs, 1)
        r_kp2ft = torch.tensor([0.0, 0.0, -0.22], device=self.device)\
            .unsqueeze(0).repeat(self.num_envs, 1)

        L1 = torch.abs(r_hp2kp[:, 2])
        L2 = torch.abs(r_kp2ft[:, 2])

        # * hip-roll angle
        L_hr2ft_yz = torch.norm(p_hr2ft_hr[:, 1:3], dim=1).clamp(min=1e-6)
        alpha_1 = torch.arcsin(
            torch.clamp(p_hr2ft_hr[:, 1] / L_hr2ft_yz, min=-1, max=1))
        beta_1 = torch.arcsin(
            torch.clamp(torch.abs(r_hr2hp[:, 1])/L_hr2ft_yz, min=-1, max=1))
        q1 = alpha_1 - beta_1

        # * hip-pitch angle
        p_hr2hp_hr = torch.zeros(self.num_envs, 3).to(self.device)
        p_hr2hp_hr[:, 1] = r_hr2hp[:, 1]*torch.cos(q1)
        p_hr2hp_hr[:, 2] = r_hr2hp[:, 1]*torch.sin(q1)
        L_hp2ft_hr = torch.clamp(
            torch.norm(p_hr2ft_hr-p_hr2hp_hr, dim=1), min=1e-6)
        vec = p_hr2ft_hr - p_hr2hp_hr
        p_hp2ft_hp = torch.stack(
            (vec[:, 0],
             vec[:, 1]*torch.cos(q1)+vec[:, 2]*torch.sin(q1),
             vec[:, 2]*torch.cos(q1)-vec[:, 1]*torch.sin(q1)), dim=1)
        alpha_2 = torch.arccos(
            torch.clamp(p_hp2ft_hp[:, 0] / L_hp2ft_hr, min=-1, max=1))
        cos_angle_2 = (L1**2 + L_hp2ft_hr**2 - L2**2)/(2*L1*L_hp2ft_hr)
        cos_angle_2 = torch.clamp(cos_angle_2, min=-1, max=1)
        beta_2 = torch.arccos(cos_angle_2)
        q2 = alpha_2 + beta_2 - torch.pi / 2

        # * knee-pitch angle
        cos_angle_3 = (L1**2+L2**2-L_hp2ft_hr**2)/(2*L1*L2)
        cos_angle_3 = torch.clamp(cos_angle_3, min=-1, max=1)
        acos_angle_3 = torch.arccos(cos_angle_3)
        q3 = -(torch.pi - acos_angle_3)

        return torch.stack((q1, -q2, -q3), dim=1)
gym/envs/mini_cheetah/mini_cheetah_osc_IK.py
mit-biomimetics/ORCAgym
from gym.envs.mini_cheetah.mini_cheetah_osc_config \
    import MiniCheetahOscCfg, MiniCheetahOscRunnerCfg

BASE_HEIGHT_REF = 0.33


class MiniCheetahOscIKCfg(MiniCheetahOscCfg):
    class env(MiniCheetahOscCfg.env):
        num_envs = 4096
        num_actuators = 12
        episode_length_s = 1.

    class terrain(MiniCheetahOscCfg.terrain):
        pass

    class init_state(MiniCheetahOscCfg.init_state):
        pass

    class control(MiniCheetahOscCfg.control):
        pass

    class osc(MiniCheetahOscCfg.osc):
        pass

    class commands:
        resampling_time = 4.  # time before command are changed[s]

        class ranges:
            lin_vel_x = [-1., 4.]  # min max [m/s]
            lin_vel_y = 1.  # max [m/s]
            yaw_vel = 3.14    # max [rad/s]

    class push_robots:
        toggle = False
        interval_s = 10
        max_push_vel_xy = 0.2

    class domain_rand(MiniCheetahOscCfg.domain_rand):
        pass

    class asset(MiniCheetahOscCfg.asset):
        pass

    class reward_settings(MiniCheetahOscCfg.reward_settings):
        pass

    class scaling(MiniCheetahOscCfg.scaling):
        ik_pos_target = 0.015


class MiniCheetahOscIKRunnerCfg(MiniCheetahOscRunnerCfg):
    seed = -1

    class policy(MiniCheetahOscRunnerCfg.policy):
        actor_hidden_dims = [256, 256, 128]
        critic_hidden_dims = [256, 256, 128]
        # * can be elu, relu, selu, crelu, lrelu, tanh, sigmoid
        activation = 'elu'

        actor_obs = [
                     "base_ang_vel",
                     "projected_gravity",
                     "commands",
                     "dof_pos_obs",
                     "dof_vel",
                     "oscillator_obs"
                     #  "oscillators_vel",
                     #  "grf",
                     #  "osc_coupling",
                     #  "osc_offset"
                     ]

        critic_obs = ["base_height",
                      "base_lin_vel",
                      "base_ang_vel",
                      "projected_gravity",
                      "commands",
                      "dof_pos_obs",
                      "dof_vel",
                      "oscillator_obs",
                      "oscillators_vel"
                      ]

        actions = ["ik_pos_target"]

        class noise(MiniCheetahOscRunnerCfg.policy.noise):
            pass

        class reward(MiniCheetahOscRunnerCfg.policy.reward):
            class weights(MiniCheetahOscRunnerCfg.policy.reward.weights):
                pass

            class termination_weight:
                termination = 15./100.

    class algorithm(MiniCheetahOscRunnerCfg.algorithm):
        pass

    class runner(MiniCheetahOscRunnerCfg.runner):
        run_name = 'IK'
        experiment_name = 'oscIK'
        max_iterations = 1000  # number of policy updates
        algorithm_class_name = 'PPO'
        num_steps_per_env = 32

gym/envs/mini_cheetah/mini_cheetah_osc_IK_config.py
mit-biomimetics/ORCAgym
from gym.envs.mini_cheetah.mini_cheetah_config \
    import MiniCheetahCfg, MiniCheetahRunnerCfg

BASE_HEIGHT_REF = 0.33


class MiniCheetahOscCfg(MiniCheetahCfg):

    class viewer:
        ref_env = 0
        pos = [6, 10, 4]  # [m]
        lookat = [6, 0, 1]  # [m]

    class env(MiniCheetahCfg.env):
        num_envs = 4096
        num_actuators = 12
        episode_length_s = 4.
        env_spacing = 3.

    class terrain(MiniCheetahCfg.terrain):
        mesh_type = 'plane'
        # mesh_type = 'trimesh'  # none, plane, heightfield or trimesh

    class init_state(MiniCheetahCfg.init_state):
        reset_mode = "reset_to_range"
        timeout_reset_ratio = 0.75
        # * default COM for basic initialization 
        pos = [0.0, 0.0, 0.35]  # x,y,z [m]
        rot = [0.0, 0.0, 0.0, 1.0]  # x,y,z,w [quat]
        lin_vel = [0.0, 0.0, 0.0]  # x,y,z [m/s]
        ang_vel = [0.0, 0.0, 0.0]  # x,y,z [rad/s]
        default_joint_angles = {
            "haa": 0.0,
            "hfe": -0.785398,
            "kfe": 1.596976,
        }

        # * initialization for random range setup
        dof_pos_range = {'haa': [-0.00001, 0.00001],
                         'hfe': [-0.785398, -0.785398],
                         'kfe': [1.596976, 1.596976]}
        dof_vel_range = {'haa': [0., 0.],
                         'hfe': [0., 0.],
                         'kfe': [0., 0.]}
        root_pos_range = [[0., 0.],       # x
                          [0., 0.],       # y
                          [0.35, 0.35],    # z
                          [0., 0.],       # roll
                          [0., 0.],       # pitch
                          [0., 0.]]       # yaw
        root_vel_range = [[-0.5, 2.],  # x
                          [0., 0.],       # y
                          [-0.05, 0.05],  # z
                          [0., 0.],       # roll
                          [0., 0.],       # pitch
                          [0., 0.]]       # yaw

    class control(MiniCheetahCfg.control):
        # * PD Drive parameters:
        stiffness = {'haa': 20., 'hfe': 20., 'kfe': 20.}
        damping = {'haa': 0.5, 'hfe': 0.5, 'kfe': 0.5}
        ctrl_frequency = 100
        desired_sim_frequency = 500

    class osc:
        process_noise_std = 0.25
        grf_threshold = 0.1  # 20. # Normalized to body weight
        # oscillator parameters
        omega = 3  # gets overwritten
        coupling = 1  # gets overwritten
        osc_bool = False  # not used in paper
        grf_bool = False  # not used in paper
        randomize_osc_params = False
        # randomization ranges: not used in paper
        omega_range = [1.0, 4.0]  # [0.0, 10.]
        coupling_range = [0.0, 1.0]  # with normalized grf, can have omega/coupling on same scale
        offset_range = [0.0, 0.0]

        # choice of oscillator parameters, see paper equation (5)
        stop_threshold = 0.5
        omega_stop = 1.
        omega_step = 2.
        omega_slope = 1.
        omega_max = 4.0
        omega_var = 0.25
        coupling_stop = 4.
        coupling_step = 1.
        coupling_slope = 0.
        coupling_max = 1.0
        offset = 1.
        coupling_var = 0.25

        init_to = 'random'
        init_w_offset = True


    class commands:
        resampling_time = 3.  # time before command are changed[s]
        var = 1.0
        class ranges:
            lin_vel_x = [-3., -1., 0., 1., 3.]  # min max [m/s]
            lin_vel_y = 1.  # [-1., 0, 1.]  # max [m/s]
            yaw_vel = 3.  # [-6., -3., 0., 3., 6.]    # max [rad/s]

    class push_robots:
        toggle = True
        interval_s = 10
        max_push_vel_xy = 0.2

    class domain_rand(MiniCheetahCfg.domain_rand):
        randomize_friction = True
        friction_range = [0.4, 1.0]
        randomize_base_mass = False
        lower_mass_offset = -0.5  # kg
        upper_mass_offset = 2.
        lower_z_offset = 0.0  # m
        upper_z_offset = 0.2
        lower_x_offset = 0.0
        upper_x_offset = 0.0

    class asset(MiniCheetahCfg.asset):
        shank_length_diff = 0  # Units in cm
        # file = "{LEGGED_GYM_ROOT_DIR}/resources/robots/" \
        #     + "mini_cheetah/urdf/mini_cheetah_" \
        #     + str(shank_length_diff) + ".urdf"
        file = ("{LEGGED_GYM_ROOT_DIR}/resources/robots/"
                + "mini_cheetah/urdf/mini_cheetah_0.urdf")
        foot_name = "foot"
        penalize_contacts_on = ["thigh", "shank"]
        terminate_after_contacts_on = ["base"]
        collapse_fixed_joints = False
        fix_base_link = False
        self_collisions = 1  # 1 to disable, 0 to enable...bitwise filter
        flip_visual_attachments = False
        disable_gravity = False
        disable_motors = False  # all torques set to 0

    class reward_settings(MiniCheetahCfg.reward_settings):
        soft_dof_pos_limit = 0.8
        soft_dof_vel_limit = 0.9
        soft_torque_limit = 0.9
        max_contact_force = 600.
        base_height_target = BASE_HEIGHT_REF+0.03
        tracking_sigma = 0.25
        switch_scale = 0.5  # 0.1

    class scaling(MiniCheetahCfg.scaling):
        base_ang_vel = [0.3, 0.3, 0.1]
        base_lin_vel = BASE_HEIGHT_REF
        # dof_vel = 100.
        # dof_vel = 4*[41, 41, 27]  # ought to be roughly max expected speed.
        dof_vel = 4*[2., 2., 4.]
        base_height = 0.3
        dof_pos = 4*[0.2, 0.3, 0.3]  # hip-abad, hip-pitch, knee
        dof_pos_obs = dof_pos
        dof_pos_target = 4*[0.2, 0.3, 0.3]
        tau_ff = 4*[18, 18, 28]  # hip-abad, hip-pitch, knee
        # commands = [base_lin_vel, base_lin_vel, base_ang_vel]
        commands = [3, 1, 3]  # [base_lin_vel, base_lin_vel, base_ang_vel]

class MiniCheetahOscRunnerCfg(MiniCheetahRunnerCfg):
    seed = -1

    class policy(MiniCheetahRunnerCfg.policy):
        actor_hidden_dims = [256, 256, 128]
        critic_hidden_dims = [256, 256, 128]
        # * can be elu, relu, selu, crelu, lrelu, tanh, sigmoid
        activation = 'elu'

        actor_obs = [
                     "base_ang_vel",
                     "projected_gravity",
                     "commands",
                     "dof_pos_obs",
                     "dof_vel",
                     "oscillator_obs",
                     "dof_pos_target",
                    #  "osc_omega",
                    #  "osc_coupling"
                     #  "oscillators_vel",
                     #  "grf",
                     ]

        critic_obs = ["base_height",
                      "base_lin_vel",
                      "base_ang_vel",
                      "projected_gravity",
                      "commands",
                      "dof_pos_obs",
                      "dof_vel",
                      "oscillator_obs",
                      "oscillators_vel",
                      "dof_pos_target",
                    #   "osc_omega",
                    #   "osc_coupling"
                    #   "osc_omega",
                    #   "osc_coupling",
                    #   "osc_offset"
                      ]

        actions = ["dof_pos_target"]

        class noise:
            scale = 1.0
            dof_pos_obs = 0.01  # can be made very low
            base_ang_vel = 0.01
            dof_pos = 0.005
            dof_vel = 0.005
            lin_vel = 0.05
            ang_vel = [0.3, 0.15, 0.4]
            gravity_vec = 0.1

        class reward:
            class weights:
                tracking_lin_vel = 4.0
                tracking_ang_vel = 2.0
                lin_vel_z = 0.
                ang_vel_xy = 0.
                orientation = 1.0
                torques = 5.e-7
                dof_vel = 0.
                min_base_height = 1.5
                collision = 0
                action_rate = 0.01  # -0.01
                action_rate2 = 0.001  # -0.001
                stand_still = 0.
                dof_pos_limits = 0.
                feet_contact_forces = 0.
                dof_near_home = 0.
                swing_grf = 5.
                stance_grf = 5.
                swing_velocity = 0.
                stance_velocity = 0.
                coupled_grf = 0.  # 8.
                enc_pace = 0.
                cursorial = 0.25
                standing_torques = 0.  # 1.e-5

            class termination_weight:
                termination = 15./100.

    class algorithm(MiniCheetahRunnerCfg.algorithm):
        # training params
        value_loss_coef = 1.0
        use_clipped_value_loss = True
        clip_param = 0.2
        entropy_coef = 0.02
        num_learning_epochs = 4
        # mini batch size = num_envs*nsteps/nminibatches
        num_mini_batches = 8
        learning_rate = 1.e-5
        schedule = 'adaptive'  # can be adaptive, fixed
        discount_horizon = 1.
        GAE_bootstrap_horizon = 2.
        desired_kl = 0.01
        max_grad_norm = 1.

    class runner(MiniCheetahRunnerCfg.runner):
        run_name = ''
        experiment_name = 'FullSend'
        max_iterations = 2000  # number of policy updates
        algorithm_class_name = 'PPO'
        num_steps_per_env = 32

gym/envs/mini_cheetah/mini_cheetah_osc_config.py
mit-biomimetics/ORCAgym
import torch
import pandas as pd
from isaacgym.torch_utils import torch_rand_float, to_torch

from gym import LEGGED_GYM_ROOT_DIR
from gym.envs.mini_cheetah.mini_cheetah import MiniCheetah


class MiniCheetahRef(MiniCheetah):
    def __init__(self, gym, sim, cfg, sim_params, sim_device, headless):
        csv_path = cfg.init_state.ref_traj.format(
            LEGGED_GYM_ROOT_DIR=LEGGED_GYM_ROOT_DIR)
        self.leg_ref = to_torch(pd.read_csv(csv_path).to_numpy(),
                                device=sim_device)
        self.omega = 2*torch.pi*cfg.control.gait_freq
        super().__init__(gym, sim, cfg, sim_params, sim_device, headless)

    def _init_buffers(self):
        super()._init_buffers()
        self.phase = torch.zeros(self.num_envs, 1, dtype=torch.float,
                                 device=self.device)
        self.phase_obs = torch.zeros(self.num_envs, 2, dtype=torch.float,
                                     device=self.device)

    def _reset_system(self, env_ids):
        super()._reset_system(env_ids)
        self.phase[env_ids] = torch_rand_float(0, torch.pi,
                                               shape=self.phase[env_ids].shape,
                                               device=self.device)

    def _post_physics_step(self):
        """ Update all states that are not handled in PhysX """
        super()._post_physics_step()
        self.phase = torch.fmod(self.phase+self.dt*self.omega, 2*torch.pi)
        self.phase_obs = torch.cat((torch.sin(self.phase),
                                    torch.cos(self.phase)), dim=1)

    def _resample_commands(self, env_ids):
        super()._resample_commands(env_ids)
        # * with 10% chance, reset to 0 commands
        rand_ids = torch_rand_float(0, 1, (len(env_ids), 1),
                                    device=self.device).squeeze(1)
        self.commands[env_ids, :3] *= (rand_ids < 0.9).unsqueeze(1)

    def _switch(self):
        c_vel = torch.linalg.norm(self.commands, dim=1)
        return torch.exp(-torch.square(torch.max(torch.zeros_like(c_vel),
                                                 c_vel-0.1))/0.1)

    def _reward_swing_grf(self):
        """Reward non-zero grf during swing (0 to pi)"""
        in_contact = \
            torch.gt(torch.norm(self.contact_forces[:, self.feet_indices, :],
                                dim=-1), 50.)
        ph_off = torch.lt(self.phase, torch.pi)
        rew = in_contact*torch.cat((ph_off, ~ph_off, ~ph_off, ph_off), dim=1)
        return -torch.sum(rew.float(), dim=1)*(1-self._switch())

    def _reward_stance_grf(self):
        """Reward non-zero grf during stance (pi to 2pi)"""
        in_contact = torch.gt(torch.norm(
            self.contact_forces[:, self.feet_indices, :], dim=-1), 50.)
        ph_off = torch.gt(self.phase, torch.pi)  # should this be in swing?
        rew = in_contact*torch.cat((ph_off, ~ph_off, ~ph_off, ph_off), dim=1)

        return torch.sum(rew.float(), dim=1)*(1-self._switch())

    def _reward_reference_traj(self):
        """REWARDS EACH LEG INDIVIDUALLY BASED ON ITS POSITION IN THE CYCLE"""
        # * dof position error
        error = self._get_ref() + self.default_dof_pos - self.dof_pos
        error /= self.scales['dof_pos']
        reward = torch.mean(self._sqrdexp(error) - torch.abs(error)*0.2, dim=1)
        # * only when commanded velocity is higher
        return reward*(1-self._switch())

    def _get_ref(self):
        leg_frame = torch.zeros_like(self.torques)
        # offset by half cycle (trot)
        ph_off = torch.fmod(self.phase+torch.pi, 2*torch.pi)
        phd_idx = (torch.round(
            self.phase * (self.leg_ref.size(dim=0)/(2*torch.pi)-1))).long()
        pho_idx = (torch.round(
            ph_off * (self.leg_ref.size(dim=0)/(2*torch.pi)-1))).long()
        leg_frame[:, 0:3] += self.leg_ref[phd_idx.squeeze(), :]
        leg_frame[:, 3:6] += self.leg_ref[pho_idx.squeeze(), :]
        leg_frame[:, 6:9] += self.leg_ref[pho_idx.squeeze(), :]
        leg_frame[:, 9:12] += self.leg_ref[phd_idx.squeeze(), :]
        return leg_frame

    def _reward_stand_still(self):
        """Penalize motion at zero commands"""
        # * normalize angles so we care about being within 5 deg
        rew_pos = torch.mean(self._sqrdexp(
            (self.dof_pos - self.default_dof_pos)/torch.pi*36), dim=1)
        rew_vel = torch.mean(self._sqrdexp(self.dof_vel), dim=1)
        rew_base_vel = torch.mean(torch.square(self.base_lin_vel), dim=1)
        rew_base_vel += torch.mean(torch.square(self.base_ang_vel), dim=1)
        return (rew_vel+rew_pos-rew_base_vel)*self._switch()

    def _reward_tracking_lin_vel(self):
        """Tracking linear velocity commands (xy axes)"""
        # just use lin_vel?
        reward = super()._reward_tracking_lin_vel()
        return reward * (1-self._switch())

gym/envs/mini_cheetah/mini_cheetah_ref.py
mit-biomimetics/ORCAgym
from gym.envs.mini_cheetah.mini_cheetah_config \
    import MiniCheetahCfg, MiniCheetahRunnerCfg

BASE_HEIGHT_REF = 0.33


class MiniCheetahRefCfg(MiniCheetahCfg):
    class env(MiniCheetahCfg.env):
        num_envs = 4096
        num_actuators = 12
        episode_length_s = 15.

    class terrain(MiniCheetahCfg.terrain):
        pass

    class init_state(MiniCheetahCfg.init_state):
        reset_mode = "reset_to_basic"
        # * default COM for basic initialization
        pos = [0.0, 0.0, 0.33]  # x,y,z [m]
        rot = [0.0, 0.0, 0.0, 1.0]  # x,y,z,w [quat]
        lin_vel = [0.0, 0.0, 0.0]  # x,y,z [m/s]
        ang_vel = [0.0, 0.0, 0.0]  # x,y,z [rad/s]
        default_joint_angles = {
            "haa": 0.0,
            "hfe": -0.785398,
            "kfe": 1.596976,
        }
        # * range
        # * initialization for random range setup
        dof_pos_range = {'haa': [-0.05, 0.05],
                         'hfe': [-0.85, -0.6],
                         'kfe': [-1.45, 1.72]}
        dof_vel_range = {'haa': [0., 0.],
                         'hfe': [0., 0.],
                         'kfe': [0., 0.]}
        root_pos_range = [[0., 0.],       # x
                          [0., 0.],       # y
                          [0.35, 0.4],    # z
                          [0., 0.],       # roll
                          [0., 0.],       # pitch
                          [0., 0.]]       # yaw
        root_vel_range = [[-0.5, 2.],  # x
                          [0., 0.],       # y
                          [-0.05, 0.05],  # z
                          [0., 0.],       # roll
                          [0., 0.],       # pitch
                          [0., 0.]]       # yaw

        ref_traj = (
            "{LEGGED_GYM_ROOT_DIR}/resources/robots/"
            + "mini_cheetah/trajectories/single_leg.csv")

    class control(MiniCheetahCfg.control):
        # * PD Drive parameters:
        stiffness = {'haa': 20., 'hfe': 20., 'kfe': 20.}
        damping = {'haa': 0.5, 'hfe': 0.5, 'kfe': 0.5}
        gait_freq = 4.
        ctrl_frequency = 100
        desired_sim_frequency = 1000

    class commands:
        resampling_time = 4.  # time before command are changed[s]

        class ranges:
            lin_vel_x = [0., 3.]  # min max [m/s]
            lin_vel_y = 1.  # max [m/s]
            yaw_vel = 3.14/2.    # max [rad/s]

    class push_robots:
        toggle = True
        interval_s = 10
        max_push_vel_xy = 0.2

    class domain_rand(MiniCheetahCfg.domain_rand):
        randomize_friction = True
        friction_range = [0.75, 1.05]  # ! zap, replace with the one below
        randomize_base_mass = True
        added_mass_range = [-1., 3.]
        friction_range = [0., 1.0]

    class asset(MiniCheetahCfg.asset):
        file = (
            "{LEGGED_GYM_ROOT_DIR}/resources/robots/"
            + "mini_cheetah/urdf/mini_cheetah_rotor.urdf")
        foot_name = "foot"
        penalize_contacts_on = ["shank"]
        terminate_after_contacts_on = ["base", "thigh"]
        collapse_fixed_joints = False
        fix_base_link = False
        self_collisions = 1  # 1 to disable, 0 to enable...bitwise filter
        flip_visual_attachments = False
        disable_gravity = False
        disable_motors = False  # all torques set to 0

    class reward_settings(MiniCheetahCfg.reward_settings):
        soft_dof_pos_limit = 0.9
        soft_dof_vel_limit = 0.9
        soft_torque_limit = 0.9
        max_contact_force = 600.
        base_height_target = BASE_HEIGHT_REF
        tracking_sigma = 0.3

    class scaling(MiniCheetahCfg.scaling):
        base_ang_vel = 3.14/(BASE_HEIGHT_REF/9.81)**0.5
        base_lin_vel = BASE_HEIGHT_REF
        dof_vel = 100.  # ought to be roughly max expected speed.
        base_height = BASE_HEIGHT_REF
        dof_pos = 4*[0.1, 1., 2]  # hip-abad, hip-pitch, knee
        dof_pos_obs = dof_pos
        dof_pos_target = 0.75  # action_scale
        tau_ff = 4*[18, 18, 28]  # hip-abad, hip-pitch, knee
        commands = [base_lin_vel, base_lin_vel, base_ang_vel]


class MiniCheetahRefRunnerCfg(MiniCheetahRunnerCfg):
    seed = -1

    class policy(MiniCheetahRunnerCfg.policy):
        actor_hidden_dims = [256, 256, 128]
        critic_hidden_dims = [256, 256, 128]
        # * can be elu, relu, selu, crelu, lrelu, tanh, sigmoid
        activation = 'elu'

        actor_obs = ["base_height",
                     "base_lin_vel",
                     "base_ang_vel",
                     "projected_gravity",
                     "commands",
                     "dof_pos_obs",
                     "dof_vel",
                     "phase_obs"
                     ]

        critic_obs = actor_obs

        actions = ["dof_pos_target"]

        class noise:
            dof_pos_obs = 0.  # 0.005  # can be made very low
            dof_vel = 0.  # 0.005
            ang_vel = 0.  # [0.1, 0.1, 0.1]  # 0.027, 0.14, 0.37
            base_ang_vel = 0.  # 0.
            dof_pos = 0.  # 0.005
            dof_vel = 0.  # 0.005
            lin_vel = 0.  # 0.
            ang_vel = 0.  # [0.3, 0.15, 0.4]
            gravity_vec = 0.  # 0.05

        class reward:
            class weights:
                tracking_lin_vel = 4.0
                tracking_ang_vel = 1.0
                lin_vel_z = 0.
                ang_vel_xy = 0.
                orientation = 1.75
                torques = 5.e-7
                dof_vel = 0.
                min_base_height = 1.5
                collision = 0.25
                action_rate = 0.01
                action_rate2 = 0.001
                stand_still = 0.
                dof_pos_limits = 0.
                feet_contact_forces = 0.
                dof_near_home = 0.
                reference_traj = 0.5
                swing_grf = 3.
                stance_grf = 3.

            class termination_weight:
                termination = 0.15

    class algorithm(MiniCheetahRunnerCfg.algorithm):
        # training params
        value_loss_coef = 1.0
        use_clipped_value_loss = True
        clip_param = 0.2
        entropy_coef = 0.01
        num_learning_epochs = 6
        # mini batch size = num_envs*nsteps/nminibatches
        num_mini_batches = 6
        learning_rate = 5.e-5
        schedule = 'adaptive'  # can be adaptive, fixed
        discount_horizon = 1.  # [s]
        GAE_bootstrap_horizon = 1.  # [s]
        desired_kl = 0.01
        max_grad_norm = 1.

    class runner(MiniCheetahRunnerCfg.runner):
        run_name = ''
        experiment_name = 'mini_cheetah_ref'
        max_iterations = 1000  # number of policy updates
        algorithm_class_name = 'PPO'
        num_steps_per_env = 32

gym/envs/mini_cheetah/mini_cheetah_ref_config.py
mit-biomimetics/ORCAgym
def adjust_settings(toggle, env_cfg, train_cfg):
    # * settings
    if train_cfg.runner.experiment_name =='':
        train_cfg.runner.experiment_name = "ORC_" + toggle
    else:
        train_cfg.runner.experiment_name = "ORC_" + toggle + '_' \
                                            +train_cfg.runner.experiment_name
    train_cfg.runner.run_name = "ORC_"+toggle
    # task_registry.set_log_dir_name(train_cfg)
    toggle = [x for x in toggle]
    if toggle[0] == '0':
        # * No osc observation
        train_cfg.policy.actor_obs.remove('oscillator_obs')
        train_cfg.policy.critic_obs.remove('oscillator_obs')
        train_cfg.policy.critic_obs.remove('oscillators_vel')
    if toggle[1] == '0':
        # * non reward
        train_cfg.policy.reward.weights.swing_grf = 0.
        train_cfg.policy.reward.weights.stance_grf = 0.
    if toggle[2] == '0':
        # * no coupling
        env_cfg.osc.coupling = 0.
        env_cfg.osc.coupling_range = [0., 0.]
        coupling_stop = 0.
        coupling_step = 0.
        coupling_slope = 0.
        env_cfg.osc.coupling_max = 0.

    return env_cfg, train_cfg
gym/scripts/ORC.py
mit-biomimetics/ORCAgym
import numpy as np
import matplotlib
import matplotlib.pyplot as plt
import os
import glob
from entro import sliding_window_entropy_multi_trial
import osc_plotters as oscplt

def get_run_names(experiment_name, base_path, data_folder):
    experiment_path = os.path.join(base_path, data_folder, experiment_name)
    return [os.path.basename(f).replace('_data.npz', '') for f in glob.glob(os.path.join(experiment_path, '*_data.npz'))]

# Initialize variables to store global min and max for each variable
global_min = None
global_max = None

all_toggles = ['111', '110', '100', '000', '001', '010', '011', '101']
# colors = ['lightsteelblue', 'yellowgreen', 'indianred', 'crimson']
# linestyle_dict = {'0': (0, (3, 1, 1, 1, 1, 1)), '1': 'solid'}
# linewidth_dict = {'0': 2, '1': 3}
plotting_options = {

    '000': {'color': 'cornflowerblue',
            'linestyle': (0, (3, 1, 1, 1, 1, 1)),
            'linewidth': 2},
    '100': {'color': 'yellowgreen',
            'linestyle': 'solid',
            'linewidth': 3},
    '110': {'color': 'goldenrod',
            'linestyle': (0, (3, 1, 1, 1, 1, 1)),
            'linewidth': 2},
    '111': {'color': 'crimson',
            'linestyle': 'solid',
            'linewidth': 3},
}



base_path = '/home/heim/Repos/trigym/gym/scripts/'
data_folder = 'FS_data'
obs2use = ['base_lin_vel', 'base_ang_vel', 'base_height']  # , 'commands']

# Changed from a tuple key to a dictionary of dictionaries
data_dict = {toggle: {} for toggle in all_toggles}

# Load all data and find global min and max for each variable
for toggle in all_toggles:
    experiment_name = "ORC_" + toggle
    run_names = get_run_names(experiment_name, base_path, data_folder)

    for run_name in run_names:
        data_path = os.path.join(base_path, data_folder, 'ORC_'+toggle, f"{run_name}_data.npz")
        with np.load(data_path, allow_pickle=True) as loaded_data:
            if len(obs2use) == 1:
                data = loaded_data[obs2use[0]]
            else:
                data = np.concatenate(list(loaded_data[name] for name in obs2use), axis=2)

        # Store in a nested dictionary
        data_dict[toggle][run_name] = data

        if global_min is None:
            global_min = np.min(data, axis=(0, 1))
            global_max = np.max(data, axis=(0, 1))
        else:
            global_min = np.minimum(global_min, np.min(data, axis=(0, 1)))
            global_max = np.maximum(global_max, np.max(data, axis=(0, 1)))

# Remaining code stays the same, with adjustments to use data_dict
# Create common bin edges based on global min and max
num_bins = [10, 10, 10]  # Update as needed
num_bins = [10] * data.shape[2]
custom_bin_edges = [np.linspace(global_min[i], global_max[i], num_bins[i] + 1)
                    for i in range(len(global_min))]

fig, ax = plt.subplots(1, 1, figsize=(12, 8))
window_size = 10  # Update as needed
step_size = 4  # Update as needed
sample_rate = 100/step_size

all_entropy_values = {toggle: [] for toggle in all_toggles}

# Loop through the toggles again for plotting
idx = 0
for toggle in all_toggles:
    for run_name, data in data_dict[toggle].items():

        assert len(num_bins) == data.shape[2], "Check your bins"

        entropy_values_multi_trial = sliding_window_entropy_multi_trial(
            data, window_size, step_size, 'digitize', custom_bin_edges)

        options = plotting_options[toggle]
        oscplt.entropy(ax, entropy_values_multi_trial,
                       f'Entropy for pushball perturbation',
                       color=options['color'],
                       linestyle=options['linestyle'],
                       linewidth=options['linewidth'],
                       sample_rate=sample_rate)

        all_entropy_values[toggle].append(entropy_values_multi_trial)
        idx += 1
        plt.tight_layout(rect=[0, 0, 0.75, 1])

# save all entropies to disk
save_path = os.path.join(base_path, f"all_entropy_values_ORC.npz")
np.savez(save_path, **all_entropy_values)

# Manually set the position of the axes to make room for the legend
ax.set_position([0.1, 0.1, 0.6, 0.8])  # [left, bottom, width, height]

# Place legend outside the plot and make it 2 columns
legend_handles = [matplotlib.lines.Line2D([0], [0],
                                          color=opts['color'],
                                          linestyle=opts['linestyle'],
                                          linewidth=opts['linewidth'])
                  for opts in plotting_options.values()]
legend = ax.legend(handles=legend_handles,
                   labels=['ORC_'+toggle for toggle in plotting_options.keys()],
                   loc='upper left', bbox_to_anchor=(1, 1), ncol=2)

observations_text = "Observations used:\n" + "\n".join(obs2use)
fig.text(0.73, 0.6, observations_text, ha='left', va='top',
         bbox={'facecolor': 'white', 'edgecolor': 'grey',
               'boxstyle': 'round, pad=0.5'})  # , fontsize=10)

# Manually set the position of the axes to make room for the legend
ax.set_position([0.1, 0.1, 0.6, 0.8])  # [left, bottom, width, height]

# Adjust layout to not overlap with the legend
# plt.subplots_adjust(right=0.7)  # leaves room for the legend
plt.subplots_adjust(right=0.6, left=0.1, top=0.9, bottom=0.1)

# Save and show the plot
save_path = os.path.join(base_path, f"entropy_plot_all_ORC.png")
# plt.savefig(save_path)
plt.savefig(save_path, bbox_extra_artists=(legend,), bbox_inches='tight')
print(f"Saved plot to {save_path}")
plt.show()
gym/scripts/ORC_entropy_analysis.py
mit-biomimetics/ORCAgym
import os

from gym.envs import __init__
from gym import LEGGED_GYM_ROOT_DIR
from gym.utils import get_args, task_registry
from gym.utils import KeyboardInterface, GamepadInterface
from ORC import adjust_settings
# torch needs to be imported after isaacgym imports in local source
import torch
from torch.multiprocessing import Process
from torch.multiprocessing import set_start_method
import numpy as np
import glob
import imageio
from datetime import datetime

RECORD_VID = False

def get_run_names(experiment_name):
    experiment_path = os.path.join(LEGGED_GYM_ROOT_DIR, 'logs',
                                   experiment_name)
    return [folder_name for folder_name in os.listdir(experiment_path)
            if os.path.isdir(os.path.join(experiment_path, folder_name))]

def create_logging_dict(runner, test_total_timesteps):
    # states_to_log = ['dof_pos_target',
    #                  'dof_pos_obs',
    #                  'dof_vel',
    #                  'torques',
    #                  'commands',
    #                  'base_lin_vel',
    #                  'base_ang_vel',
    #                  'oscillators',
    #                  'grf',
    #                  'base_height']
    states_to_log = [
                     'commands',
                     'base_lin_vel',
                     'base_ang_vel',
                     'oscillators',
                     'base_height']
    states_to_log_dict = {}

    for state in states_to_log:
        array_dim = runner.get_obs_size([state, ])
        states_to_log_dict[state] = torch.zeros((runner.env.num_envs,
                                                 test_total_timesteps,
                                                 array_dim),
                                                device=runner.env.device)
    return states_to_log, states_to_log_dict

def setup(toggle, run_name):
    args = get_args()
    env_cfg, train_cfg = task_registry.create_cfgs(args)
    env_cfg.env.num_envs = 1800
    if hasattr(env_cfg, "push_robots"):
        env_cfg.push_robots.toggle = False
    env_cfg.commands.resampling_time = 999999
    env_cfg.env.episode_length_s = 100.
    env_cfg.init_state.timeout_reset_ratio = 1.
    env_cfg.domain_rand.randomize_base_mass = False
    env_cfg.domain_rand.randomize_friction = False
    env_cfg.terrain.mesh_type = "plane"
    env_cfg.osc.init_to = 'random'
    env_cfg.osc.init_w_offset = False
    env_cfg.osc.process_noise = 0.
    env_cfg.osc.omega_var = 0.
    env_cfg.osc.coupling_var = 0.
    env_cfg.commands.ranges.lin_vel_x = [0., 0.]
    env_cfg.commands.ranges.lin_vel_y = 0.
    env_cfg.commands.ranges.yaw_vel = 0.
    env_cfg.commands.var = 0.
    env_cfg.env.env_spacing = 0.0
    train_cfg.policy.noise.scale = 0.0
    train_cfg.runner.run_name = run_name
    train_cfg.runner.load_run = run_name
    env_cfg, train_cfg = adjust_settings(toggle=toggle,
                                         env_cfg=env_cfg,
                                         train_cfg=train_cfg)
    env_cfg.init_state.reset_mode = "reset_to_basic"
    task_registry.set_log_dir_name(train_cfg)

    task_registry.make_gym_and_sim()
    env = task_registry.make_env(args.task, env_cfg)
    train_cfg.runner.resume = True
    runner = task_registry.make_alg_runner(env, train_cfg)

    # * switch to evaluation mode (dropout for example)
    runner.switch_to_eval()
    # ... [rest of your setup code]
    return env, runner, train_cfg


def play(env, runner, train_cfg, png_folder):
    protocol_name = 'pushball_run'+train_cfg.runner.run_name

    # * set up logging
    # log_file_path = os.path.join(LEGGED_GYM_ROOT_DIR, 'gym', 'scripts',
    #                              "osc_data", train_cfg.runner.run_name,
    #                              protocol_name+".npz")
    # if not os.path.exists(os.path.dirname(log_file_path)):
    #     os.makedirs(os.path.dirname(log_file_path))

    push_interval = 0.5
    stagger_interval = 0.01
    num_staggers = int(push_interval/stagger_interval)
    stagger_timesteps = int(stagger_interval/env.cfg.control.ctrl_dt)
    n_directions = 36
    n_trials = num_staggers*n_directions  # ! this must match num_envs
    # print(f"number of trials: {n_trials}")
    # print(f"number of loaded environments: {env.num_envs}")
    assert n_trials == env.num_envs, f"number of trials: {n_trials}, number of loaded environments: {env.num_envs}"
    env.commands[:, 0] = 3.
    push_magnitude = 3.0
    angles = torch.linspace(0, 2*np.pi, n_directions, device=env.device )
    push_ball = torch.stack((torch.cos(angles),
                             torch.sin(angles),
                             torch.zeros_like(angles)), dim=1)*push_magnitude

    kdx = 0  # trial counter
    img_idx = 0

    entrainment_time = 5.

    test_start_time = 1.0
    test_total_time = 9.0 + test_start_time
    test_total_timesteps = int(test_total_time/env.cfg.control.ctrl_dt)

    states_to_log, states_to_log_dict = create_logging_dict(runner,
                                                        test_total_timesteps)

    for t in range(int(entrainment_time/env.cfg.control.ctrl_dt)):
        # * sim step
        runner.set_actions(runner.get_inference_actions())
        env.step()

    for t in range(int(test_total_time/env.cfg.control.ctrl_dt)):

        # * record pngs
        if RECORD_VID and t % 5:
            filename = os.path.join(png_folder,
                                    f"{img_idx}.png")
            env.gym.write_viewer_image_to_file(env.viewer, filename)
            img_idx += 1

        # * do your thing here
        if t == int(test_start_time/env.cfg.control.ctrl_dt) + stagger_timesteps*kdx:
            env_ids = torch.arange(kdx*n_directions,
                                   (kdx+1)*n_directions,
                                   dtype=torch.int64,
                                   device=env.device)
            env.perturb_base_velocity(push_ball, env_ids)
            kdx += 1
        if kdx == num_staggers:  # turn off
            kdx = 0
        # * sim step
        runner.set_actions(runner.get_inference_actions())
        env.step()

        # * log
        for state in states_to_log:
            states_to_log_dict[state][:, t, :] = getattr(env, state)

    # * save data
    return states_to_log_dict
    # np.savez_compressed(log_file_path, **states_to_log_dict)
    # print("saved to ", log_file_path)


def save_gif(png_folder, gif_folder, frame_rate):
    png_files = sorted(glob.glob(f"{png_folder}/*.png"),
                       key=os.path.getmtime)
    images = [imageio.imread(f) for f in png_files]
    gif_path = os.path.join(gif_folder, 'output.gif')
    imageio.mimsave(gif_path, images, fps=frame_rate)


def worker(toggle, run_name):
    with torch.no_grad():
        env, runner, train_cfg = setup(toggle, run_name)
        # ... [rest of your worker code]

        # Adjust log_file_path to include run_name
        log_file_path = os.path.join(LEGGED_GYM_ROOT_DIR, 'gym', 'scripts',
                                     "FS_data", train_cfg.runner.run_name,
                                     run_name+"_data.npz")

        os.makedirs(os.path.dirname(log_file_path), exist_ok=True)

        png_folder = os.path.join(LEGGED_GYM_ROOT_DIR, 'temp', run_name)
        os.makedirs(png_folder, exist_ok=True)

        states_to_log_dict = play(env, runner=runner, train_cfg=train_cfg,
                                  png_folder=png_folder)

        states_to_log_dict_cpu = {k: v.detach().cpu().numpy()
                                  for k, v in states_to_log_dict.items()}

        np.savez_compressed(log_file_path, **states_to_log_dict_cpu)

        gif_folder = os.path.join(LEGGED_GYM_ROOT_DIR, 'gym', 'scripts',
                                  "osc_data", train_cfg.runner.run_name,
                                  run_name)
        if RECORD_VID:
            save_gif(png_folder, gif_folder, frame_rate=20)


if __name__ == '__main__':
    set_start_method('spawn')
    # all_toggles = ['111']
    all_toggles = ['000', '111', '110', '100', '101', '001', '010', '011']
    experiment_name = "ORC_" + all_toggles[0] + "_FullSend"
    run_names = get_run_names(experiment_name)
    worker(all_toggles[0], run_names[0])


    for toggle in all_toggles:
        experiment_name = "ORC_" + toggle+ "_FullSend"
        run_names = get_run_names(experiment_name)

        for run_name in run_names:
            p = Process(target=worker, args=(toggle, run_name))
            p.start()
            p.join()

            # Free up any resources if needed
            p.terminate()
            p.close()

gym/scripts/ORC_protocol_pushball.py
mit-biomimetics/ORCAgym
import os

from gym.envs import __init__
from gym import LEGGED_GYM_ROOT_DIR
from gym.utils import get_args, task_registry
from gym.utils import KeyboardInterface, GamepadInterface
from ORC import *
# torch needs to be imported after isaacgym imports in local source
import torch
from torch.multiprocessing import Process
from torch.multiprocessing import set_start_method
import numpy as np
import glob
import imageio
from datetime import datetime

RECORD_VID = False

def get_run_names(experiment_name):
    experiment_path = os.path.join(LEGGED_GYM_ROOT_DIR, 'logs',
                                   experiment_name)
    return [folder_name for folder_name in os.listdir(experiment_path)
            if os.path.isdir(os.path.join(experiment_path, folder_name))]

def create_logging_dict(runner, test_total_timesteps):
    # states_to_log = ['dof_pos_target',
    #                  'dof_pos_obs',
    #                  'dof_vel',
    #                  'torques',
    #                  'commands',
    #                  'base_lin_vel',
    #                  'base_ang_vel',
    #                  'oscillators',
    #                  'grf',
    #                  'base_height']
    states_to_log = [
                     'commands',
                     'base_lin_vel',
                     'base_ang_vel',
                     'oscillators',
                     'base_height']
    states_to_log_dict = {}

    for state in states_to_log:
        array_dim = runner.get_obs_size([state, ])
        states_to_log_dict[state] = torch.zeros((runner.env.num_envs,
                                                 test_total_timesteps,
                                                 array_dim),
                                                device=runner.env.device)
    return states_to_log, states_to_log_dict

def setup(toggle, run_name):
    args = get_args()
    env_cfg, train_cfg = task_registry.create_cfgs(args)
    env_cfg.env.num_envs = 1800
    if hasattr(env_cfg, "push_robots"):
        env_cfg.push_robots.toggle = False
    env_cfg.commands.resampling_time = 999999
    env_cfg.env.episode_length_s = 100.
    env_cfg.init_state.timeout_reset_ratio = 1.
    env_cfg.domain_rand.randomize_base_mass = False
    env_cfg.domain_rand.randomize_friction = False
    env_cfg.terrain.mesh_type = "plane"
    env_cfg.osc.init_to = 'random'
    env_cfg.osc.init_w_offset = False
    env_cfg.osc.process_noise = 0.
    env_cfg.osc.omega_var = 0.
    env_cfg.osc.coupling_var = 0.
    env_cfg.commands.ranges.lin_vel_x = [0., 0.]
    env_cfg.commands.ranges.lin_vel_y = 0.
    env_cfg.commands.ranges.yaw_vel = 0.
    env_cfg.commands.var = 0.
    env_cfg.env.env_spacing = 3.0
    train_cfg.policy.noise.scale = 0.0
    train_cfg.runner.run_name = run_name
    train_cfg.runner.load_run = run_name
    env_cfg, train_cfg = adjust_settings(toggle=toggle,
                                         env_cfg=env_cfg,
                                         train_cfg=train_cfg)
    env_cfg.init_state.reset_mode = "reset_to_basic"
    task_registry.set_log_dir_name(train_cfg)

    task_registry.make_gym_and_sim()
    env = task_registry.make_env(args.task, env_cfg)
    train_cfg.runner.resume = True
    runner = task_registry.make_alg_runner(env, train_cfg)

    # * switch to evaluation mode (dropout for example)
    runner.switch_to_eval()
    # ... [rest of your setup code]
    return env, runner, train_cfg


def play(env, runner, train_cfg, log_file_path, run_name):
    push_mags = torch.arange(0.5, 4., 0.5)

    for push_mag in push_mags:
        push_interval = 0.5
        stagger_interval = 0.01
        num_staggers = int(push_interval/stagger_interval)
        stagger_timesteps = int(stagger_interval/env.cfg.control.ctrl_dt)
        n_directions = 36
        n_trials = num_staggers*n_directions  # ! this must match num_envs
        # print(f"number of trials: {n_trials}")
        # print(f"number of loaded environments: {env.num_envs}")
        assert n_trials == env.num_envs, f"number of trials: {n_trials}, number of loaded environments: {env.num_envs}"
        env.commands[:, 0] = 3.
        angles = torch.linspace(0, 2*np.pi, n_directions, device=env.device )
        push_ball = torch.stack((torch.cos(angles),
                                torch.sin(angles),
                                torch.zeros_like(angles)), dim=1)*push_mag

        states_to_log_dict = run_protocol(env, runner, push_ball,
                                          num_staggers, stagger_timesteps,
                                          n_directions)

        # * save data
        states_to_log_dict_cpu = {k: v.detach().cpu().numpy()
                                    for k, v in states_to_log_dict.items()}

        save_path = os.path.join(log_file_path, f"{push_mag:.1f}")
        os.makedirs(save_path, exist_ok=True)
        save_name = os.path.join(save_path, f"{run_name}_data.npz")
        np.savez_compressed(save_name, **states_to_log_dict_cpu)
        env.reset()

def run_protocol(env, runner, push_ball, num_staggers, stagger_timesteps,
                 n_directions):

    kdx = 0  # trial counter
    entrainment_time = 5.

    test_start_time = 1.0
    test_total_time = 9.0 + test_start_time
    test_total_timesteps = int(test_total_time/env.cfg.control.ctrl_dt)

    states_to_log, states_to_log_dict = create_logging_dict(runner,
                                                        test_total_timesteps)

    for t in range(int(entrainment_time/env.cfg.control.ctrl_dt)):
        # * sim step
        runner.set_actions(runner.get_inference_actions())
        env.step()

    for t in range(int(test_total_time/env.cfg.control.ctrl_dt)):

        # * do your thing here
        if t == int(test_start_time/env.cfg.control.ctrl_dt) + stagger_timesteps*kdx:
            env_ids = torch.arange(kdx*n_directions,
                                   (kdx+1)*n_directions,
                                   dtype=torch.int64,
                                   device=env.device)
            env.perturb_base_velocity(push_ball, env_ids)
            kdx += 1
        if kdx == num_staggers:  # turn off
            kdx = 0
        # * sim step
        runner.set_actions(runner.get_inference_actions())
        env.step()

        # * log
        for state in states_to_log:
            states_to_log_dict[state][:, t, :] = getattr(env, state)

    return states_to_log_dict


def worker(toggle, run_name):
    with torch.no_grad():
        env, runner, train_cfg = setup(toggle, run_name)
        # ... [rest of your worker code]

        # Adjust log_file_path to include run_name
        log_file_path = os.path.join(LEGGED_GYM_ROOT_DIR, 'gym', 'scripts',
                                     "FS_data", train_cfg.runner.run_name)

        os.makedirs(os.path.dirname(log_file_path), exist_ok=True)

        png_folder = os.path.join(LEGGED_GYM_ROOT_DIR, 'temp', run_name)
        os.makedirs(png_folder, exist_ok=True)

        play(env, runner=runner, train_cfg=train_cfg,
             log_file_path=log_file_path, run_name=run_name)


if __name__ == '__main__':
    set_start_method('spawn')
    # all_toggles = ['111']
    all_toggles = ['000', '111', '110', '100', '101', '001', '010', '011']
    experiment_name = "ORC_" + all_toggles[0] + "_FullSend"
    run_names = get_run_names(experiment_name)
    worker(all_toggles[0], run_names[0])


    for toggle in all_toggles:
        experiment_name = "ORC_" + toggle+ "_FullSend"
        run_names = get_run_names(experiment_name)

        for run_name in run_names:
            p = Process(target=worker, args=(toggle, run_name))
            p.start()
            p.join()

            # Free up any resources if needed
            p.terminate()
            p.close()

gym/scripts/ORC_protocol_pushball_sweep.py
mit-biomimetics/ORCAgym
import numpy as np
import matplotlib.pyplot as plt
import os
from entro import sliding_window_entropy_multi_trial
import osc_plotters as oscplt
import gc
import glob

def get_run_names(experiment_name, base_path):
    experiment_path = os.path.join(base_path, data_folder, experiment_name)
    return [os.path.basename(f).replace('_data.npz', '') for f in glob.glob(os.path.join(experiment_path, '*_data.npz'))]


# Setup
all_toggles = ['000', '001', '010', '011', '100', '101', '110', '111']
colors = ['lightsteelblue', 'mediumaquamarine', 'yellowgreen', 'forestgreen', 'sandybrown', 'peru', 'indianred', 'crimson']
toggle_settings = dict(zip(all_toggles, colors))
base_path = '/home/heim/Repos/trigym/gym/scripts/'
data_folder = 'FS_data'  # Adjust to your actual data folder name
obs2use = ['base_lin_vel', 'base_ang_vel', 'commands']

# Initialize data dictionary
data_dict = {}

fig, ax = plt.subplots()
# Load data

fail_rate = {}

for toggle in all_toggles:
    experiment_name = "ORC_" + toggle
    run_names = get_run_names(experiment_name, base_path)

    fail_rate[toggle] = []

    for run_name in run_names:
        data_path = os.path.join(base_path, data_folder, 'ORC_'+toggle, f"{run_name}_data.npz")
        with np.load(data_path, allow_pickle=True) as loaded_data:
            data_dict = {key: loaded_data[key] for key in obs2use}
        fails = np.sum(data_dict['commands'][:, -2, 0] < 0.1)
    # * plotting and analysis
        # print(f"Toggle: {toggle}, Run: {run_name}, Failed: {fails}")

        ax.plot(data_dict['base_lin_vel'][:, :, 0].T,
                linewidth=1,
                color=toggle_settings[toggle],
                alpha=0.5)

        os.makedirs("FS_plots", exist_ok=True)
        plt.savefig(os.path.join("FS_plots", f"{run_name}.png"))
        fail_rate[toggle].append(fails)

        ax.clear()

        # # Explicitly delete the large objects to free memory
        # del data_dict
        # del loaded_data
        # gc.collect()  # Run the garbage collector
    print(' ')

for dataset_name, fails in fail_rate.items():
    a = np.array(fails)
    mean_val = a.mean()
    std_val = a.std()
    print(f"{dataset_name}: {mean_val:.4f} +/- {std_val:.4f}")



# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# # Plotting
# fig, ax = plt.subplots()
# legend_handles = []

# for (toggle, run_name), data in data_dict.items():
#     print(f"Toggle: {toggle}, Run: {run_name}, Failed: {np.sum(data['commands'][:, -2, 0] < 1)}")

#     ax.plot(data['base_lin_vel'][:, :, 0].T,
#             linewidth=1,
#             color=toggle_settings[toggle],
#             alpha=0.5)
#     plt.savefig(f"{run_name}.png")


gym/scripts/ORC_pushball_analysis.py
mit-biomimetics/ORCAgym
import numpy as np
import matplotlib.pyplot as plt
import os
from entro import sliding_window_entropy_multi_trial
import osc_plotters as oscplt
import gc
import glob

def get_run_names(experiment_name, base_path):
    experiment_path = os.path.join(base_path, data_folder, experiment_name)
    return [os.path.basename(f).replace('_data.npz', '') for f in glob.glob(os.path.join(experiment_path, '*_data.npz'))]


# Setup
all_toggles = ["000", "001", "100", "101", "010", "110", "011", "111"]

all_pushmags = ['0.5', '1.0', '1.5', '2.0', '2.5', '3.0', '3.5']

colors = ['lightsteelblue', 'mediumaquamarine', 'yellowgreen', 'forestgreen', 'sandybrown', 'peru', 'indianred', 'crimson']
toggle_settings = dict(zip(all_toggles, colors))
base_path = '/home/heim/Repos/trigym/gym/scripts/'
data_folder = 'FS_data'  # Adjust to your actual data folder name
obs2use = ['base_lin_vel', 'base_ang_vel', 'commands']

# Initialize data dictionary
data_dict = {}

fig, ax = plt.subplots()
# Load data

# set up a dict of dicts, containing fail rates for each toggle, and each pushmag
fail_rate = {toggle: {pushmag: [] for pushmag in all_pushmags}
                 for toggle in all_toggles}

for toggle in all_toggles:
    for pushmag in all_pushmags:
        subpath_name = f"ORC_{toggle}"
        experiment_name = pushmag
        save_name = subpath_name + "_" + experiment_name
        run_names = get_run_names(os.path.join(subpath_name, experiment_name),
                                  base_path)

        for run_name in run_names:
            data_path = os.path.join(base_path, data_folder,
                                     subpath_name, experiment_name,
                                     run_name+"_data.npz")
            with np.load(data_path, allow_pickle=True) as loaded_data:
                # data_dict = {key: loaded_data[key] for key in obs2use}
                fails = np.sum(loaded_data['commands'][:, -2, 0] < 0.1)
        # * plotting and analysis
            # print(f"Toggle: {toggle}, Run: {run_name}, Failed: {fails}")

                # ax.plot(loaded_data['base_lin_vel'][:, :, 0].T,
                #         linewidth=1,
                #         color=toggle_settings[toggle],
                #         alpha=0.5)

                # os.makedirs("FS_plots", exist_ok=True)
                # plt.savefig(os.path.join("FS_plots", f"{save_name}.png"))
                fail_rate[toggle][pushmag].append(fails)

            ax.clear()
        # print(' ')

num_envs = 1800
for dataset_name, fails_for_dataset in fail_rate.items():
    print(f"{dataset_name}")
    for pushmag, fails in fails_for_dataset.items():
        a = np.array(fails)
        mean_val = a.mean()/num_envs*100
        std_val = a.std()/num_envs*100
        print(f"{pushmag}: {mean_val:.4f}% +/- {std_val:.4f}%")
    print(' ')

# * let's also save the entire fail_rate dict, so it is easier to load later
np.savez_compressed(os.path.join("FS_plots", "fail_rate.npz"), **fail_rate)
gym/scripts/ORC_pushball_analysis_sweep.py
mit-biomimetics/ORCAgym
import wandb


# download the specified WandB artifact to the artifacts parent directory
def download_artifact_code(artifact_name):
    run = wandb.init()
    artifact = \
        run.use_artifact(artifact_name, type='code')
    artifact.download('../../artifacts/'+artifact_name)


# the name WandB name of the artifact to download
artifact_name = 'ajm4/logging_test/source-367gxqhp:v0'
download_artifact_code(artifact_name)

gym/scripts/download.py
mit-biomimetics/ORCAgym

import numpy as np
import matplotlib.pyplot as plt
import scipy.stats as stats


def compute_entropy(counts, method='manual'):
    if method == 'manual':
        total = np.sum(counts)
        probabilities = counts / total
        # Add a small constant to avoid log(0)
        log_probabilities = np.log(probabilities + 1e-9)
        return -np.sum(probabilities * log_probabilities)
    elif method == 'scipy':
        return stats.entropy(counts)
    else:
        raise ValueError("Invalid method. Choose 'manual' or 'scipy'.")

def get_counts(data, bin_edges, method='digitize'):
    M, D = data.shape
    if method == 'digitize':
        digi_data = [np.digitize(data[:, i], bins=bin_edges[i])
                     for i in range(D)]
        counts = np.unique(np.column_stack(digi_data), return_counts=True)[1]
    elif method == 'histogramdd':
        counts = np.histogramdd(data, density=True, bins=bin_edges)[0]
    else:
        'No appropriate method found'
        return None
    # NOTE: could use raveled using np.ravel_multi_index. Should be the same.
    return counts.flatten()

def sliding_window_entropy_single_trial(data, window_size, step_size, method, bin_edges):
    M, D = data.shape
    entropy_values = []

    for start in range(0, M - window_size + 1, step_size):
        window_data = data[start:start + window_size, :]
        flat_counts = get_counts(window_data, bin_edges, method=method)
        # flat_counts = counts.flatten()
        flat_counts = flat_counts[flat_counts > 0]  # * needed?
        window_entropy = compute_entropy(flat_counts, method='manual')
        entropy_values.append(window_entropy)
    return entropy_values

def sliding_window_entropy_multi_trial(data, window_size, step_size, method, bin_edges):
    N, M, D = data.shape
    entropy_values = []

    for start in range(0, M - window_size + 1, step_size):
        window_data = data[:, start:start + window_size, :].reshape((-1, D))
        flat_counts = get_counts(window_data, bin_edges, method=method)
        # flat_counts = counts.flatten()
        flat_counts = flat_counts[flat_counts > 0]  # * needed?
        window_entropy = compute_entropy(flat_counts, method='manual')
        entropy_values.append(window_entropy)
    return np.array(entropy_values)

if __name__ == "__main__":
    data_path = '/home/heim/Repos/trigym/gym/scripts/osc_data/ORC_111/noise_step_1.0.npz'
    # dimensions = [0, 1, 2]
    window_size = 72
    step_size = 1
    num_bins = 10
    method = 'digitize'  # Choose between 'digitize' and 'histogramdd'
    plot_trials = True

    loaded_data = np.load(data_path, allow_pickle=True)
    # data = np.concatenate((loaded_data['base_lin_vel'],
    #                     loaded_data['base_ang_vel']
    #                     ),
    #                    axis=2)
    data = loaded_data['base_lin_vel']

    # number of bins for each dimension in data
    num_bins = [10, 10, 10]
    if len(num_bins) == 1:
        num_bins = num_bins * data.shape[2]

    assert len(num_bins) == data.shape[2], "check your bin number"

    # ! needs to be set the same across all experiments
    custom_bin_edges = [np.linspace(np.min(data[:, :, i]), np.max(data[:, :, i]), num_bins[i] + 1) for i in range(data.shape[2])]
    dimensions = data.shape[2]

    all_trials_entropy_values = []

    for trial in range(data.shape[0]):
        single_trial_data = data[trial, :, :]
        # reshaped_single_trial_data = single_trial_data.reshape(-1, dimensions)

        entropy_values_single_trial = sliding_window_entropy_single_trial(
            single_trial_data, window_size, step_size, method, custom_bin_edges
        )
        all_trials_entropy_values.append(entropy_values_single_trial)

# * now analyze all bots together
    multi_trial_values = []
    entropy_values_multi_trial = sliding_window_entropy_multi_trial(
        data, window_size, step_size, method, custom_bin_edges)
    multi_trial_values.append(entropy_values_multi_trial)

    if plot_trials:
        # plt.figure(figsize=(14, 8))
        # for i, entropy_values in enumerate(all_trials_entropy_values):
        #     plt.plot(entropy_values, label=f'Trial {i+1}', linewidth=1)

        # plt.title(f'Sliding Window Entropy Over Time for All Trials (Window Size = {window_size})')
        # plt.xlabel('Time Step')
        # plt.ylabel('Entropy')
        # plt.legend(loc='upper right', ncol=5)
        # plt.ylim([0, 4])
        # plt.show()

        entropy_values = np.array(all_trials_entropy_values)
        n_trials = entropy_values.shape[0]  # Number of trials
        n_timesteps = entropy_values.shape[1]  # Number of time steps
        mean_data = np.mean(entropy_values, axis=0)  # shape: (D, )
        # Compute the standard error of the mean (SEM) to get the confidence interval
        # Note: You can adjust the confidence level as needed.
        sem_data = np.std(entropy_values, axis=0) / np.sqrt(n_trials)  # shape: (D, )

        # Create 95% confidence interval bounds
        confidence_level = 0.95
        z_value = 1.96  # z-value for 95% confidence (from the standard normal distribution)
        lower_bound = mean_data - z_value * sem_data
        upper_bound = mean_data + z_value * sem_data

        # Create a plot
        fig, ax = plt.subplots()
        ax.plot(mean_data, label='Mean', color='b')
        ax.fill_between(range(n_timesteps), lower_bound, upper_bound,
                        color='b', alpha=0.2, label='95% CI')

        ax.set_xlabel('Time Steps')
        ax.set_ylabel('Mean Value')
        ax.set_title('Mean and 95% Confidence Interval Over Time Steps')
        ax.legend()

        fig2, ax2 = plt.subplots()
        ax2.plot(entropy_values_multi_trial, linewidth=1)
        plt.xlabel('Time step')
        plt.ylabel('Entropy')
        plt.title('Entropy across multiple robots')
        # plt.ylim([0, 4])

        plt.show()

gym/scripts/entropy.py
mit-biomimetics/ORCAgym
import numpy as np

import matplotlib
# matplotlib.use("TkAgg")

import matplotlib.colors as colors
import matplotlib.cbook as cbook
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
import pickle
import os

matplotlib.rcParams['pdf.fonttype'] = 42
matplotlib.rcParams['ps.fonttype'] = 42
matplotlib.rcParams['mathtext.fontset'] = 'stix'
matplotlib.rcParams['font.family'] = 'STIXGeneral'

font = {'size': 24}
matplotlib.rc('font', **font)


# Function to plot phase vs phase velocity in a polar subplot with progressive colormap and legend
def plot_polar_phase(dataset, robot_idx, r_variable='grf', segment=None,
                              title=''):
    # n_points = len(ref_phase)

    def plot_individual_oscillator(ax, ref_phase, phase_velocity_data,
                                   oscillator_phase, title):
        n_points = len(ref_phase)
        # Create a progressive colormap for each cycle
        cmap = plt.cm.Set2
        num_cycles = np.sum(np.diff(oscillator_phase) < 0) + 1  # Count of phase wraps + 1
        colors = cmap(np.linspace(0, 1, num_cycles))

        start_idx = 0
        color_idx = 0

        # Loop through each point and plot with progressive color for each cycle
        for i in range(1, n_points):
            if oscillator_phase[i] < oscillator_phase[i-1]:  # Check for phase wrap to start new cycle
                ax.plot(ref_phase[start_idx:i+1], phase_velocity_data[start_idx:i+1], color=colors[color_idx % len(colors)], label=f"Cycle {color_idx+1}")
                color_idx += 1
                start_idx = i

        # Marking the zero-points and pi-crossings of the corresponding oscillator
        max_velocity = max(phase_velocity_data)
        zero_points = []
        pi_points = []
        for i in range(n_points):
            if oscillator_phase[i] < oscillator_phase[i-1]:
                zero_points.append(i)
            if oscillator_phase[i] > np.pi and oscillator_phase[i-1] < np.pi:
                pi_points.append(i)
        zero_points = np.array(zero_points)
        pi_points = np.array(pi_points)

        # * 
        # zero_points = np.where(np.diff(oscillator_phase) < 0)[0]
        # pi_points = np.where(np.abs(np.diff(oscillator_phase)) > np.pi)[0]
        for idx in zero_points:
            ax.plot([ref_phase[idx], ref_phase[idx]], [0, max_velocity],
                    color='red', alpha=0.5)
            ax.plot(ref_phase[idx], max_velocity, 'ro', markersize=5)
        for idx in pi_points:
            ax.plot([ref_phase[idx], ref_phase[idx]], [0, max_velocity],
                    color='blue', alpha=0.5)
            ax.plot(ref_phase[idx], max_velocity, 'bo', markersize=5)

        ax.set_title(title)
        ax.set_rlabel_position(-22.5)  # Move radial labels to avoid overlap with line
        ax.grid(True)
        ax.legend(loc="upper right", bbox_to_anchor=(1.3, 1.1))

    # * ##################
    # unpack dataset
    ref_phase = dataset['oscillators'][robot_idx][:, 0]
    # phase_velocity = dataset['oscillators_vel'][robot_idx][:, 0]

    if segment is None:
        start_idx = 0
        end_idx = len(ref_phase)
    else:
        start_idx = int(segment[0] * fs)
        end_idx = int(segment[1] * fs)

    # Extracting the recorded phase velocity dataset for all oscillators during the transition to 3 m/s
    ref_phase = ref_phase[start_idx:end_idx]

    # Combined plots for all oscillators in a 2x2 grid
    fig, axs = plt.subplots(2, 2, figsize=(16, 16),
                            subplot_kw={'projection': 'polar'})
    axs = axs.ravel()
    plt.title(title)
    # Plotting for each oscillator
    for i in range(4):
        oscillator_phase_i = dataset['oscillators'][robot_idx, start_idx:end_idx, i]
        r_variable_i = dataset[r_variable][robot_idx, start_idx:end_idx, i]
        plot_individual_oscillator(axs[i], ref_phase,
                                   r_variable_i,
                                   oscillator_phase_i,
                                   f"Oscillator {i+1}")

    plt.tight_layout()
    return fig, axs
    # plt.show()
    from scipy.stats import entropy

def compute_multi_var_shannon_entropy(data, variable_names, num_bins=50):
    """
    Compute multi-variable Shannon entropy for specified variables in the dataset.
    Parameters:
    - data: np.ndarray, the dataset
    - variable_names: list, names of the variables to include in the multi-variable X
    - num_bins: int, number of bins for histogram (used in entropy calculation)
    Returns:
    - entropies: np.array, Shannon entropy for the multi-variable X at each time step across all bots
    """
    # Initialize entropy array
    num_time_steps = data[variable_names[0]].shape[1]
    entropies = np.zeros(num_time_steps)
    for t in range(num_time_steps):
        # Concatenate specified variables across bots to form multi-variable X
        multi_var_X = np.concatenate([data[var][:, t] for var in variable_names], axis=1)
        # Flatten the multi-variable data for this time step
        flat_data = multi_var_X.flatten()
        # Compute histogram
        hist, _ = np.histogram(flat_data, bins=num_bins, density=True)
        # Compute Shannon entropy for this time step
        entropies[t] = entropy(hist)

    return entropies

def entropy_(ax, entropy, title, color='b', linestyle='-', linewidth=1,
            sample_rate=100):
    time = np.linspace(0, len(entropy) / sample_rate, len(entropy))
    ax.plot(time, entropy,
            linewidth=linewidth,
            linestyle=linestyle,
            color=color)

    ax.set_xlabel('Time (s)')  # Now in seconds
    ax.set_ylabel('Entropy (bits)')
    ax.set_title(title)

    ax.grid(True)  # Add grid lines for better readability

    # Optional: If you are plotting multiple lines, a legend can be helpful
    # ax.legend(loc='best')

def entropy(ax, entropy, title, label='', color='b', linestyle='-', linewidth=1, sample_rate=100):
    time = np.linspace(0, entropy.size / sample_rate, entropy.size)
    ax.plot(time, entropy,
            linewidth=linewidth,
            linestyle=linestyle,
            color=color,
            alpha=0.6,
            label=label)  # Add label to the plot for the legend

    ax.set_xlabel('Time (s)')  # Now in seconds
    ax.set_ylabel('Entropy (bits)')
    ax.set_title(title)
    ax.grid(True)  # Add grid lines for better readability

    # Place legend outside the plot and make it 2 columns
    # ax.legend(loc='upper left', bbox_to_anchor=(1, 1), ncol=2)

gym/scripts/osc_plotters.py
mit-biomimetics/ORCAgym
import os

from gym.envs import __init__
from gym import LEGGED_GYM_ROOT_DIR
from gym.utils import get_args, task_registry
from gym.utils import KeyboardInterface, GamepadInterface

# torch needs to be imported after isaacgym imports in local source
import torch


def setup(args):
    env_cfg, train_cfg = task_registry.create_cfgs(args)
    env_cfg.env.num_envs = min(env_cfg.env.num_envs, 16)
    if hasattr(env_cfg, "push_robots"):
        env_cfg.push_robots.toggle = False
    env_cfg.commands.resampling_time = 9999
    env_cfg.env.episode_length_s = 9999
    env_cfg.init_state.timeout_reset_ratio = 1.
    env_cfg.domain_rand.randomize_base_mass = True
    env_cfg.domain_rand.randomize_friction = True
    env_cfg.terrain.mesh_type = "plane"

    task_registry.make_gym_and_sim()
    env = task_registry.make_env(args.task, env_cfg)
    env.cfg.init_state.reset_mode = "reset_to_basic"
    train_cfg.runner.resume = True
    runner = task_registry.make_alg_runner(env, train_cfg)

    # * switch to evaluation mode (dropout for example)
    runner.switch_to_eval()
    if EXPORT_POLICY:
        path = os.path.join(LEGGED_GYM_ROOT_DIR, 'logs',
                            train_cfg.runner.experiment_name, 'exported')
        runner.export(path)
    return env, runner, train_cfg


def play(env, runner, train_cfg):

    # * set up interface: GamepadInterface(env) or KeyboardInterface(env)
    COMMANDS_INTERFACE = hasattr(env, "commands")
    if COMMANDS_INTERFACE:
        # interface = GamepadInterface(env)
        interface = KeyboardInterface(env)
    for i in range(10*int(env.max_episode_length)):
        if COMMANDS_INTERFACE:
            interface.update(env)
        runner.set_actions(runner.get_inference_actions())
        env.step()


if __name__ == '__main__':
    EXPORT_POLICY = True
    args = get_args()
    with torch.no_grad():
        env, runner, train_cfg = setup(args)
        play(env, runner, train_cfg)

gym/scripts/play.py
mit-biomimetics/ORCAgym
import os

from gym.envs import __init__
from gym import LEGGED_GYM_ROOT_DIR
from gym.utils import get_args, task_registry
from gym.utils import KeyboardInterface, GamepadInterface
from ORC import adjust_settings
# torch needs to be imported after isaacgym imports in local source
import torch


def setup(args):
    env_cfg, train_cfg = task_registry.create_cfgs(args)
    env_cfg.env.num_envs = 50
    if hasattr(env_cfg, "push_robots"):
        env_cfg.push_robots.toggle = False
    env_cfg.commands.resampling_time = 9999
    env_cfg.env.episode_length_s = 9999
    env_cfg.init_state.timeout_reset_ratio = 1.
    env_cfg.domain_rand.randomize_base_mass = False
    env_cfg.domain_rand.randomize_friction = False
    env_cfg.terrain.mesh_type = "plane"
    env_cfg.osc.init_to = 'random'
    env_cfg.osc.process_noise = 0.
    env_cfg.osc.omega_var = 0.
    env_cfg.osc.coupling_var = 0.
    env_cfg.commands.ranges.lin_vel_x = [0., 0.]
    env_cfg.commands.ranges.lin_vel_y = 0.
    env_cfg.commands.ranges.yaw_vel = 0.
    env_cfg.commands.var = 0.

    train_cfg.policy.noise.scale = 1.0

    env_cfg, train_cfg = adjust_settings(toggle=args.ORC_toggle,
                                         env_cfg=env_cfg,
                                         train_cfg=train_cfg)
    env_cfg.init_state.reset_mode = "reset_to_basic"
    task_registry.set_log_dir_name(train_cfg)

    task_registry.make_gym_and_sim()
    env = task_registry.make_env(args.task, env_cfg)
    train_cfg.runner.resume = True
    runner = task_registry.make_alg_runner(env, train_cfg)

    # * switch to evaluation mode (dropout for example)
    runner.switch_to_eval()
    if EXPORT_POLICY:
        path = os.path.join(LEGGED_GYM_ROOT_DIR, 'logs',
                            train_cfg.runner.experiment_name, 'exported')
        runner.export(path)
    return env, runner, train_cfg


def play(env, runner, train_cfg):
    RECORD_FRAMES = False
    # * set up interface: GamepadInterface(env) or KeyboardInterface(env)
    COMMANDS_INTERFACE = hasattr(env, "commands")
    if COMMANDS_INTERFACE:
        # interface = GamepadInterface(env)
        interface = KeyboardInterface(env)
    img_idx = 0

    for i in range(10*int(env.max_episode_length)):

        if RECORD_FRAMES:
            if i % 5:
                filename = os.path.join(LEGGED_GYM_ROOT_DIR,
                                        'gym', 'scripts', 'gifs',
                                        train_cfg.runner.experiment_name,
                                        f"{img_idx}.png")
                env.gym.write_viewer_image_to_file(env.viewer, filename)
                img_idx += 1

        if COMMANDS_INTERFACE:
            interface.update(env)
        runner.set_actions(runner.get_inference_actions())
        env.step()


if __name__ == '__main__':
    EXPORT_POLICY = True
    args = get_args()
    with torch.no_grad():
        env, runner, train_cfg = setup(args)
        play(env, runner, train_cfg)

gym/scripts/play_ORC.py
mit-biomimetics/ORCAgym
import os
import json
import wandb

from gym import LEGGED_GYM_ROOT_DIR
from gym.scripts.train import train, setup
from gym.utils import get_args
from gym.utils.logging_and_saving import wandb_singleton

# torch needs to be imported after isaacgym imports in local source
from torch.multiprocessing import Process
from torch.multiprocessing import set_start_method


def load_sweep_config(file_name):
    return json.load(open(os.path.join(
        LEGGED_GYM_ROOT_DIR, 'gym', 'scripts',
        'sweep_configs', file_name)))


def train_with_sweep_cfg():
    train_cfg, policy_runner = setup()
    train(train_cfg=train_cfg, policy_runner=policy_runner)


def sweep_wandb_mp():
    ''' start a new process for each train function '''

    p = Process(target=train_with_sweep_cfg, args=())
    p.start()
    p.join()
    p.kill()


def start_sweeps(args):
    # * required for multiprocessing CUDA workloads
    set_start_method('spawn')

    # * load sweep_config from JSON file
    if args.wandb_sweep_config is not None:
        sweep_config = load_sweep_config(args.wandb_sweep_config)
    else:
        sweep_config = load_sweep_config('sweep_config_example.json')
    # * set sweep_id if you have a previous id to use
    sweep_id = None
    if args.wandb_sweep_id is not None:
        sweep_id = args.wandb_sweep_id

    wandb_helper = wandb_singleton.WandbSingleton()
    wandb_helper.set_wandb_values(args)

    if wandb_helper.is_wandb_enabled():
        if sweep_id is None:
            sweep_id = wandb.sweep(
                sweep_config,
                entity=wandb_helper.get_entity_name(),
                project=wandb_helper.get_project_name())
        wandb.agent(
            sweep_id,
            sweep_wandb_mp,
            entity=wandb_helper.get_entity_name(),
            project=wandb_helper.get_project_name(),
            count=sweep_config['run_cap'])
    else:
        print('ERROR: No WandB project and entity provided for sweeping')


if __name__ == '__main__':
    args = get_args()
    os.environ[wandb.env.DIR] = os.path.join(LEGGED_GYM_ROOT_DIR, 'logs')
    start_sweeps(args)

gym/scripts/sweep.py
mit-biomimetics/ORCAgym
from gym.envs import __init__
from gym.utils import get_args, task_registry
from gym.utils.logging_and_saving \
    import local_code_save_helper, wandb_singleton


def setup():
    args = get_args()
    wandb_helper = wandb_singleton.WandbSingleton()

    # * prepare environment
    env_cfg, train_cfg = task_registry.create_cfgs(args)
    task_registry.make_gym_and_sim()
    wandb_helper.setup_wandb(env_cfg=env_cfg, train_cfg=train_cfg, args=args)
    env = task_registry.make_env(name=args.task, env_cfg=env_cfg)
    # * then make env
    policy_runner = task_registry.make_alg_runner(env, train_cfg)

    local_code_save_helper.log_and_save(
        env, env_cfg, train_cfg, policy_runner)
    wandb_helper.attach_runner(policy_runner=policy_runner)

    return train_cfg, policy_runner


def train(train_cfg, policy_runner):
    wandb_helper = wandb_singleton.WandbSingleton()

    policy_runner.learn(
        num_learning_iterations=train_cfg.runner.max_iterations,
        init_at_random_ep_len=True)

    wandb_helper.close_wandb()


if __name__ == '__main__':
    train_cfg, policy_runner = setup()
    train(train_cfg=train_cfg, policy_runner=policy_runner)

gym/scripts/train.py
mit-biomimetics/ORCAgym
from gym.envs import __init__
from gym.utils import get_args, task_registry
from gym.utils.logging_and_saving \
    import local_code_save_helper, wandb_singleton
from ORC import adjust_settings

def setup():
    args = get_args()
    wandb_helper = wandb_singleton.WandbSingleton()

    # * prepare environment
    env_cfg, train_cfg = task_registry.create_cfgs(args)
    env_cfg, train_cfg = adjust_settings(toggle=args.ORC_toggle,
                                         env_cfg=env_cfg,
                                         train_cfg=train_cfg)
    task_registry.set_log_dir_name(train_cfg)

    task_registry.make_gym_and_sim()
    wandb_helper.setup_wandb(env_cfg=env_cfg, train_cfg=train_cfg, args=args)
    env = task_registry.make_env(name=args.task, env_cfg=env_cfg)
    # * then make env
    policy_runner = task_registry.make_alg_runner(env, train_cfg)

    local_code_save_helper.log_and_save(
        env, env_cfg, train_cfg, policy_runner)
    wandb_helper.attach_runner(policy_runner=policy_runner)

    return train_cfg, policy_runner


def train(train_cfg, policy_runner):
    wandb_helper = wandb_singleton.WandbSingleton()

    policy_runner.learn(
        num_learning_iterations=train_cfg.runner.max_iterations,
        init_at_random_ep_len=True)

    wandb_helper.close_wandb()


if __name__ == '__main__':
    train_cfg, policy_runner = setup()
    train(train_cfg=train_cfg, policy_runner=policy_runner)

gym/scripts/train_ORC.py
mit-biomimetics/ORCAgym
from gym.envs import __init__
from gym.utils import get_args, task_registry
from gym.utils.logging_and_saving \
    import local_code_save_helper, wandb_singleton
from ORC import adjust_settings

from torch.multiprocessing import Process
from torch.multiprocessing import set_start_method

def setup(toggle):
    args = get_args()
    wandb_helper = wandb_singleton.WandbSingleton()

    # * prepare environment
    env_cfg, train_cfg = task_registry.create_cfgs(args)
    env_cfg, train_cfg = adjust_settings(toggle=toggle,
                                         env_cfg=env_cfg,
                                         train_cfg=train_cfg)
    task_registry.set_log_dir_name(train_cfg)

    task_registry.make_gym_and_sim()
    wandb_helper.setup_wandb(env_cfg=env_cfg, train_cfg=train_cfg, args=args)
    env = task_registry.make_env(name=args.task, env_cfg=env_cfg)
    # * then make env
    policy_runner = task_registry.make_alg_runner(env, train_cfg)

    local_code_save_helper.log_and_save(
        env, env_cfg, train_cfg, policy_runner)
    wandb_helper.attach_runner(policy_runner=policy_runner)

    return train_cfg, policy_runner


def train(train_cfg, policy_runner):
    wandb_helper = wandb_singleton.WandbSingleton()

    policy_runner.learn(
        num_learning_iterations=train_cfg.runner.max_iterations,
        init_at_random_ep_len=True)

    wandb_helper.close_wandb()

def worker(toggle):
    train_cfg, policy_runner = setup(toggle)
    train(train_cfg=train_cfg, policy_runner=policy_runner)

if __name__ == '__main__':
    all_toggles = ['000', '010', '011', '100', '101', '110', '111']
    set_start_method('spawn')

    processes = []
    for toggle in all_toggles:
        p = Process(target=worker, args=(toggle,))
        p.start()
        p.join()  # Wait for the process to finish

        # Free up any resources if needed
        p.terminate()
        p.close()

    print("All learning runs are completed.")

gym/scripts/train_ORC_all.py
mit-biomimetics/ORCAgym
# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
# 
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice, this
# list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
#
# 3. Neither the name of the copyright holder nor the names of its
# contributors may be used to endorse or promote products derived from
# this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#
# Copyright (c) 2021 ETH Zurich, Nikita Rudin

from .helpers import (class_to_dict, get_load_path, get_args, set_seed,
                      update_class_from_dict, select_run)

from .task_registry import task_registry
from .math import *
from .terrain import Terrain
from .interfaces.interfaces import KeyboardInterface, GamepadInterface
from .logger import Logger

gym/utils/__init__.py
mit-biomimetics/ORCAgym
# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice,
# this list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
#
# 3. Neither the name of the copyright holder nor the names of its
# contributors may be used to endorse or promote products derived from
# this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF
# THE POSSIBILITY OF SUCH DAMAGE.
#
# Copyright (c) 2021 ETH Zurich, Nikita Rudin

import os
import numpy as np
import random
from isaacgym import gymapi
from isaacgym import gymutil
from gym import LEGGED_GYM_ROOT_DIR
import torch
from isaacgym.torch_utils import to_torch


def class_to_dict(obj, torch_device=None) -> dict:
    """
    Convert a multi-level class (config) into a dict of dicts.
    Keyword arguments:
    torch_device (opt): if specified, converts each element to a torch tensor.
    """
    if not hasattr(obj, "__dict__"):
        return obj
    result = {}
    for key in dir(obj):
        if key.startswith("_"):
            continue
        element = []
        val = getattr(obj, key)
        if isinstance(val, list):
            for item in val:
                element.append(class_to_dict(item))
        else:
            element = class_to_dict(val)
        if torch_device:
            result[key] = to_torch(element, device=torch_device)
        else:
            result[key] = element
    return result


def update_class_from_dict(obj, dict):
    for key, val in dict.items():
        attr = getattr(obj, key, None)
        if isinstance(attr, type):
            update_class_from_dict(attr, val)
        else:
            setattr(obj, key, val)
    return


def set_seed(seed):
    if seed == -1:
        seed = np.random.randint(0, 10000)
    print("Setting seed: {}".format(seed))

    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)


def parse_sim_params(args, cfg):
    # * code from Isaac Gym Preview 2
    # * initialize sim params
    sim_params = gymapi.SimParams()

    # * set some values from args
    if args.physics_engine == gymapi.SIM_FLEX:
        if args.device != "cpu":
            print("WARNING: Using Flex with GPU instead of PHYSX!")
    elif args.physics_engine == gymapi.SIM_PHYSX:
        sim_params.physx.use_gpu = args.use_gpu
        sim_params.physx.num_subscenes = args.subscenes
    sim_params.use_gpu_pipeline = args.use_gpu_pipeline

    # * if sim options are provided in cfg, parse them and
    # * update/override above:
    if "sim" in cfg:
        gymutil.parse_sim_config(cfg["sim"], sim_params)

    # * Override num_threads if passed on the command line
    if args.physics_engine == gymapi.SIM_PHYSX and args.num_threads > 0:
        sim_params.physx.num_threads = args.num_threads

    return sim_params


def get_load_path(name, load_run=-1, checkpoint=-1):
    root = os.path.join(LEGGED_GYM_ROOT_DIR, 'logs', name)
    run_path = select_run(root, load_run)
    model_name = select_model(run_path, checkpoint)
    load_path = os.path.join(run_path, model_name)
    return load_path


def select_run(root, load_run):
    try:
        runs = sorted(
            os.listdir(root),
            key=lambda x: os.path.getctime(os.path.join(root, x)))
        if 'exported' in runs:
            runs.remove('exported')
        last_run = os.path.join(root, runs[-1])
    except:  # ! no bare excepts!!!
        raise ValueError("No runs in this directory: " + root)

    if load_run == -1:
        load_run = last_run
    else:
        load_run = os.path.join(root, load_run)
    return load_run


def select_model(load_run, checkpoint):
    if checkpoint == -1:
        models = [file for file in os.listdir(load_run) if 'model' in file]
        models.sort(key=lambda m: '{0:0>15}'.format(m))
        model = models[-1]
    else:
        model = "model_{}.pt".format(checkpoint)
    return model


def update_cfg_from_args(env_cfg, train_cfg, args):
    # * seed
    if env_cfg is not None:
        # * num envs
        if args.num_envs is not None:
            env_cfg.env.num_envs = args.num_envs
        if args.omega is not None:
            env_cfg.osc.omega = args.omega
            env_cfg.osc.omega_range = [args.omega, args.omega]
        if args.coupling is not None:
            env_cfg.osc.coupling = args.coupling
            env_cfg.osc.coupling_range = [args.coupling, args.coupling]
        if args.offset is not None:
            env_cfg.osc.offset = args.offset
            env_cfg.osc.offset_range = [args.offset, args.offset]
        if args.velocity is not None:
            env_cfg.commands.ranges.lin_vel_x = [args.velocity, args.velocity]
            env_cfg.commands.ranges.lin_vel_y = 0.
            env_cfg.commands.ranges.yaw_vel = 0.
        if args.shank is not None:
            env_cfg.asset.shank_length_diff = args.shank
            env_cfg.asset.file = "{LEGGED_GYM_ROOT_DIR}/resources/robots/" \
            + "mini_cheetah/urdf/mini_cheetah_" \
            + str(int(args.shank)) + ".urdf"
    if train_cfg is not None:
        if args.seed is not None:
            train_cfg.seed = args.seed
        # * alg runner parameters
        if args.max_iterations is not None:
            train_cfg.runner.max_iterations = args.max_iterations
        if args.resume:
            train_cfg.runner.resume = args.resume
        if args.experiment_name is not None:
            train_cfg.runner.experiment_name = args.experiment_name
        if args.run_name is not None:
            train_cfg.runner.run_name = args.run_name
        if args.load_run is not None:
            train_cfg.runner.load_run = args.load_run
        if args.checkpoint is not None:
            train_cfg.runner.checkpoint = args.checkpoint
        if args.rl_device is not None:
            train_cfg.runner.device = args.rl_device


def get_args():
    custom_parameters = [
        {"name": "--task", "type": str, "default": "mini_cheetah_osc", "help":
            "Resume training or start testing from a checkpoint. "
            "Overrides config file if provided."},
        {"name": "--resume", "action": "store_true", "default": False, "help":
            "Resume training from a checkpoint"},
        {"name": "--experiment_name", "type": str,  "help":
            "Name of the experiment to run or load. "
            "Overrides config file if provided."},
        {"name": "--run_name", "type": str,  "help":
            "Name of the run. Overrides config file if provided."},
        {"name": "--load_run", "type": str,  "help":
            "Name of the run to load when resume=True. "
            "If -1: will load the last run. "
            "Overrides config file if provided."},
        {"name": "--checkpoint", "type": int,  "help":
            "Saved model checkpoint number. "
            "If -1: will load the last checkpoint. "
            "Overrides config file if provided."},

        {"name": "--headless", "action": "store_true", "default": False,
         "help": "Force display off at all times"},
        {"name": "--horovod", "action": "store_true", "default": False, "help":
            "Use horovod for multi-gpu training"},
        {"name": "--rl_device", "type": str, "default": "cuda:0", "help":
            'Device used by the RL algorithm, '
            '(cpu, gpu, cuda:0, cuda:1 etc..)'},
        {"name": "--num_envs", "type": int, "help":
            "Number of environments to create. "
            "Overrides config file if provided."},
        {"name": "--seed", "type": int, "help":
            "Random seed. Overrides config file if provided."},
        {"name": "--max_iterations", "type": int, "help":
            "Maximum number of training iterations. "
            "Overrides config file if provided."},
        {"name": "--wandb_project", "type": str, "help":
            "Enter the name of your project for better WandB tracking."},
        {"name": "--wandb_entity", "type": str, "help":
            "Enter your wandb entity username to track your "
            "experiment on your account."},
        {"name": "--wandb_sweep_id", "type": str, "help":
            "Enter a WandB sweep ID to continue an existing sweep."},
        {"name": "--wandb_sweep_config", "type": str, "help":
            "Enter the name of a JSON config file defining the WandB sweep."},
        {"name": "--disable_wandb", "action": "store_true", "default": False,
         "help": "Disable WandB logging for debugging."},

         # * oscillator specific cli args
        {"name": "--omega", "type": float, "default": None, "help":
            "Omega value for oscillators"
            "Overrides config file if provided."},
        {"name": "--coupling", "type": float, "default": None, "help":
            "coupling value for oscillators"
            "Overrides config file if provided."},
        {"name": "--offset", "type": float, "default": None, "help":
            "offset value for oscillators"
            "Overrides config file if provided."},
        {"name": "--velocity", "type": float, "default": None, "help":
            "Velocity command (no range)"
            "Overrides config file if provided."},
        {"name": "--shank", "type": float, "default": None, "help":
            "Shank length difference from default"
            "Overrides config file if provided."},
        {"name": "--ORC_toggle", "type": str, "default": "111", "help":
            "Observations, Rewards, Coupling toggle"
            "Overrides config file if provided."}
    ]
    # * parse arguments
    args = gymutil.parse_arguments(
        description="RL Policy",
        custom_parameters=custom_parameters)

    # * name allignment
    args.sim_device_id = args.compute_device_id
    args.sim_device = args.sim_device_type
    if args.sim_device == 'cuda':
        args.sim_device += f":{args.sim_device_id}"
    return args

gym/utils/helpers.py
mit-biomimetics/ORCAgym
import os
os.environ['PYGAME_HIDE_SUPPORT_PROMPT'] = "hide"
import pygame

class GamepadInterface():
    def __init__(self, env):
        pygame.init()
        pygame.joystick.init()
        print("__________________________________________________________")
        if pygame.joystick.get_count() == 1:
            print("Using gamepad interface, "
                  + "overriding default comand settings")
            print("left joystick: forward, strafe left, "
                  "backward, strafe right")
            print("right joystick (left/right): yaw left/right")
            print("back button: reset environments")
            print("start button: quit")
            self.joystick = pygame.joystick.Joystick(0)

            env.commands[:] = 0.
            env.cfg.commands.resampling_time = env.max_episode_length_s + 1
            self.max_vel_backward = -1.
            self.max_vel_forward = 4.
            self.max_vel_sideways = 1.
            self.max_yaw_vel = 2.

        elif pygame.joystick.get_count() > 1:
            print("WARNING: you have more than one gamepad plugged in."
                  "Please unplug one.")
        else:
            print("WARNING: failed to initialize gamepad.")
        print("__________________________________________________________")

    def update(self, env):
        for event in pygame.event.get():
            if event.type == pygame.JOYBUTTONDOWN:
                if event.button == 7:  # press start
                    exit()
                if event.button == 6:  # press back
                    env.episode_length_buf[:] = env.max_episode_length+1
            if event.type == pygame.JOYAXISMOTION:
                # left joystick
                if event.axis == 1:  # up-down
                    if event.value >= 0:
                        env.commands[:, 0] = self.max_vel_backward*event.value
                    else:
                        env.commands[:, 0] = -self.max_vel_forward*event.value
                if event.axis == 0:  # left_right
                    env.commands[:, 1] = -self.max_vel_sideways*event.value
                # right joystick
                if event.axis == 3:  # left-right
                    env.commands[:, 2] = -self.max_yaw_vel*event.value

gym/utils/interfaces/GamepadInterface.py
mit-biomimetics/ORCAgym
from isaacgym import gymapi
import torch


class KeyboardInterface():
    def __init__(self, env):
        env.gym.subscribe_viewer_keyboard_event(env.viewer,
                                                gymapi.KEY_W, 'forward')
        env.gym.subscribe_viewer_keyboard_event(env.viewer,
                                                gymapi.KEY_A, 'left')
        env.gym.subscribe_viewer_keyboard_event(env.viewer,
                                                gymapi.KEY_D, 'right')
        env.gym.subscribe_viewer_keyboard_event(env.viewer,
                                                gymapi.KEY_S, 'back')
        env.gym.subscribe_viewer_keyboard_event(env.viewer,
                                                gymapi.KEY_Q, 'yaw_left')
        env.gym.subscribe_viewer_keyboard_event(env.viewer,
                                                gymapi.KEY_E, 'yaw_right')
        env.gym.subscribe_viewer_keyboard_event(env.viewer,
                                                gymapi.KEY_R, "RESET")
        env.gym.subscribe_viewer_keyboard_event(env.viewer,
                                                gymapi.KEY_ESCAPE, "QUIT")
        print("______________________________________________________________")
        print("Using keyboard interface, overriding default comand settings")
        print("commands are in 1/5 increments of max.")
        print("WASD: forward, strafe left, "
              "backward, strafe right")
        print("QE: yaw left/right")
        print("R: reset environments")
        print("ESC: quit")
        print("______________________________________________________________")

        env.commands[:] = 0.
        env.cfg.commands.resampling_time = env.max_episode_length_s + 1
        self.max_vel_backward = -1.
        self.max_vel_forward = 4.
        self.increment_x = (self.max_vel_forward-self.max_vel_backward)*0.5

        self.max_vel_sideways = 1.0
        self.increment_y = self.max_vel_sideways*0.2

        self.max_vel_yaw = 2.0
        self.increment_yaw = self.max_vel_yaw*0.2

    def update(self, env):
        for evt in env.gym.query_viewer_action_events(env.viewer):
            if evt.value == 0:
                continue
            if evt.action == 'forward':
                if env.commands[0, 0] < -1.:
                    env.commands[:, 0] = -1.
                elif env.commands [0, 0] < 0.:
                    env.commands[:, 0] = 0.
                elif env.commands [0, 0] < 1.:
                    env.commands[:, 0] = 1.
                else:
                    env.commands[:, 0] = 4.
                # env.commands[:, 0] = torch.clamp(env.commands[:, 0]
                #                                  + self.increment_x,
                #                                  max=self.max_vel_forward)
            elif evt.action == 'back':
                if env.commands[0, 0] > 1.:
                    env.commands[:, 0] = 1.
                elif env.commands [0, 0] > 0.:
                    env.commands[:, 0] = 0.
                elif env.commands [0, 0] > -1.:
                    env.commands[:, 0] = -1.
                else:
                    env.commands[:, 0] = -4.
                # env.commands[:, 0] = torch.clamp(env.commands[:, 0]
                #                                  - self.increment_x,
                #                                  min=self.max_vel_backward)
            elif evt.action == 'left':
                # env.commands[:, 1] = torch.clamp(env.commands[:, 1]
                #                                  + self.increment_y,
                #                                  min=-self.max_vel_sideways)
                # similar to above, increment at -1, 0, 1
                if env.commands[0, 1] <= 0.:
                    env.commands[:, 1] = 1.
                else:
                    env.commands[:, 1] = 0.
            elif evt.action == 'right':
                # env.commands[:, 1] = torch.clamp(env.commands[:, 1]
                #                                  - self.increment_y,
                #                                  max=self.max_vel_sideways)
                if env.commands[0, 1] >= 0.:
                    env.commands[:, 1] = 0.
                else:
                    env.commands[:, 1] = -1.
            elif evt.action == 'yaw_left':
                # env.commands[:, 2] = torch.clamp(env.commands[:, 2]
                #                                  - self.increment_yaw,
                #                                  min=-self.max_vel_yaw)
                if env.commands[0, 2] < -3.:
                    env.commands[:, 2] = -3.
                elif env.commands[0, 2] < 0.:
                    env.commands[:, 2] = 0.
                elif env.commands[0, 2] < 3.:
                    env.commands[:, 2] = 3.
                else:
                    env.commands[:, 2] = 4.

            elif evt.action == 'yaw_right':
                # env.commands[:, 2] = torch.clamp(env.commands[:, 2]
                #                                  + self.increment_yaw,
                #                                  max=self.max_vel_yaw)
                if env.commands[0, 2] > 3.:
                    env.commands[:, 2] = 3.
                elif env.commands[0, 2] > 0.:
                    env.commands[:, 2] = 0.
                elif env.commands[0, 2] > -3.:
                    env.commands[:, 2] = -3.
                else:
                    env.commands[:, 2] = -4.

            elif evt.action == "QUIT":
                exit()
            elif evt.action == "RESET":
                env.timed_out[:] = True
                env.reset()
                env.commands[:] = 0.

gym/utils/interfaces/KeyboardInterface.py
mit-biomimetics/ORCAgym
from .KeyboardInterface import KeyboardInterface
from .GamepadInterface import GamepadInterface

gym/utils/interfaces/interfaces.py
mit-biomimetics/ORCAgym
# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
# 
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice, this
# list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
#
# 3. Neither the name of the copyright holder nor the names of its
# contributors may be used to endorse or promote products derived from
# this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#
# Copyright (c) 2021 ETH Zurich, Nikita Rudin

import matplotlib.pyplot as plt
import numpy as np
from collections import defaultdict
from multiprocessing import Process, Value

class Logger:
    def __init__(self, dt):
        self.state_log = defaultdict(list)
        self.rew_log = defaultdict(list)
        self.dt = dt
        self.num_episodes = 0
        self.plot_process = None

    def log_state(self, key, value):
        self.state_log[key].append(value)

    def log_states(self, dict):
        for key, value in dict.items():
            self.log_state(key, value)

    def log_rewards(self, dict, num_episodes):
        for key, value in dict.items():
            try:
                self.rew_log[key].append(value.item() * num_episodes)
            except AttributeError:
                self.rew_log[key].append(value * num_episodes)
        self.num_episodes += num_episodes

    def reset(self):
        self.state_log.clear()
        self.rew_log.clear()

    def plot_states(self, log_file_path, visualize=True):
        self.plot_process = Process(target=self._plot(log_file_path, visualize))
        self.plot_process.start()

    def _plot(self, log_file_path, visualize):
        nb_rows = 3
        nb_cols = 4
        fig, axs = plt.subplots(nb_rows, nb_cols)
        for key, value in self.state_log.items():
            time = np.linspace(0, len(value)*self.dt, len(value))
            break
        log = self.state_log

        if log["oscillators"] and log["contact_forces_z"]:
            oscillators = np.array(log["oscillators"])
            oscillators_vel = np.array(log["oscillators_vel"])
            scaled_grf = np.array(log["grf"])

            forces = np.array(log["contact_forces_z"])
            time = np.transpose([time])
            vel_x = np.transpose(np.array([log["base_vel_x"]]))
            vel_target = np.transpose(np.array([log["command_x"]]))
            torques = np.array(log["torques"])
            velocities = np.array(log["velocities"])
            np.savetxt(log_file_path, np.hstack((time, vel_x, vel_target, oscillators, oscillators_vel, forces, scaled_grf, torques, velocities)), delimiter=",")
            for i in range(oscillators.shape[1]):
                #TODO: Change so that each [0,2pi) cycle is in a diff color
                a = axs[0, i]
                a.plot(oscillators[:, i], scaled_grf[:, i], label=f'force {i}')
                a.set(xlabel='phase [/]', ylabel='Forces z [/]', title='Vertical Contact forces normalized by body weight')
            for i in range(oscillators.shape[1]):
                #TODO: Currently shows jump from 2pi -> 0 clamp
                a = axs[1, i]
                a.scatter(oscillators[:, i], oscillators_vel[:, i], label=f'phase rate of change {i}')
                a.set(xlabel='phase [/]', ylabel='dt * phase rate of change [/]', title='Phase')
            for i in range(oscillators.shape[1]):
                a = axs[2, i]
                a.plot(time, oscillators[:, i], label=f'Phase over time {i}')
                a.set(xlabel='time [s]', ylabel='phase [/]', title='Phase')
        if visualize:
            plt.show()

        nb_rows = 3
        nb_cols = 3
        fig, axs = plt.subplots(nb_rows, nb_cols)
        # plot joint targets and measured positions
        a = axs[1, 0]
        if log["dof_pos"]: a.plot(time, log["dof_pos"], label='measured')
        if log["dof_pos_target"]: a.plot(time, log["dof_pos_target"], label='target')
        a.set(xlabel='time [s]', ylabel='Position [rad]', title='DOF Position')
        a.legend()
        # plot joint velocity
        a = axs[1, 1]
        if log["dof_vel"]: a.plot(time, log["dof_vel"], label='measured')
        if log["dof_vel_target"]: a.plot(time, log["dof_vel_target"], label='target')
        a.set(xlabel='time [s]', ylabel='Velocity [rad/s]', title='Joint Velocity')
        a.legend()
        # plot base vel x
        a = axs[0, 0]
        if log["base_vel_x"]: a.plot(time, log["base_vel_x"], label='measured')
        if log["command_x"]: a.plot(time, log["command_x"], label='commanded')
        a.set(xlabel='time [s]', ylabel='base lin vel [m/s]', title='Base velocity x')
        a.legend()
        # plot base vel y
        a = axs[0, 1]
        if log["base_vel_y"]: a.plot(time, log["base_vel_y"], label='measured')
        if log["command_y"]: a.plot(time, log["command_y"], label='commanded')
        a.set(xlabel='time [s]', ylabel='base lin vel [m/s]', title='Base velocity y')
        a.legend()
        # plot base vel yaw
        a = axs[0, 2]
        if log["base_vel_yaw"]: a.plot(time, log["base_vel_yaw"], label='measured')
        if log["command_yaw"]: a.plot(time, log["command_yaw"], label='commanded')
        a.set(xlabel='time [s]', ylabel='base ang vel [rad/s]', title='Base velocity yaw')
        a.legend()
        # plot base vel z
        a = axs[1, 2]
        if log["base_vel_z"]: a.plot(time, log["base_vel_z"], label='measured')
        a.set(xlabel='time [s]', ylabel='base lin vel [m/s]', title='Base velocity z')
        a.legend()
        # plot contact forces
        a = axs[2, 0]
        if log["contact_forces_z"]:
            forces = np.array(log["contact_forces_z"])
            for i in range(forces.shape[1]):
                a.plot(time, forces[:, i], label=f'force {i}')
        a.set(xlabel='time [s]', ylabel='Forces z [N]', title='Vertical Contact forces')
        a.legend()
        # plot torque/vel curves
        a = axs[2, 1]
        if log["dof_vel"]!=[] and log["dof_torque"]!=[]: a.plot(log["dof_vel"], log["dof_torque"], 'x', label='measured')
        a.set(xlabel='Joint vel [rad/s]', ylabel='Joint Torque [Nm]', title='Torque/velocity curves')
        a.legend()
        # plot torques
        a = axs[2, 2]
        if log["dof_torque"]!=[]: a.plot(time, log["dof_torque"], label='measured')
        a.set(xlabel='time [s]', ylabel='Joint Torque [Nm]', title='Torque')
        a.legend()
        if visualize:
            plt.show()

    def print_rewards(self, reward_file_path):
        print("Average rewards per second:")
        reward_keys = []
        reward_means = []
        for key, mean in self.rew_log.items():
            reward_keys = np.append(reward_keys, key)
            reward_means = np.append(reward_means, mean[0])
            print(f" - {key}: {mean}")
        reward_keys = np.append(reward_keys, "Total number of episodes")
        reward_means = np.append(reward_means, self.num_episodes)
        print(f"Total number of episodes: {self.num_episodes}")
        np.savetxt(reward_file_path, np.column_stack((reward_keys, reward_means)), delimiter=",", fmt="%s")

    def __del__(self):
        if self.plot_process is not None:
            self.plot_process.kill()
gym/utils/logger.py
mit-biomimetics/ORCAgym
import os
from gym import LEGGED_GYM_ROOT_DIR


def log_and_save(env, env_cfg, train_cfg, runner):
    """Configure local and cloud code logging"""

    # setup local code saving if enabled
    if check_local_saving_flag(train_cfg):
        save_paths = get_local_save_paths(env, env_cfg)
        runner.logger.configure_local_files(save_paths)


def check_local_saving_flag(train_cfg):
    """Check if enable_local_saving is set to true in the training_config"""

    if hasattr(train_cfg, 'logging') and \
       hasattr(train_cfg.logging, 'enable_local_saving'):
        enable_local_saving = train_cfg.logging.enable_local_saving
    else:
        enable_local_saving = False
    return enable_local_saving


def get_local_save_paths(env, env_cfg):
    """Create a save_paths object for saving code locally"""

    learning_dir = os.path.join(LEGGED_GYM_ROOT_DIR, 'learning')
    learning_target = os.path.join('learning')

    gym_dir = os.path.join(LEGGED_GYM_ROOT_DIR, 'gym')
    gym_target = os.path.join('gym')

    # list of things to copy
    # source paths need the full path and target are relative to log_dir
    save_paths = [
        {'type': 'dir', 'source_dir': learning_dir,
                        'target_dir': learning_target,
            'include_patterns': ('*.py', '*.json')},
        {'type': 'dir', 'source_dir': gym_dir,
                        'target_dir': gym_target,
            'include_patterns': ('*.py', '*.json')}
    ]

    return save_paths

gym/utils/logging_and_saving/local_code_save_helper.py
mit-biomimetics/ORCAgym
import os
import json
import wandb
from gym import LEGGED_GYM_ROOT_DIR


class WandbSingleton(object):
    def __new__(self):
        if not hasattr(self, 'instance'):
            self.instance = super(WandbSingleton, self).__new__(self)
            self.entity_name = None
            self.project_name = None
            self.experiment_name = ''
            self.enabled = False
            self.parameters_dict = None

        return self.instance

    def set_wandb_values(self, args, train_cfg=None):
        # build the path for the wandb_config.json file
        wandb_config_path = os.path.join(
            LEGGED_GYM_ROOT_DIR, 'user', 'wandb_config.json')

        # load entity and project name from JSON if it exists
        if os.path.exists(wandb_config_path):
            json_data = json.load(open(wandb_config_path))
            print('Loaded WandB entity and project from JSON.')
            self.entity_name = json_data['entity']
            self.project_name = json_data['project']
        # override entity and project by commandline args if provided
        if args.wandb_entity is not None:
            print('Received WandB entity from arguments.')
            self.entity_name = args.wandb_entity
        if args.wandb_project is not None:
            print('Received WandB project from arguments.')
            self.project_name = args.wandb_project
        # assume WandB is off if entity or project is None and short-circuit
        if train_cfg is not None:
            self.experiment_name = train_cfg.log_dir.split('/')[-1]

        if (self.entity_name is None or self.project_name is None
           or args.disable_wandb):
            self.enabled = False
        else:
            print(f'Setting WandB project name: {self.project_name}\n' +
                  f'Setting WandB entitiy name: {self.entity_name}\n')
            self.enabled = True

    def set_wandb_sweep_cfg_values(self, env_cfg, train_cfg):
        if not self.is_wandb_enabled():
            return

        # * update the config settings based off the sweep_dict
        for key, value in self.parameters_dict.items():
            print('Setting: ' + key + ' = ' + str(value))
            locs = key.split('.')

            if locs[0] == 'train_cfg':
                attr = train_cfg
            elif locs[0] == 'env_cfg':
                attr = env_cfg
            else:
                print('Unrecognized cfg: ' + locs[0])
                break

            for loc in locs[1:-1]:
                attr = getattr(attr, loc)

            setattr(attr, locs[-1], value)

    def is_wandb_enabled(self):
        return self.enabled

    def get_entity_name(self):
        return self.entity_name

    def get_project_name(self):
        return self.project_name

    def setup_wandb(self, env_cfg, train_cfg, args, is_sweep=False):
        self.set_wandb_values(args, train_cfg)

        # short-circuit if the values say WandB is turned off
        if not self.is_wandb_enabled():
            print('WARNING: WandB is disabled and will not save or log.')
            return

        wandb.config = {}

        if is_sweep:
            wandb.init(dir=os.path.join(LEGGED_GYM_ROOT_DIR, 'logs'),
                       config=wandb.config,
                       name=self.experiment_name)
        else:
            wandb.init(project=self.project_name,
                       entity=self.entity_name,
                       dir=os.path.join(LEGGED_GYM_ROOT_DIR, 'logs'),
                       config=wandb.config,
                       name=self.experiment_name)

        wandb.run.log_code(
            os.path.join(LEGGED_GYM_ROOT_DIR, 'gym'))

        self.parameters_dict = wandb.config
        self.set_wandb_sweep_cfg_values(env_cfg=env_cfg, train_cfg=train_cfg)

    def attach_runner(self, policy_runner):
        if not self.is_wandb_enabled():
            return
        policy_runner.attach_to_wandb(wandb)

    def close_wandb(self):
        if self.enabled:
            wandb.finish()

gym/utils/logging_and_saving/wandb_singleton.py
mit-biomimetics/ORCAgym
# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice,
# this list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
#
# 3. Neither the name of the copyright holder nor the names of its
# contributors may be used to endorse or promote products derived from
# this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF
# THE POSSIBILITY OF SUCH DAMAGE.
#
# Copyright (c) 2021 ETH Zurich, Nikita Rudin

import torch
from torch import Tensor
import numpy as np
from isaacgym.torch_utils import quat_apply, normalize, torch_rand_float
from typing import Tuple


# @ torch.jit.script
def quat_apply_yaw(quat, vec):
    quat_yaw = quat.clone().view(-1, 4)
    quat_yaw[:, :2] = 0.
    quat_yaw = normalize(quat_yaw)
    return quat_apply(quat_yaw, vec)


# @ torch.jit.script
def wrap_to_pi(angles):
    angles %= 2*np.pi
    angles -= 2*np.pi * (angles > np.pi)
    return angles


# @ torch.jit.script
def torch_rand_sqrt_float(lower, upper, shape, device):
    # type: (float, float, Tuple[int, int], str) -> Tensor
    r = 2*torch.rand(*shape, device=device) - 1
    r = torch.where(r < 0., -torch.sqrt(-r), torch.sqrt(r))
    r = (r + 1.) / 2.
    return (upper - lower) * r + lower


# @ torch.jit.script
def random_sample(env_ids, low, high, device):
    """
    Generate random samples for each entry of env_ids
    """
    rand_pos = torch_rand_float(0, 1, (len(env_ids), len(low)),
                                device=device)
    diff_pos = (high - low).repeat(len(env_ids), 1)
    random_dof_pos = rand_pos*diff_pos + low.repeat(len(env_ids), 1)
    return random_dof_pos


# @ torch.jit.script
def exp_avg_filter(x, avg, alpha=0.8):
    """
    Simple exponential average filter
    """
    avg = alpha*x + (1-alpha)*avg
    return avg

gym/utils/math.py
mit-biomimetics/ORCAgym
# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice,
# this list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
#
# 3. Neither the name of the copyright holder nor the names of its
# contributors may be used to endorse or promote products derived from
# this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF
# THE POSSIBILITY OF SUCH DAMAGE.
#
# Copyright (c) 2021 ETH Zurich, Nikita Rudin

import os
from isaacgym import gymapi
from isaacgym import gymutil
from datetime import datetime
from typing import Tuple

from learning.env import VecEnv
from learning.runners import OnPolicyRunner
from learning.utils import set_discount_from_horizon

from gym import LEGGED_GYM_ROOT_DIR
from .helpers import \
    update_cfg_from_args, class_to_dict, get_load_path, set_seed
from gym.envs.base.legged_robot_config import (
    LeggedRobotCfg, LeggedRobotRunnerCfg)
from gym.envs.base.base_config import BaseConfig
from gym.envs.base.sim_config import SimCfg


class TaskRegistry():
    def __init__(self):
        self.task_classes = {}
        self.env_cfgs = {}
        self.train_cfgs = {}
        self.sim_cfg = class_to_dict(SimCfg)
        self.sim = {}

    def register(self, name: str, task_class: VecEnv,
                 env_cfg: BaseConfig, train_cfg: LeggedRobotRunnerCfg):
        self.task_classes[name] = task_class
        self.env_cfgs[name] = env_cfg
        self.train_cfgs[name] = train_cfg

    def get_task_class(self, name: str) -> VecEnv:
        return self.task_classes[name]

    def get_cfgs(self, name) -> Tuple[LeggedRobotCfg, LeggedRobotRunnerCfg]:
        env_cfg = self.env_cfgs[name]
        train_cfg = self.train_cfgs[name]
        # * copy seed
        env_cfg.seed = train_cfg.seed
        return env_cfg, train_cfg

    def create_cfgs(self, args):
        env_cfg, train_cfg = self.get_cfgs(name=args.task)
        self.update_and_parse_cfgs(env_cfg, train_cfg, args)
        self.set_log_dir_name(train_cfg)
        return env_cfg, train_cfg

    def set_log_dir_name(self, train_cfg, log_root="default"):
        if log_root == "default":
            log_root = os.path.join(
                LEGGED_GYM_ROOT_DIR, 'logs', train_cfg.runner.experiment_name)
            log_dir = os.path.join(
                log_root, datetime.now().strftime('%b%d_%H-%M-%S')
                + '_' + train_cfg.runner.run_name)
        elif log_root is None:
            log_dir = None
        train_cfg.log_dir = log_dir

    def update_and_parse_cfgs(self, env_cfg, train_cfg, args):
        update_cfg_from_args(env_cfg, train_cfg, args)
        self.convert_frequencies_to_params(env_cfg, train_cfg)
        self.update_sim_cfg(args)

    def convert_frequencies_to_params(self, env_cfg, train_cfg):
        self.set_control_and_sim_dt(env_cfg, train_cfg)
        self.set_discount_rates(train_cfg, env_cfg.control.ctrl_dt)

    def set_control_and_sim_dt(self, env_cfg, train_cfg):
        env_cfg.control.decimation = int(env_cfg.control.desired_sim_frequency
                                         / env_cfg.control.ctrl_frequency)
        env_cfg.control.ctrl_dt = 1.0 / env_cfg.control.ctrl_frequency
        env_cfg.sim_dt = env_cfg.control.ctrl_dt / env_cfg.control.decimation
        self.sim_cfg["dt"] = env_cfg.sim_dt
        if env_cfg.sim_dt != 1.0/env_cfg.control.desired_sim_frequency:
            print(f'****** Simulation dt adjusted from '
                  f'{1.0/env_cfg.control.desired_sim_frequency}'
                  f' to {env_cfg.sim_dt}.')

    def set_discount_rates(self, train_cfg, dt):
        if hasattr(train_cfg.algorithm, 'discount_horizon'):
            hrzn = train_cfg.algorithm.discount_horizon
            train_cfg.algorithm.gamma = set_discount_from_horizon(dt, hrzn)

        if hasattr(train_cfg.algorithm, 'GAE_bootstrap_horizon'):
            hrzn = train_cfg.algorithm.GAE_bootstrap_horizon
            train_cfg.algorithm.lam = set_discount_from_horizon(dt, hrzn)

    def update_sim_cfg(self, args):
        self.sim["sim_device"] = args.sim_device
        self.sim["sim_device_id"] = args.sim_device_id
        self.sim["graphics_device_id"] = args.graphics_device_id
        self.sim["physics_engine"] = args.physics_engine
        self.sim["headless"] = args.headless
        if self.sim["headless"]:
            self.sim["graphics_device_id"] = -1
        self.sim["params"] = gymapi.SimParams()
        self.sim["params"].physx.use_gpu = args.use_gpu
        self.sim["params"].physx.num_subscenes = args.subscenes
        self.sim["params"].use_gpu_pipeline = args.use_gpu_pipeline
        gymutil.parse_sim_config(self.sim_cfg, self.sim["params"])

    def make_gym_and_sim(self):
        self.make_gym()
        self.make_sim()

    def make_gym(self):
        self._gym = gymapi.acquire_gym()

    def make_sim(self):
        self._sim = self._gym.create_sim(
            self.sim["sim_device_id"],
            self.sim["graphics_device_id"],
            self.sim["physics_engine"],
            self.sim["params"])

    def make_env(self, name, env_cfg) -> VecEnv:
        if name in self.task_classes:
            task_class = self.get_task_class(name)
        else:
            raise ValueError(f"Task with name: {name} was not registered")
        set_seed(env_cfg.seed)
        env = task_class(gym=self._gym, sim=self._sim, cfg=env_cfg,
                         sim_params=self.sim["params"],
                         sim_device=self.sim["sim_device"],
                         headless=self.sim["headless"])
        return env

    def make_alg_runner(self, env, train_cfg):
        train_cfg_dict = class_to_dict(train_cfg)
        runner = OnPolicyRunner(env, train_cfg_dict, train_cfg.runner.device)

        # * save resume path before creating a new log_dir
        if train_cfg.runner.resume:
            resume_path = get_load_path(name=train_cfg.runner.experiment_name,
                                        load_run=train_cfg.runner.load_run,
                                        checkpoint=train_cfg.runner.checkpoint)
            print(f"Loading model from: {resume_path}")
            runner.load(resume_path)
        return runner


# make global task registry
task_registry = TaskRegistry()

gym/utils/task_registry.py
mit-biomimetics/ORCAgym
# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice,
# this list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
#
# 3. Neither the name of the copyright holder nor the names of its
# contributors may be used to endorse or promote products derived from
# this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF
# THE POSSIBILITY OF SUCH DAMAGE.
#
# Copyright (c) 2021 ETH Zurich, Nikita Rudin

import numpy as np

from isaacgym import terrain_utils
from gym.envs.base.legged_robot_config import LeggedRobotCfg


class Terrain:
    def __init__(self, cfg: LeggedRobotCfg.terrain, num_robots) -> None:

        self.cfg = cfg
        self.num_robots = num_robots
        self.type = cfg.mesh_type
        if self.type in ["none", "plane"]:
            return
        self.env_length = cfg.terrain_length
        self.env_width = cfg.terrain_width
        if np.sum(cfg.terrain_proportions) != 1.0:
            raise ValueError("Terrain proportions provided do not sum to 1.0")
        self.proportions = [np.sum(cfg.terrain_proportions[:i+1])
                            for i in range(len(cfg.terrain_proportions))]

        self.cfg.num_sub_terrains = cfg.num_rows * cfg.num_cols
        self.env_origins = np.zeros((cfg.num_rows, cfg.num_cols, 3))

        self.width_per_env_pixels = int(
            self.env_width / cfg.horizontal_scale)
        self.length_per_env_pixels = int(
            self.env_length / cfg.horizontal_scale)

        self.border = int(
            cfg.border_size/self.cfg.horizontal_scale)
        self.tot_cols = int(
            cfg.num_cols * self.width_per_env_pixels) + 2 * self.border
        self.tot_rows = int(
            cfg.num_rows * self.length_per_env_pixels) + 2 * self.border

        self.height_field_raw = np.zeros(
            (self.tot_rows, self.tot_cols), dtype=np.int16)
        if cfg.curriculum:
            self.curiculum()
        elif cfg.selected:
            self.selected_terrain()
        else:
            self.randomized_terrain()

        self.heightsamples = self.height_field_raw
        if self.type == "trimesh":
            self.vertices, self.triangles = \
                terrain_utils.convert_heightfield_to_trimesh(
                    self.height_field_raw,
                    self.cfg.horizontal_scale,
                    self.cfg.vertical_scale,
                    self.cfg.slope_treshold)

    def randomized_terrain(self):
        for k in range(self.cfg.num_sub_terrains):
            # Env coordinates in the world
            (i, j) = np.unravel_index(k, (self.cfg.num_rows, self.cfg.num_cols))

            choice = np.random.uniform(0, 1)
            difficulty = np.random.choice([0.5, 0.75, 0.9])
            terrain = self.make_terrain(choice, difficulty)
            self.add_terrain_to_map(terrain, i, j)

    def curiculum(self):
        for j in range(self.cfg.num_cols):
            for i in range(self.cfg.num_rows):
                difficulty = i / self.cfg.num_rows
                choice = j / self.cfg.num_cols + 0.001

                terrain = self.make_terrain(choice, difficulty)
                self.add_terrain_to_map(terrain, i, j)

    def selected_terrain(self):
        terrain_type = self.cfg.terrain_kwargs.pop('type')
        for k in range(self.cfg.num_sub_terrains):
            # Env coordinates in the world
            (i, j) = np.unravel_index(k, (self.cfg.num_rows, self.cfg.num_cols))

            terrain = terrain_utils.SubTerrain(
                "terrain",
                width=self.width_per_env_pixels,
                length=self.width_per_env_pixels,
                vertical_scale=self.vertical_scale,
                horizontal_scale=self.horizontal_scale)

            eval(terrain_type)(
                terrain, **self.cfg.terrain_kwargs.terrain_kwargs)
            self.add_terrain_to_map(terrain, i, j)

    def make_terrain(self, choice, difficulty):
        terrain = terrain_utils.SubTerrain(
            "terrain",
            width=self.width_per_env_pixels,
            length=self.length_per_env_pixels,
            vertical_scale=self.cfg.vertical_scale,
            horizontal_scale=self.cfg.horizontal_scale)
        slope = difficulty * 0.2
        step_height = 0.05 + 0.05 * difficulty
        discrete_obstacles_height = 0.05 + difficulty * 0.1
        stepping_stones_size = 1.5 * (1.05 - difficulty)
        stone_distance = 0.05 if difficulty == 0 else 0.1
        gap_size = .3 * difficulty
        pit_depth = .3 * difficulty
        if choice < self.proportions[0]:
            if choice < self.proportions[0] / 2:
                slope *= -1
            terrain_utils.pyramid_sloped_terrain(
                terrain, slope=slope, platform_size=3.)
        elif choice < self.proportions[1]:
            # terrain_utils.pyramid_sloped_terrain(
            #     terrain, slope=slope, platform_size=3.)
            terrain_utils.random_uniform_terrain(
                terrain, min_height=-0.05, max_height=0.05,
                step=0.005, downsampled_scale=0.2)
        elif choice < self.proportions[3]:
            if choice < self.proportions[2]:
                step_height *= -1
            terrain_utils.pyramid_stairs_terrain(
                terrain, step_width=0.31,
                step_height=step_height, platform_size=3.)
        elif choice < self.proportions[4]:
            num_rectangles = 20
            rectangle_min_size = 1.
            rectangle_max_size = 2.
            terrain_utils.discrete_obstacles_terrain(
                terrain, discrete_obstacles_height, rectangle_min_size,
                rectangle_max_size, num_rectangles, platform_size=3.)
        elif choice < self.proportions[5]:
            terrain_utils.stepping_stones_terrain(
                terrain, stone_size=stepping_stones_size,
                stone_distance=stone_distance, max_height=0., platform_size=4.)
        elif choice < self.proportions[6]:
            gap_terrain(terrain, gap_size=gap_size, platform_size=3.)
        else:
            pit_terrain(terrain, depth=pit_depth, platform_size=4.)

        return terrain

    def add_terrain_to_map(self, terrain, row, col):
        i = row
        j = col
        # map coordinate system
        start_x = self.border + i * self.length_per_env_pixels
        end_x = self.border + (i + 1) * self.length_per_env_pixels
        start_y = self.border + j * self.width_per_env_pixels
        end_y = self.border + (j + 1) * self.width_per_env_pixels
        self.height_field_raw[start_x: end_x, start_y:end_y] = \
            terrain.height_field_raw

        env_origin_x = (i + 0.5) * self.env_length
        env_origin_y = (j + 0.5) * self.env_width
        x1 = int((self.env_length/2. - 1) / terrain.horizontal_scale)
        x2 = int((self.env_length/2. + 1) / terrain.horizontal_scale)
        y1 = int((self.env_width/2. - 1) / terrain.horizontal_scale)
        y2 = int((self.env_width/2. + 1) / terrain.horizontal_scale)
        env_origin_z = np.max(
            terrain.height_field_raw[x1:x2, y1:y2])*terrain.vertical_scale
        self.env_origins[i, j] = [env_origin_x, env_origin_y, env_origin_z]


def gap_terrain(terrain, gap_size, platform_size=1.):
    gap_size = int(gap_size / terrain.horizontal_scale)
    platform_size = int(platform_size / terrain.horizontal_scale)

    center_x = terrain.length // 2
    center_y = terrain.width // 2
    x1 = (terrain.length - platform_size) // 2
    x2 = x1 + gap_size
    y1 = (terrain.width - platform_size) // 2
    y2 = y1 + gap_size

    terrain.height_field_raw[
        center_x-x2: center_x + x2, center_y-y2: center_y + y2] = -1000
    terrain.height_field_raw[
        center_x-x1: center_x + x1, center_y-y1: center_y + y1] = 0


def pit_terrain(terrain, depth, platform_size=1.):
    depth = int(depth / terrain.vertical_scale)
    platform_size = int(platform_size / terrain.horizontal_scale / 2)
    x1 = terrain.length // 2 - platform_size
    x2 = terrain.length // 2 + platform_size
    y1 = terrain.width // 2 - platform_size
    y2 = terrain.width // 2 + platform_size
    terrain.height_field_raw[x1:x2, y1:y2] = -depth

gym/utils/terrain.py
mit-biomimetics/ORCAgym
# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice,
# this list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
#
# 3. Neither the name of the copyright holder nor the names of its
# contributors may be used to endorse or promote products derived from
# this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF
# THE POSSIBILITY OF SUCH DAMAGE.
#
# Copyright (c) 2021 ETH Zurich, Nikita Rudin

learning/__init__.py
mit-biomimetics/ORCAgym
import torch.nn as nn
import torch.optim as optim

from learning.modules import StateEstimatorNN
from learning.storage import SERolloutStorage


class StateEstimator:
    """ This class provides a learned state estimator.
    This is trained with supervised learning, using only on-policy data
    collected in a rollout storage.
    predict() function provides state estimation for RL given the observation
    update() function optimizes for the nn params
    process_env_step() function stores values in a rollout storage
    """
    state_estimator: StateEstimatorNN

    def __init__(self,
                 state_estimator,    # nn module
                 learning_rate=1e-3,
                 num_mini_batches=1,
                 num_learning_epochs=1,
                 device='cpu',
                 **kwargs
                 ):

        # general parameters
        self.device = device
        self.learning_rate = learning_rate
        self.num_mini_batches = num_mini_batches
        self.num_learning_epochs = num_learning_epochs

        # SE storage
        self.transition = SERolloutStorage.Transition()
        self.storage = None

        # SE network and optimizer
        self.state_estimator = state_estimator
        self.state_estimator.to(self.device)
        self.optimizer = optim.Adam(self.state_estimator.parameters(),
                                    lr=learning_rate)
        self.SE_loss_fn = nn.MSELoss()

    def init_storage(self, num_envs, num_transitions_per_env,
                     obs_shape, se_shape):
        self.storage = SERolloutStorage(num_envs, num_transitions_per_env,
                                        obs_shape, se_shape,
                                        device=self.device)

    def predict(self, obs):
        return self.state_estimator.evaluate(obs)

    def process_env_step(self, obs, SE_targets):
        # Record the transition
        self.transition.SE_targets = SE_targets
        self.transition.observations = obs
        self.storage.add_transitions(self.transition)
        self.transition.clear()

    def update(self):
        """ Update the SE neural network weights via supervised learning """
        generator = self.storage.mini_batch_generator(self.num_mini_batches,
                                                      self.num_learning_epochs)
        mean_loss = 0
        for obs_batch, SE_target_batch in generator:

            SE_estimate_batch = self.state_estimator.evaluate(obs_batch)

            SE_loss = self.SE_loss_fn(SE_estimate_batch, SE_target_batch)
            self.optimizer.zero_grad()
            SE_loss.backward()
            self.optimizer.step()
            mean_loss += SE_loss.item()

        num_updates = self.num_learning_epochs * self.num_mini_batches
        mean_loss /= num_updates
        self.storage.clear()

        return mean_loss

    def export(self, path):
        self.state_estimator.export(path)

learning/algorithms/SE.py
mit-biomimetics/ORCAgym
# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice,
# this list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
#
# 3. Neither the name of the copyright holder nor the names of its
# contributors may be used to endorse or promote products derived from
# this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF
# THE POSSIBILITY OF SUCH DAMAGE.
#
# Copyright (c) 2021 ETH Zurich, Nikita Rudin

from .ppo import PPO
from .SE import StateEstimator

learning/algorithms/__init__.py
mit-biomimetics/ORCAgym
# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice,
# this list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
#
# 3. Neither the name of the copyright holder nor the names of its
# contributors may be used to endorse or promote products derived from
# this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF
# THE POSSIBILITY OF SUCH DAMAGE.
#
# Copyright (c) 2021 ETH Zurich, Nikita Rudin

import torch
import torch.nn as nn
import torch.optim as optim

from learning.modules import ActorCritic
from learning.storage import RolloutStorage


class PPO:
    actor_critic: ActorCritic

    def __init__(self,
                 actor_critic,
                 num_learning_epochs=1,
                 num_mini_batches=1,
                 clip_param=0.2,
                 gamma=0.998,
                 lam=0.95,
                 value_loss_coef=1.0,
                 entropy_coef=0.0,
                 learning_rate=1e-3,
                 max_grad_norm=1.0,
                 use_clipped_value_loss=True,
                 schedule="fixed",
                 desired_kl=0.01,
                 device='cpu',
                 **kwargs
                 ):

        self.device = device

        self.desired_kl = desired_kl
        self.schedule = schedule
        self.learning_rate = learning_rate

        # * PPO components
        self.actor_critic = actor_critic
        self.actor_critic.to(self.device)
        self.storage = None  # initialized later
        self.optimizer = optim.Adam(self.actor_critic.parameters(),
                                    lr=learning_rate)
        self.transition = RolloutStorage.Transition()

        # * PPO parameters
        self.clip_param = clip_param
        self.num_learning_epochs = num_learning_epochs
        self.num_mini_batches = num_mini_batches
        self.value_loss_coef = value_loss_coef
        self.entropy_coef = entropy_coef
        self.gamma = gamma
        self.lam = lam
        self.max_grad_norm = max_grad_norm
        self.use_clipped_value_loss = use_clipped_value_loss

    def init_storage(self, num_envs, num_transitions_per_env, actor_obs_shape,
                     critic_obs_shape, action_shape):
        self.storage = RolloutStorage(num_envs, num_transitions_per_env,
                                      actor_obs_shape, critic_obs_shape,
                                      action_shape, self.device)

    def test_mode(self):
        self.actor_critic.test()

    def train_mode(self):
        self.actor_critic.train()

    def act(self, obs, critic_obs):
        # * Compute the actions and values
        self.transition.actions = self.actor_critic.act(obs).detach()
        self.transition.values = \
            self.actor_critic.evaluate(critic_obs).detach()
        self.transition.actions_log_prob = \
            self.actor_critic.get_actions_log_prob(
                self.transition.actions).detach()
        self.transition.action_mean = self.actor_critic.action_mean.detach()
        self.transition.action_sigma = self.actor_critic.action_std.detach()
        # * need to record obs and critic_obs before env.step()
        self.transition.observations = obs
        self.transition.critic_observations = critic_obs
        return self.transition.actions

    def process_env_step(self, rewards, dones, timed_out=None):
        self.transition.rewards = rewards.clone()
        self.transition.dones = dones
        # * Bootstrapping on time outs
        if timed_out is not None:
            self.transition.rewards += self.gamma * self.transition.values \
                                      * timed_out

        # * Record the transition
        self.storage.add_transitions(self.transition)
        self.transition.clear()

    def compute_returns(self, last_critic_obs):
        last_values = self.actor_critic.evaluate(last_critic_obs).detach()
        self.storage.compute_returns(last_values, self.gamma, self.lam)

    def update(self):
        mean_value_loss = 0
        mean_surrogate_loss = 0
        generator = self.storage.mini_batch_generator(
            self.num_mini_batches, self.num_learning_epochs)
        for (obs_batch, critic_obs_batch, actions_batch, target_values_batch,
             advantages_batch, returns_batch, old_actions_log_prob_batch,
             old_mu_batch, old_sigma_batch) in generator:

            self.actor_critic.act(obs_batch)
            actions_log_prob_batch = \
                self.actor_critic.get_actions_log_prob(actions_batch)
            value_batch = self.actor_critic.evaluate(critic_obs_batch)
            mu_batch = self.actor_critic.action_mean
            sigma_batch = self.actor_critic.action_std
            entropy_batch = self.actor_critic.entropy

            # * KL
            if self.desired_kl is not None and self.schedule == 'adaptive':
                with torch.inference_mode():
                    kl = torch.sum(
                        torch.log(sigma_batch / old_sigma_batch + 1.e-5)
                        + (torch.square(old_sigma_batch)
                           + torch.square(old_mu_batch - mu_batch))
                        / (2.0 * torch.square(sigma_batch)) - 0.5, axis=-1)
                    kl_mean = torch.mean(kl)

                    if kl_mean > self.desired_kl * 2.0:
                        self.learning_rate = max(
                            1e-5, self.learning_rate / 1.5)
                    elif kl_mean < self.desired_kl / 2.0 and kl_mean > 0.0:
                        self.learning_rate = min(
                            1e-2, self.learning_rate * 1.5)

                    for param_group in self.optimizer.param_groups:
                        param_group['lr'] = self.learning_rate

            # * Surrogate loss
            ratio = torch.exp(
                actions_log_prob_batch
                - torch.squeeze(old_actions_log_prob_batch))
            surrogate = -torch.squeeze(advantages_batch) * ratio
            surrogate_clipped = \
                -torch.squeeze(advantages_batch) * torch.clamp(
                    ratio, 1.0 - self.clip_param, 1.0 + self.clip_param)
            surrogate_loss = torch.max(surrogate, surrogate_clipped).mean()

            # * Value function loss
            if self.use_clipped_value_loss:
                value_clipped = (
                    target_values_batch
                    + (value_batch - target_values_batch).clamp(
                        -self.clip_param, self.clip_param))
                value_losses = (value_batch - returns_batch).pow(2)
                value_losses_clipped = (value_clipped - returns_batch).pow(2)
                value_loss = torch.max(
                    value_losses, value_losses_clipped).mean()
            else:
                value_loss = (returns_batch - value_batch).pow(2).mean()

            loss = (surrogate_loss
                    + self.value_loss_coef * value_loss
                    - self.entropy_coef * entropy_batch.mean())

            # * Gradient step
            self.optimizer.zero_grad()
            loss.backward()
            nn.utils.clip_grad_norm_(
                self.actor_critic.parameters(), self.max_grad_norm)
            self.optimizer.step()

            mean_value_loss += value_loss.item()
            mean_surrogate_loss += surrogate_loss.item()

        num_updates = self.num_learning_epochs * self.num_mini_batches
        mean_value_loss /= num_updates
        mean_surrogate_loss /= num_updates
        self.storage.clear()

        return mean_value_loss, mean_surrogate_loss

learning/algorithms/ppo.py
mit-biomimetics/ORCAgym
# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice,
# this list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
#
# 3. Neither the name of the copyright holder nor the names of its
# contributors may be used to endorse or promote products derived from
# this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF
# THE POSSIBILITY OF SUCH DAMAGE.
#
# Copyright (c) 2021 ETH Zurich, Nikita Rudin

from .vec_env import VecEnv

learning/env/__init__.py
mit-biomimetics/ORCAgym
# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice,
# this list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
#
# 3. Neither the name of the copyright holder nor the names of its
# contributors may be used to endorse or promote products derived from
# this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF
# THE POSSIBILITY OF SUCH DAMAGE.
#
# Copyright (c) 2021 ETH Zurich, Nikita Rudin

from abc import ABC, abstractmethod
import torch
from typing import Tuple, Union


# minimal interface of the environment
class VecEnv(ABC):
    num_envs: int
    num_obs: int
    num_privileged_obs: int
    num_actions: int
    max_episode_length: int
    privileged_obs_buf: torch.Tensor
    obs_buf: torch.Tensor
    rew_buf: torch.Tensor
    reset_buf: torch.Tensor
    episode_length_buf: torch.Tensor  # current episode duration
    extras: dict
    device: torch.device

    @abstractmethod
    def step(self, actions: torch.Tensor) \
        -> Tuple[
            torch.Tensor, Union[torch.Tensor, None],
            torch.Tensor, torch.Tensor, dict]:
        pass

    @abstractmethod
    def reset(self, env_ids: Union[list, torch.Tensor]):
        pass

    @abstractmethod
    def get_observations(self) -> torch.Tensor:
        pass

    @abstractmethod
    def get_privileged_observations(self) -> Union[torch.Tensor, None]:
        pass

learning/env/vec_env.py
mit-biomimetics/ORCAgym
# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice,
# this list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
#
# 3. Neither the name of the copyright holder nor the names of its
# contributors may be used to endorse or promote products derived from
# this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF
# THE POSSIBILITY OF SUCH DAMAGE.
#
# Copyright (c) 2021 ETH Zurich, Nikita Rudin

from .actor_critic import ActorCritic
from .state_estimator import StateEstimatorNN
from .actor import Actor

learning/modules/__init__.py
mit-biomimetics/ORCAgym
import torch
import torch.nn as nn
from torch.distributions import Normal
from .utils import RunningMeanStd
from .utils import create_MLP
from .utils import export_network


class Actor(nn.Module):
    def __init__(self,
                 num_obs,
                 num_actions,
                 hidden_dims,
                 activation="relu",
                 init_noise_std=1.0,
                 normalize_obs=True,
                 **kwargs):

        if kwargs:
            print("Actor.__init__ got unexpected arguments, "
                  "which will be ignored: "
                  + str([key for key in kwargs.keys()]))
        super().__init__()

        self._normalize_obs = normalize_obs
        if self._normalize_obs:
            self.obs_rms = RunningMeanStd(num_obs)

        self.num_obs = num_obs
        self.num_actions = num_actions
        self.NN = create_MLP(num_obs, num_actions, hidden_dims, activation)

        # Action noise
        self.std = nn.Parameter(init_noise_std * torch.ones(num_actions))
        self.distribution = None
        # disable args validation for speedup
        Normal.set_default_validate_args = False

    @property
    def action_std(self):
        return self.std

    @property
    def action_mean(self):
        return self.distribution.mean

    @property
    def action_std(self):
        return self.distribution.stddev

    @property
    def entropy(self):
        return self.distribution.entropy().sum(dim=-1)

    def update_distribution(self, observations):
        if self._normalize_obs:
            observations = self.norm_obs(observations)
        mean = self.NN(observations)
        self.distribution = Normal(mean, mean*0. + self.std)

    def act(self, observations):
        self.update_distribution(observations)
        return self.distribution.sample()

    def get_actions_log_prob(self, actions):
        return self.distribution.log_prob(actions).sum(dim=-1)

    def act_inference(self, observations):
        if self._normalize_obs:
            observations = self.norm_obs(observations)
        actions_mean = self.NN(observations)
        return actions_mean

    def export(self, path):
        export_network(self.NN, "policy", path, self.num_obs)

    def norm_obs(self, observation):
        with torch.no_grad():
            return self.obs_rms(observation) if self._normalize_obs else observation

learning/modules/actor.py
mit-biomimetics/ORCAgym
import torch.nn as nn

from .actor import Actor
from .critic import Critic


class ActorCritic(nn.Module):
    def __init__(self, num_actor_obs, num_critic_obs, num_actions,
                 actor_hidden_dims=[256, 256, 256],
                 critic_hidden_dims=[256, 256, 256],
                 activation="relu",
                 init_noise_std=1.0,
                 **kwargs):

        if kwargs:
            print(
                "ActorCritic.__init__ got unexpected arguments, "
                "which will be ignored: "
                + str([key for key in kwargs.keys()]))
        super(ActorCritic, self).__init__()

        self.actor = Actor(num_actor_obs,
                          num_actions,
                          actor_hidden_dims,
                          activation,
                          init_noise_std)

        self.critic = Critic(num_critic_obs,
                             critic_hidden_dims,
                             activation)

        print(f"Actor MLP: {self.actor.NN}")
        print(f"Critic MLP: {self.critic.NN}")

    @property
    def action_mean(self):
        return self.actor.action_mean

    @property
    def action_std(self):
        return self.actor.action_std

    @property
    def entropy(self):
        return self.actor.entropy

    @property
    def std(self):
        return self.actor.std

    def update_distribution(self, observations):
        self.actor.update_distribution(observations)

    def act(self, observations, **kwargs):
        return self.actor.act(observations)

    def get_actions_log_prob(self, actions):
        return self.actor.get_actions_log_prob(actions)

    def act_inference(self, observations):
        return self.actor.act_inference(observations)

    def evaluate(self, critic_observations, **kwargs):
        return self.critic.evaluate(critic_observations)

    def export_policy(self, path):
        self.actor.export(path)

learning/modules/actor_critic.py
mit-biomimetics/ORCAgym
import torch
import torch.nn as nn
from .utils import create_MLP
from .utils import RunningMeanStd


class Critic(nn.Module):
    def __init__(self,
                 num_obs,
                 hidden_dims,
                 activation="relu",
                 normalize_obs=True,
                 **kwargs):

        if kwargs:
            print("Critic.__init__ got unexpected arguments, "
                  "which will be ignored: "
                  + str([key for key in kwargs.keys()]))
        super().__init__()

        self.NN = create_MLP(num_obs, 1, hidden_dims, activation)

        self._normalize_obs = normalize_obs
        if self._normalize_obs:
            self.obs_rms = RunningMeanStd(num_obs)

    def evaluate(self, critic_observations):
        if self._normalize_obs:
            observations = self.norm_obs(critic_observations)
        return self.NN(observations).squeeze()

    def norm_obs(self, observation):
        with torch.no_grad():
            return self.obs_rms(observation) if self._normalize_obs else observation

learning/modules/critic.py
mit-biomimetics/ORCAgym
import torch
import torch.nn as nn
from .utils import create_MLP
from .utils import export_network


class StateEstimatorNN(nn.Module):
    """ Set up a neural network for state-estimation

    Keyword arguments:
    num_inputs -- the number of observations used
    num_outputs -- the number of estimated states
    hidden_dims -- list of hidden layers and their sizes (default [256, 128])
    activation -- activation type (default 'elu')
    dropouts -- list of dropout rate for each layer (default None)
    """
    def __init__(self,  num_inputs, num_outputs, hidden_dims=[256, 128],
                 activation='elu', dropouts=None, **kwargs):
        if kwargs:
            print("StateEstimator.__init__ got unexpected arguments, "
                  "which will be ignored: "
                  + str([key for key in kwargs.keys()]))
        super().__init__()

        self.num_inputs = num_inputs
        self.num_outputs = num_outputs
        self.NN = create_MLP(num_inputs, num_outputs, hidden_dims,
                             activation, dropouts)
        print(f"State Estimator MLP: {self.NN}")

    def evaluate(self, observations):
        return self.NN(observations)

    def export(self, path):
        export_network(self.NN, "state_estimator", path, self.num_inputs)

learning/modules/state_estimator.py
mit-biomimetics/ORCAgym
from .neural_net import create_MLP, export_network
from .running_mean_std import RunningMeanStd, RunningMeanStdObs

learning/modules/utils/__init__.py
mit-biomimetics/ORCAgym
import torch
import os
import copy


def create_MLP(num_inputs, num_outputs, hidden_dims, activation,
               dropouts=None):

    activation = get_activation(activation)

    if dropouts is None:
        dropouts = [0]*len(hidden_dims)

    layers = []
    if not hidden_dims:  # handle no hidden layers
        add_layer(layers, num_inputs, num_outputs)
    else:
        add_layer(layers, num_inputs, hidden_dims[0], activation, dropouts[0])
        for i in range(len(hidden_dims)):
            if i == len(hidden_dims) - 1:
                add_layer(layers, hidden_dims[i], num_outputs)
            else:
                add_layer(layers, hidden_dims[i], hidden_dims[i+1],
                          activation, dropouts[i+1])
    return torch.nn.Sequential(*layers)


def get_activation(act_name):
    if act_name == "elu":
        return torch.nn.ELU()
    elif act_name == "selu":
        return torch.nn.SELU()
    elif act_name == "relu":
        return torch.nn.ReLU()
    elif act_name == "crelu":
        return torch.nn.ReLU()
    elif act_name == "lrelu":
        return torch.nn.LeakyReLU()
    elif act_name == "tanh":
        return torch.nn.Tanh()
    elif act_name == "sigmoid":
        return torch.nn.Sigmoid()
    else:
        print("invalid activation function!")
        return None


def add_layer(layer_list, num_inputs, num_outputs, activation=None, dropout=0):
    layer_list.append(torch.nn.Linear(num_inputs, num_outputs))
    if dropout > 0:
        layer_list.append(torch.nn.Dropout(p=dropout))
    if activation is not None:
        layer_list.append(activation)


def export_network(network, network_name, path, num_inputs):
    """
    Thsi function traces and exports the given network module in .pt and
    .onnx file formats. These can be used for evaluation on other systems
    without needing a Pytorch environment.

    :param network:         PyTorch neural network module
    :param network_name:    (string) Network will be saved with this name
    :path:                  (string) Network will be saved to this location
    :param num_inputs:      (int) Number of inputs to the network module
    """

    os.makedirs(path, exist_ok=True)
    path_TS = os.path.join(path, network_name + '.pt')   # TorchScript path
    path_onnx = os.path.join(path, network_name + '.onnx')   # ONNX path
    model = copy.deepcopy(network).to('cpu')
    # To trace model, must be evaluated once with arbitrary input
    model.eval()
    dummy_input = torch.rand(num_inputs,)
    model_traced = torch.jit.trace(model, dummy_input)
    torch.jit.save(model_traced, path_TS)
    torch.onnx.export(model_traced, dummy_input, path_onnx)
learning/modules/utils/neural_net.py
mit-biomimetics/ORCAgym
import torch
import torch.nn as nn
import numpy as np

def get_mean_var_with_masks(values, masks):
    sum_mask = masks.sum()
    values_mask = values * masks
    values_mean = values_mask.sum() / sum_mask
    min_sqr = ((((values_mask)**2)/sum_mask).sum() - ((values_mask/sum_mask).sum())**2)
    values_var = min_sqr * sum_mask / (sum_mask-1)
    return values_mean, values_var

class RunningMeanStd(nn.Module):
    def __init__(self, insize, epsilon=1e-05, per_channel=False, norm_only=False):
        super(RunningMeanStd, self).__init__()
        print('RunningMeanStd: ', insize)
        self.insize = insize
        self.epsilon = epsilon

        self.norm_only = norm_only
        self.per_channel = per_channel
        if per_channel:
            if len(self.insize) == 3:
                self.axis = [0,2,3]
            if len(self.insize) == 2:
                self.axis = [0,2]
            if len(self.insize) == 1:
                self.axis = [0]
            in_size = self.insize[0] 
        else:
            self.axis = [0]
            in_size = insize

        self.register_buffer("running_mean", torch.zeros(in_size, dtype = torch.float64))
        self.register_buffer("running_var", torch.ones(in_size, dtype = torch.float64))
        self.register_buffer("count", torch.ones((), dtype = torch.float64))

    def _update_mean_var_count_from_moments(self, mean, var, count, batch_mean, batch_var, batch_count):
        delta = batch_mean - mean
        tot_count = count + batch_count

        new_mean = mean + delta * batch_count / tot_count
        m_a = var * count
        m_b = batch_var * batch_count
        M2 = m_a + m_b + delta**2 * count * batch_count / tot_count
        new_var = M2 / tot_count
        new_count = tot_count
        return new_mean, new_var, new_count

    def forward(self, input, denorm=False, mask=None):
        if self.training:
            if mask is not None:
                mean, var = get_mean_var_with_masks(input, mask)
            else:
                mean = input.mean(self.axis) # along channel axis
                var = input.var(self.axis)
            self.running_mean, self.running_var, self.count = self._update_mean_var_count_from_moments(self.running_mean, self.running_var, self.count, 
                                                    mean, var, input.size()[0] )

        # change shape
        if self.per_channel:
            if len(self.insize) == 3:
                current_mean = self.running_mean.view([1, self.insize[0], 1, 1]).expand_as(input)
                current_var = self.running_var.view([1, self.insize[0], 1, 1]).expand_as(input)
            if len(self.insize) == 2:
                current_mean = self.running_mean.view([1, self.insize[0], 1]).expand_as(input)
                current_var = self.running_var.view([1, self.insize[0], 1]).expand_as(input)
            if len(self.insize) == 1:
                current_mean = self.running_mean.view([1, self.insize[0]]).expand_as(input)
                current_var = self.running_var.view([1, self.insize[0]]).expand_as(input)        
        else:
            current_mean = self.running_mean
            current_var = self.running_var
        # get output


        if denorm:
            y = torch.clamp(input, min=-5.0, max=5.0)
            y = torch.sqrt(current_var.float() + self.epsilon)*y + current_mean.float()
        else:
            if self.norm_only:
                y = input/ torch.sqrt(current_var.float() + self.epsilon)
            else:
                y = (input - current_mean.float()) / torch.sqrt(current_var.float() + self.epsilon)
                y = torch.clamp(y, min=-5.0, max=5.0)
        return y

class RunningMeanStdObs(nn.Module):
    def __init__(self, insize, epsilon=1e-05, per_channel=False, norm_only=False):
        assert(isinstance(insize, dict))
        super(RunningMeanStdObs, self).__init__()
        self.running_mean_std = nn.ModuleDict({
            k : RunningMeanStd(v, epsilon, per_channel, norm_only) for k,v in insize.items()
        })
    
    def forward(self, input, denorm=False):
        res = {k : self.running_mean_std[k](v, denorm) for k,v in input.items()}
        return res
learning/modules/utils/running_mean_std.py
mit-biomimetics/ORCAgym
# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice,
# this list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
#
# 3. Neither the name of the copyright holder nor the names of its
# contributors may be used to endorse or promote products derived from
# this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF
# THE POSSIBILITY OF SUCH DAMAGE.
#
# Copyright (c) 2021 ETH Zurich, Nikita Rudin

from .on_policy_runner import OnPolicyRunner

learning/runners/__init__.py
mit-biomimetics/ORCAgym
# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice,
# this list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
#
# 3. Neither the name of the copyright holder nor the names of its
# contributors may be used to endorse or promote products derived from
# this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF
# THE POSSIBILITY OF SUCH DAMAGE.
#
# Copyright (c) 2021 ETH Zurich, Nikita Rudin

import time
import os
import wandb
import torch
from isaacgym.torch_utils import torch_rand_float

from learning.algorithms import PPO
from learning.modules import ActorCritic
from learning.env import VecEnv
from learning.utils import remove_zero_weighted_rewards
from learning.utils import Logger


class OnPolicyRunner:

    def __init__(self,
                 env: VecEnv,
                 train_cfg,
                 device='cpu'):

        self.device = device
        self.env = env
        self.parse_train_cfg(train_cfg)

        num_actor_obs = self.get_obs_size(self.policy_cfg["actor_obs"])
        num_critic_obs = self.get_obs_size(self.policy_cfg["critic_obs"])
        num_actions = self.get_action_size(self.policy_cfg["actions"])
        actor_critic = ActorCritic(num_actor_obs,
                                   num_critic_obs,
                                   num_actions,
                                   **self.policy_cfg).to(self.device)

        alg_class = eval(self.cfg["algorithm_class_name"])
        self.alg: PPO = alg_class(actor_critic,
                                  device=self.device, **self.alg_cfg)

        self.num_steps_per_env = self.cfg["num_steps_per_env"]
        self.save_interval = self.cfg["save_interval"]
        self.tot_timesteps = 0
        self.tot_time = 0
        self.it = 0

        # * init storage and model
        self.init_storage()

        # * Log
        self.log_dir = train_cfg["log_dir"]
        self.SE_path = os.path.join(self.log_dir, 'SE')  # log_dir for SE
        self.logger = Logger(self.log_dir, self.env.max_episode_length_s,
                             self.device)

        reward_keys_to_log = \
            list(self.policy_cfg["reward"]["weights"].keys()) \
            + list(self.policy_cfg["reward"]["termination_weight"].keys())

        reward_keys_to_log += ["Total_reward"]

        # ! gait tracking
        self.gait_weights = {'trot': 1.,
                             'pace': 1.,
                             'pronk': 1.,
                             'bound': 1,
                             'asymettric': 1.}
        reward_keys_to_log += list(self.gait_weights.keys())
        reward_keys_to_log += ['Gait_score']
        self.logger.initialize_buffers(self.env.num_envs, reward_keys_to_log)

    def parse_train_cfg(self, train_cfg):
        self.cfg = train_cfg['runner']
        self.alg_cfg = train_cfg['algorithm']

        remove_zero_weighted_rewards(train_cfg['policy']['reward']['weights'])
        self.policy_cfg = train_cfg['policy']

    def init_storage(self):
        num_actor_obs = self.get_obs_size(self.policy_cfg["actor_obs"])
        num_critic_obs = self.get_obs_size(self.policy_cfg["critic_obs"])
        num_actions = self.get_action_size(self.policy_cfg["actions"])
        self.alg.init_storage(self.env.num_envs,
                              self.num_steps_per_env,
                              actor_obs_shape=[num_actor_obs],
                              critic_obs_shape=[num_critic_obs],
                              action_shape=[num_actions])

    def attach_to_wandb(self, wandb, log_freq=100, log_graph=True):
        wandb.watch((self.alg.actor_critic.actor,
                    self.alg.actor_critic.critic),
                    log_freq=log_freq,
                    log_graph=log_graph)

    def learn(self, num_learning_iterations, init_at_random_ep_len=False):

        if init_at_random_ep_len:
            self.env.episode_length_buf = torch.randint_like(
                self.env.episode_length_buf,
                high=int(self.env.max_episode_length))

        actor_obs = self.get_obs(self.policy_cfg["actor_obs"])
        critic_obs = self.get_obs(self.policy_cfg["critic_obs"])
        self.alg.actor_critic.train()
        self.num_learning_iterations = num_learning_iterations
        self.tot_iter = self.it + num_learning_iterations

        self.save()

        reward_weights = self.policy_cfg['reward']['weights']
        termination_weight = self.policy_cfg['reward']['termination_weight']
        rewards = 0.*self.get_rewards(reward_weights)

        # * simulate 2 seconds to let the robots fall
        for i in range(int(2/self.env.dt)):
            self.env.step()

        for self.it in range(self.it+1, self.tot_iter+1):
            start = time.time()
            # * Rollout
            with torch.inference_mode():

                for i in range(self.num_steps_per_env):
                    actions = self.alg.act(actor_obs, critic_obs)
                    self.set_actions(actions)

                    self.env.step()

                    actor_obs = self.get_noisy_obs(
                        self.policy_cfg['actor_obs'],
                        self.policy_cfg['noise'])
                    critic_obs = self.get_obs(self.policy_cfg['critic_obs'])
                    # * get time_outs
                    timed_out = self.get_timed_out()
                    terminated = self.get_terminated()
                    dones = timed_out | terminated

                    rewards += self.get_and_log_rewards(reward_weights,
                                                        modifier=self.env.dt,
                                                        mask=~terminated)
                    rewards += self.get_and_log_rewards(termination_weight,
                                                        mask=terminated)
                    self.logger.log_current_reward('Total_reward', rewards)

                    # * log for gaits
                    gait_score = self.get_and_log_rewards(self.gait_weights,
                                                          modifier=self.env.dt,
                                                          mask=~terminated)
                    self.logger.log_current_reward('Gait_score', gait_score)
                    self.alg.process_env_step(rewards, dones, timed_out)
                    self.logger.update_episode_buffer(dones)
                    rewards *= 0.

                stop = time.time()
                self.collection_time = stop - start
                # * Learning step
                start = stop
                self.alg.compute_returns(critic_obs)

            self.mean_value_loss, self.mean_surrogate_loss = self.alg.update()
            stop = time.time()
            self.learn_time = stop - start
            self.tot_timesteps += self.num_steps_per_env * self.env.num_envs
            self.tot_time += self.collection_time + self.learn_time

            self.log()

            if self.it % self.save_interval == 0:
                self.save()

        # * save
        self.save()

    def get_noise(self, obs_list, noise_dict):
        noise_vec = torch.zeros(self.get_obs_size(obs_list),
                                device=self.device)
        obs_index = 0
        for obs in obs_list:
            obs_size = self.get_obs_size([obs])
            if obs in noise_dict.keys():
                noise_tensor = torch.ones(obs_size).to(self.device) \
                               * torch.tensor(noise_dict[obs]).to(self.device)
                if obs in self.env.scales.keys():
                    noise_tensor /= self.env.scales[obs]
                noise_vec[obs_index:obs_index+obs_size] = noise_tensor
            obs_index += obs_size
        return torch_rand_float(-1., 1., (self.env.num_envs, len(noise_vec)),
                                self.device) * noise_vec * noise_dict["scale"]

    def get_noisy_obs(self, obs_list, noise_dict):
        observation = self.get_obs(obs_list)
        return observation + self.get_noise(obs_list, noise_dict)

    def get_obs(self, obs_list):
        observation = self.env.get_states(obs_list).to(self.device)
        return observation

    def set_actions(self, actions):
        if self.policy_cfg['disable_actions']:
            return
        if hasattr(self.env.cfg.scaling, "clip_actions"):
            actions = torch.clip(actions,
                                 -self.env.cfg.scaling.clip_actions,
                                 self.env.cfg.scaling.clip_actions)
        self.env.set_states(self.policy_cfg["actions"], actions)

    def get_timed_out(self):
        return self.env.get_states(['timed_out']).to(self.device)

    def get_terminated(self):
        return self.env.get_states(['terminated']).to(self.device)

    def get_obs_size(self, obs_list):
        # todo make unit-test to assert len(shape)==1 always
        return self.get_obs(obs_list)[0].shape[0]

    def get_action_size(self, action_list):
        return self.env.get_states(action_list)[0].shape[0]

    def get_and_log_rewards(self, reward_weights, modifier=1,
                            mask=None):
        '''
        Computes each reward on the fly, sends them to logging, and returns the
        total reward.
        reward_weights: dict with reward name, and weighting
        modifier: an additional weighting applied to all rewards
        mask: a boolean tensor of shape (num_envs), to toggle which rewards are
              computed
        '''

        if mask is None:
            mask = 1.0
        total_rewards = torch.zeros(self.env.num_envs,
                                    device=self.device, dtype=torch.float)
        for name, weight in reward_weights.items():
            reward = mask * self.get_rewards({name: weight}, modifier)
            total_rewards += reward
            self.logger.log_current_reward(name, reward)
        return total_rewards

    def get_rewards(self, reward_weights, modifier=1):
        return modifier*self.env.compute_reward(reward_weights).to(self.device)

    def log(self):
        fps = int(self.num_steps_per_env * self.env.num_envs
                  / (self.collection_time+self.learn_time))
        mean_noise_std = self.alg.actor_critic.std.mean().item()
        self.logger.add_log(self.logger.mean_rewards)
        self.logger.add_log({
            'Loss/value_function': self.mean_value_loss,
            'Loss/surrogate': self.mean_surrogate_loss,
            'Loss/learning_rate': self.alg.learning_rate,
            'Policy/mean_noise_std': mean_noise_std,
            'Perf/total_fps': fps,
            'Perf/collection_time': self.collection_time,
            'Perf/learning_time': self.learn_time,
            'Train/mean_reward': self.logger.total_mean_reward,
            'Train/mean_episode_length': self.logger.mean_episode_length,
            'Train/total_timesteps': self.tot_timesteps,
            'Train/iteration_time': self.collection_time+self.learn_time,
            'Train/time': self.tot_time,
            })
        self.logger.update_iterations(self.it, self.tot_iter,
                                      self.num_learning_iterations)

        # TODO: iterate through the config for any extra things
        # TODO: you might want to log

        if wandb.run is not None:
            self.logger.log_to_wandb()
        self.logger.print_to_terminal()

    def get_infos(self):
        return self.env.extras

    def save(self):
        path = os.path.join(self.log_dir, 'model_{}.pt'.format(self.it))
        torch.save({
            'model_state_dict': self.alg.actor_critic.state_dict(),
            'optimizer_state_dict': self.alg.optimizer.state_dict(),
            'iter': self.it},
                   path)

    def load(self, path, load_optimizer=True):
        loaded_dict = torch.load(path)
        self.alg.actor_critic.load_state_dict(loaded_dict['model_state_dict'])
        if load_optimizer:
            self.alg.optimizer.load_state_dict(
                loaded_dict['optimizer_state_dict'])
        self.it = loaded_dict['iter']

    def switch_to_eval(self):
        self.alg.actor_critic.eval()

    def get_inference_actions(self):
        obs = self.get_noisy_obs(self.policy_cfg["actor_obs"],
                                 self.policy_cfg['noise'])
        return self.alg.actor_critic.actor.act_inference(obs)

    def export(self, path):
        self.alg.actor_critic.export_policy(path)

learning/runners/on_policy_runner.py
mit-biomimetics/ORCAgym
import torch

from .base_storage import BaseStorage


class SERolloutStorage(BaseStorage):
    """ Store episodic data for supervised learning of the state-estimator.
    """

    class Transition:
        def __init__(self):
            self.observations = None
            self.SE_targets = None

        def clear(self):
            self.__init__()

    def __init__(self, num_envs, num_transitions_per_env, actor_obs_shape,
                 se_shape, device='cpu'):

        self.device = device
        self.num_transitions_per_env = num_transitions_per_env
        self.num_envs = num_envs
        self.fill_count = 0

        self.actor_obs_shape = actor_obs_shape    # raw states for actor
        self.se_shape = se_shape                  # SE prediction states

        self.observations = torch.zeros(num_transitions_per_env, num_envs,
                                        *actor_obs_shape, device=self.device)
        self.SE_targets = torch.zeros(num_transitions_per_env, num_envs,
                                      *se_shape, device=self.device)

    def add_transitions(self, transition: Transition):
        if self.fill_count >= self.num_transitions_per_env:
            raise AssertionError("Rollout buffer overflow")
        self.observations[self.fill_count].copy_(transition.observations)
        self.SE_targets[self.fill_count].copy_(transition.SE_targets)
        self.fill_count += 1

    def clear(self):
        self.fill_count = 0

    def mini_batch_generator(self, num_mini_batches, num_epochs=8):
        """ Generate mini batch for learning
        """
        batch_size = self.num_envs * self.num_transitions_per_env
        mini_batch_size = batch_size // num_mini_batches
        indices = torch.randperm(num_mini_batches*mini_batch_size,
                                 requires_grad=False, device=self.device)

        observations = self.observations.flatten(0, 1)
        SE_targets = self.SE_targets.flatten(0, 1)
        # if self.privileged_observations is not None:
        #     critic_observations = self.privileged_observations.flatten(0, 1)
        # else:
        #     critic_observations = observations

        for epoch in range(num_epochs):
            for i in range(num_mini_batches):
                start = i*mini_batch_size
                end = (i+1)*mini_batch_size
                batch_idx = indices[start:end]

                obs_batch = observations[batch_idx]
                # critic_observations_batch = critic_observations[batch_idx]
                SE_target_batch = SE_targets[batch_idx]
                yield obs_batch, SE_target_batch

learning/storage/SE_storage.py
mit-biomimetics/ORCAgym
#  Copyright 2021 ETH Zurich, NVIDIA CORPORATION
#  SPDX-License-Identifier: BSD-3-Clause

from .rollout_storage import RolloutStorage
from .SE_storage import SERolloutStorage

learning/storage/__init__.py
mit-biomimetics/ORCAgym
""" Store data for different loop
Design consideration: change the original "has-a" relationship of
LT->rollout->transition to parallel. Transition stores all the needed data for
every step and will be cleared/updated every step and passed to rollout and
longterm Rollout store data for every iteration, generate batch for learning
Longterm stores needed data from transition, and probably also generate batches
"""


class BaseStorage:
    """ This class is a skeleton for an arbitrary storage type. """

    class Transition:
        """ Transition storage class.
        i.e. store data for each STEP of ALL agents
        """
        def __init__(self):
            """ Define all the data you need to store in __init__
            """
            raise NotImplementedError

        def clear(self):
            self.__init__()

    def __init__(self, max_storage, device='cpu'):
        self.device = device
        self.max_storage = max_storage
        # fill_count keeps track of how much storage is filled with actual data
        # anything after the fill_count is stale and should be ignored
        self.fill_count = 0

    def add_transitions(self, transition: Transition):
        """ Add current transition to LT storage
        Store variables according to the __init__
        """
        self.fill_count += 1
        raise NotImplementedError

    def clear(self):
        self.fill_count = 0

    def mini_batch_generator(self):
        """ Generate mini batch for learning
        """
        raise NotImplementedError

learning/storage/base_storage.py
mit-biomimetics/ORCAgym
# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice,
# this list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
#
# 3. Neither the name of the copyright holder nor the names of its
# contributors may be used to endorse or promote products derived from
# this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF
# THE POSSIBILITY OF SUCH DAMAGE.
#
# Copyright (c) 2021 ETH Zurich, Nikita Rudin

import torch

from .base_storage import BaseStorage


class RolloutStorage(BaseStorage):
    """ A standard rollout storage, implemented for for PPO.
    """
    class Transition:
        def __init__(self):
            self.observations = None
            self.critic_observations = None
            self.actions = None
            self.rewards = None
            self.dones = None
            self.values = None
            self.actions_log_prob = None
            self.action_mean = None
            self.action_sigma = None

        def clear(self):
            self.__init__()

    def __init__(self, num_envs, num_transitions_per_env, obs_shape,
                 privileged_obs_shape, actions_shape, device='cpu'):

        self.device = device

        self.obs_shape = obs_shape
        self.privileged_obs_shape = privileged_obs_shape
        self.actions_shape = actions_shape

        # * Core
        self.observations = torch.zeros(num_transitions_per_env, num_envs,
                                        *obs_shape, device=self.device)
        if privileged_obs_shape[0] is not None:
            self.privileged_observations = torch.zeros(
                num_transitions_per_env, num_envs, *privileged_obs_shape,
                device=self.device)
        else:
            self.privileged_observations = None

        self.rewards = torch.zeros(num_transitions_per_env, num_envs,
                                   device=self.device)
        self.actions = torch.zeros(num_transitions_per_env, num_envs,
                                   *actions_shape, device=self.device)
        self.dones = torch.zeros(num_transitions_per_env, num_envs,
                                 device=self.device).byte()

        # * For PPO
        self.actions_log_prob = torch.zeros(num_transitions_per_env, num_envs,
                                            1, device=self.device)
        self.values = torch.zeros(num_transitions_per_env, num_envs,
                                  device=self.device)
        self.returns = torch.zeros(num_transitions_per_env, num_envs,
                                   device=self.device)
        self.advantages = torch.zeros(num_transitions_per_env, num_envs,
                                      device=self.device)
        self.mu = torch.zeros(num_transitions_per_env, num_envs,
                              *actions_shape, device=self.device)
        self.sigma = torch.zeros(num_transitions_per_env, num_envs,
                                 *actions_shape, device=self.device)

        self.num_transitions_per_env = num_transitions_per_env
        self.num_envs = num_envs
        # * fill count indexes the current fill of the storage.
        # * for rollout storage, this implicitly indicates the step as well.
        self.fill_count = 0

    def add_transitions(self, transition: Transition):
        if self.fill_count >= self.num_transitions_per_env:
            raise AssertionError("Rollout buffer overflow")
        self.observations[self.fill_count].copy_(transition.observations)
        if self.privileged_observations is not None:
            self.privileged_observations[self.fill_count].copy_(
                transition.critic_observations)
        self.actions[self.fill_count].copy_(transition.actions)
        self.rewards[self.fill_count].copy_(transition.rewards)
        self.dones[self.fill_count].copy_(transition.dones)
        self.values[self.fill_count].copy_(transition.values)
        self.actions_log_prob[self.fill_count].copy_(
            transition.actions_log_prob.view(-1, 1))
        self.mu[self.fill_count].copy_(transition.action_mean)
        self.sigma[self.fill_count].copy_(transition.action_sigma)
        self.fill_count += 1

    def clear(self):
        self.fill_count = 0

    def compute_returns(self, last_values, gamma, lam):
        advantage = 0
        for fill_count in reversed(range(self.num_transitions_per_env)):
            if fill_count == self.num_transitions_per_env - 1:
                next_values = last_values
            else:
                next_values = self.values[fill_count + 1]
            next_is_not_terminal = 1.0 - self.dones[fill_count].float()
            delta = (self.rewards[fill_count]
                     + next_is_not_terminal * gamma * next_values
                     - self.values[fill_count])
            advantage = delta + next_is_not_terminal * gamma * lam * advantage
            self.returns[fill_count] = advantage + self.values[fill_count]

        # * Compute and normalize the advantages
        self.advantages = self.returns - self.values
        self.advantages = ((self.advantages - self.advantages.mean())
                           / (self.advantages.std() + 1e-8))

    def get_statistics(self):
        done = self.dones
        done[-1] = 1
        flat_dones = done.permute(1, 0, 2).reshape(-1, 1)
        done_indices = torch.cat((flat_dones.new_tensor([-1],
                                                        dtype=torch.int64),
                                  flat_dones.nonzero()[:, 0]))
        trajectory_lengths = (done_indices[1:] - done_indices[:-1])
        return trajectory_lengths.float().mean(), self.rewards.mean()

    def mini_batch_generator(self, num_mini_batches, num_epochs=8):
        batch_size = self.num_envs * self.num_transitions_per_env
        mini_batch_size = batch_size // num_mini_batches
        indices = torch.randperm(num_mini_batches*mini_batch_size,
                                 requires_grad=False,
                                 device=self.device)

        observations = self.observations.flatten(0, 1)
        if self.privileged_observations is not None:
            critic_observations = self.privileged_observations.flatten(0, 1)
        else:
            critic_observations = observations

        actions = self.actions.flatten(0, 1)
        values = self.values.flatten(0, 1)
        returns = self.returns.flatten(0, 1)
        old_actions_log_prob = self.actions_log_prob.flatten(0, 1)
        advantages = self.advantages.flatten(0, 1)
        old_mu = self.mu.flatten(0, 1)
        old_sigma = self.sigma.flatten(0, 1)

        for epoch in range(num_epochs):
            for i in range(num_mini_batches):

                start = i*mini_batch_size
                end = (i+1)*mini_batch_size
                batch_idx = indices[start:end]

                obs_batch = observations[batch_idx]
                critic_observations_batch = critic_observations[batch_idx]
                actions_batch = actions[batch_idx]
                target_values_batch = values[batch_idx]
                returns_batch = returns[batch_idx]
                old_actions_log_prob_batch = old_actions_log_prob[batch_idx]
                advantages_batch = advantages[batch_idx]
                old_mu_batch = old_mu[batch_idx]
                old_sigma_batch = old_sigma[batch_idx]
                yield obs_batch, critic_observations_batch, actions_batch, \
                    target_values_batch, advantages_batch, returns_batch, \
                    old_actions_log_prob_batch, old_mu_batch, old_sigma_batch

learning/storage/rollout_storage.py
mit-biomimetics/ORCAgym
# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice,
# this list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
#
# 3. Neither the name of the copyright holder nor the names of its
# contributors may be used to endorse or promote products derived from
# this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF
# THE POSSIBILITY OF SUCH DAMAGE.
#
# Copyright (c) 2021 ETH Zurich, Nikita Rudin

from .utils import (split_and_pad_trajectories, unpad_trajectories,
                    remove_zero_weighted_rewards, set_discount_from_horizon)
from .logging import Logger

learning/utils/__init__.py
mit-biomimetics/ORCAgym
import os
import wandb
import torch
import shutil
import fnmatch
from collections import deque
from statistics import mean


class Logger:
    def __init__(self, log_dir, max_episode_length_s, device):
        self.log_dir = log_dir
        self.device = device
        self.avg_window = 100
        self.log = {}
        self.it = 0
        self.tot_iter = 0
        self.learning_iter = 0
        self.mean_episode_length = 0.
        self.total_mean_reward = 0.
        self.max_episode_length_s = max_episode_length_s

    def initialize_buffers(self, num_envs, reward_keys):
        self.current_episode_return = {
            name: torch.zeros(
                num_envs, dtype=torch.float,
                device=self.device, requires_grad=False)
            for name in reward_keys}
        self.current_episode_length = torch.zeros(
            num_envs,
            dtype=torch.float, device=self.device)
        self.avg_return_buffer = {
            name:  deque(maxlen=self.avg_window)
            for name in reward_keys}
        self.avg_length_buffer = deque(maxlen=self.avg_window)
        self.mean_rewards = {"Episode/"+name:  0. for name in reward_keys}

    def log_to_wandb(self):
        wandb.log(self.log)

    def add_log(self, log_dict):
        self.log.update(log_dict)

    def update_iterations(self, it, tot_iter, learning_iter):
        self.it = it
        self.tot_iter = tot_iter
        self.learning_iter = learning_iter

    def log_current_reward(self, name, reward):
        if name in self.current_episode_return.keys():
            self.current_episode_return[name] += reward

    def update_episode_buffer(self, dones):
        self.current_episode_length += 1
        terminated_ids = torch.where(dones == True)[0]
        for name in self.current_episode_return.keys():
            self.avg_return_buffer[name].extend(
                self.current_episode_return[name]
                [terminated_ids].cpu().numpy().tolist())
            self.current_episode_return[name][terminated_ids] = 0.

        self.avg_length_buffer.extend(
            self.current_episode_length[terminated_ids].cpu().numpy().tolist())
        self.current_episode_length[terminated_ids] = 0
        if (len(self.avg_length_buffer) > 0):
            self.calculate_reward_avg()

    def calculate_reward_avg(self):
        self.mean_episode_length = mean(self.avg_length_buffer)
        self.mean_rewards = {
            "Episode/"+name:
                mean(self.avg_return_buffer[name]) / self.max_episode_length_s
            for name in self.current_episode_return.keys()}
        self.total_mean_reward = sum(list(self.mean_rewards.values()))

    def print_to_terminal(self):
        width = 80
        pad = 35
        str = f" \033[1m Learning iteration {self.it}/{self.tot_iter} \033[0m "

        log_string = (
            f"""{'#' * width}\n"""
            f"""{str.center(width, ' ')}\n\n"""
            f"""{'Computation:':>{pad}} {
                self.log['Perf/total_fps']:.0f} steps/s (collection: {
                self.log['Perf/collection_time']:.3f}s, learning {
                self.log['Perf/learning_time']:.3f}s)\n"""
            f"""{'Value function loss:':>{pad}} {
                self.log['Loss/value_function']:.4f}\n"""
            f"""{'Surrogate loss:':>{pad}} {
                self.log['Loss/surrogate']:.4f}\n"""
            f"""{'Mean action noise std:':>{pad}} {
                self.log['Policy/mean_noise_std']:.2f}\n"""
            f"""{'Mean episode length:':>{pad}} {
                self.log['Train/mean_episode_length']:.2f}\n"""
            f"""{'Mean reward:':>{pad}} {self.total_mean_reward:.2f}\n""")
        log_string += (f"""{'-' * width}\n""")
        for key, value in self.mean_rewards.items():
            log_string += f"""{
                f'Mean episode {key[8:]}:':>{pad}} {value:.4f}\n"""
        log_string += (
            f"""{'-' * width}\n"""
            f"""{'Total timesteps:':>{pad}} {
                self.log['Train/total_timesteps']}\n"""
            f"""{'Iteration time:':>{pad}} {
                self.log['Train/iteration_time']:.2f}s\n"""
            f"""{'Total time:':>{pad}} {self.log['Train/time']:.2f}s\n"""
            f"""{'ETA:':>{pad}} {self.log['Train/time'] / (self.it + 1) * (
                self.learning_iter - self.it):.1f}s\n""")

        print(log_string)

    def configure_local_files(self, save_paths):
        # create ignore patterns dynamically based on include patterns
        def create_ignored_pattern_except(*patterns):
            def _ignore_patterns(path, names):
                keep = set(name for pattern in patterns for name in
                           fnmatch.filter(names, pattern))
                ignore = set(name for name in names if name not in keep and
                             not os.path.isdir(os.path.join(path, name)))
                return ignore
            return _ignore_patterns

        def remove_empty_folders(path, removeRoot=True):
            if not os.path.isdir(path):
                return
            # remove empty subfolders
            files = os.listdir(path)
            if len(files):
                for f in files:
                    fullpath = os.path.join(path, f)
                    if os.path.isdir(fullpath):
                        remove_empty_folders(fullpath)
            # if folder empty, delete it
            files = os.listdir(path)
            if len(files) == 0 and removeRoot:
                os.rmdir(path)

        # copy the relevant source files to the local logs for records
        save_dir = self.log_dir+'/files/'
        for save_path in save_paths:
            if save_path['type'] == 'file':
                os.makedirs(save_dir+save_path['target_dir'],
                            exist_ok=True)
                shutil.copy2(save_path['source_file'],
                             save_dir+save_path['target_dir'])
            elif save_path['type'] == 'dir':
                shutil.copytree(
                    save_path['source_dir'],
                    save_dir+save_path['target_dir'],
                    ignore=create_ignored_pattern_except(
                        *save_path['include_patterns']))
            else:
                print('WARNING: uncaught save path type:', save_path['type'])
        remove_empty_folders(save_dir)

learning/utils/logging.py
mit-biomimetics/ORCAgym
import pytest

from learning.utils import (set_discount_from_horizon,
                            remove_zero_weighted_rewards)


class TestUtils:
    def test_set_discount_from_horizon(self):
        dt = 0.1
        horizon = 1
        assert set_discount_from_horizon(dt, horizon) == 0.9
        assert set_discount_from_horizon(dt, 0.) == 0
        with pytest.raises(AssertionError):
            set_discount_from_horizon(dt, dt/2.)
        with pytest.raises(AssertionError):
            set_discount_from_horizon(dt, -1)
        with pytest.raises(AssertionError):
            set_discount_from_horizon(-dt, dt/2.)

    def test_remove_zero_weighted_rewards(self):
        reward_weights = {'a': 0, 'b': 1}
        remove_zero_weighted_rewards(reward_weights)
        assert reward_weights == {'b': 1}

learning/utils/tests/test_utils.py
mit-biomimetics/ORCAgym
# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice,
# this list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
#
# 3. Neither the name of the copyright holder nor the names of its
# contributors may be used to endorse or promote products derived from
# this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF
# THE POSSIBILITY OF SUCH DAMAGE.
#
# Copyright (c) 2021 ETH Zurich, Nikita Rudin

import torch


def split_and_pad_trajectories(tensor, dones):
    """
    Splits trajectories at done indices. Then concatenates them and pads
    with zeros up to the length og the longest trajectory.
    Returns masks corresponding to valid parts of the trajectories
    Example:
        Input: [ [a1, a2, a3, a4 | a5, a6],
                 [b1, b2 | b3, b4, b5 | b6]
                ]

        Output:[ [a1, a2, a3, a4], | [  [True, True, True, True],
                 [a5, a6, 0, 0],   |    [True, True, False, False],
                 [b1, b2, 0, 0],   |    [True, True, False, False],
                 [b3, b4, b5, 0],  |    [True, True, True, False],
                 [b6, 0, 0, 0]     |    [True, False, False, False],
                ]                  | ]

    Assumes that the inputy has the following dimension order:
    [time, number of envs, aditional dimensions]
    """
    dones = dones.clone()
    dones[-1] = 1
    # * Permute the buffers to have order
    # * (num_envs, num_transitions_per_env, ...), for correct reshaping
    flat_dones = dones.transpose(1, 0).reshape(-1, 1)

    # * Get length of trajectory by counting the number of
    # * successive not done elements
    done_indices = torch.cat(
        (flat_dones.new_tensor(
            [-1], dtype=torch.int64), flat_dones.nonzero()[:, 0]))
    trajectory_lengths = done_indices[1:] - done_indices[:-1]
    trajectory_lengths_list = trajectory_lengths.tolist()
    # * Extract the individual trajectories
    trajectories = torch.split(
        tensor.transpose(1, 0).flatten(0, 1), trajectory_lengths_list)
    padded_trajectories = torch.nn.utils.rnn.pad_sequence(trajectories)

    trajectory_masks = trajectory_lengths > torch.arange(
        0, tensor.shape[0], device=tensor.device).unsqueeze(1)
    return padded_trajectories, trajectory_masks


def unpad_trajectories(trajectories, masks):
    """ Does the inverse operation of  split_and_pad_trajectories()
    """
    # * Need to transpose before and after the masking to have proper reshaping
    return trajectories.transpose(1, 0)[masks.transpose(1, 0)].view(
        -1, trajectories.shape[0], trajectories.shape[-1]).transpose(1, 0)


def remove_zero_weighted_rewards(reward_weights):
    for name in list(reward_weights.keys()):
        if reward_weights[name] == 0:
            reward_weights.pop(name)


def set_discount_from_horizon(dt, horizon):
    """ Calculate a discount-factor from the desired discount horizon,
    and the time-step (dt).
    """
    assert (dt > 0), "Invalid time-step"
    if horizon == 0:
        discount_factor = 0
    else:
        assert (horizon >= dt), "Invalid discounting horizon"
        discrete_time_horizon = int(horizon/dt)
        discount_factor = 1 - 1/discrete_time_horizon

    return discount_factor

learning/utils/utils.py
mit-biomimetics/ORCAgym
from setuptools import find_packages
from distutils.core import setup

setup(
    name='gpuGym',
    version='1.0.1',
    author='Biomimetic Robotics Lab',
    license="BSD-3-Clause",
    packages=find_packages(),
    description='Isaac Gym environments for Legged Robots',
    install_requires=['isaacgym',
                      'matplotlib',
                      'pandas',
                      'tensorboard',
                      'setuptools==59.5.0',
                      'torch>=1.4.0',
                      'torchvision>=0.5.0',
                      'numpy>=1.16.4']
)

setup.py
mit-biomimetics/ORCAgym
import isaacgym
import torch
import sys
import pytest
from unittest.mock import patch

from gym.envs import *
from gym.utils import get_args
from gym.utils import task_registry as _task_registry

from torch.multiprocessing import set_start_method

environment_dict = {}
environment_list = []


def pytest_sessionstart(session):
    set_start_method('spawn')
    with patch.object(sys, 'argv', sys.argv[:1]):
        args = get_args()
        args.headless = True
        _task_registry.make_gym()
        env_names = _task_registry.task_classes.keys()
        env_names =[]
        for env_name in env_names:
            print(env_name)
            args.task = env_name
            env_cfg, _ = _task_registry.create_cfgs(args)
            _task_registry.update_sim_cfg(args)
            _task_registry.make_sim()
            env = _task_registry.make_env(name=env_name,
                                          env_cfg=env_cfg)
            environment_dict[env_name] = env
            environment_list.append(environment_dict[env_name])
            _task_registry._gym.destroy_sim(_task_registry._sim)


@pytest.fixture
def task_registry():
    return _task_registry


@pytest.fixture
def args(monkeypatch):
    monkeypatch.setattr(sys, 'argv', sys.argv[:1])
    return get_args()


@pytest.fixture
def env_cfg():
    return env_cfg


@pytest.fixture
def train_cfg():
    return train_cfg


@pytest.fixture
def env_dict():
    return environment_dict


@pytest.fixture
def env_list():
    return environment_list


@pytest.fixture
def mini_cheetah():
    return environment_dict['mini_cheetah']
tests/conftest.py
mit-biomimetics/ORCAgym
import os
import isaacgym
import torch

from gym.envs import __init__
from gym.utils.logging_and_saving import local_code_save_helper
from gym.utils import task_registry

from multiprocessing import Process, Manager


def _run_integration_test(return_queue, args):
    # * do initial setup
    env_cfg, train_cfg = task_registry.create_cfgs(args)

    task_registry.make_gym_and_sim()
    env = task_registry.make_env(name=args.task, env_cfg=env_cfg)
    policy_runner = task_registry.make_alg_runner(env=env, train_cfg=train_cfg)

    local_code_save_helper.log_and_save(env, env_cfg, train_cfg, policy_runner)

    # * get the default test values before learning
    default_actions = policy_runner.env.get_states(train_cfg.policy.actions)

    # * take a full iteration step
    policy_runner.learn(num_learning_iterations=args.max_iterations)

    # * get the test values after learning
    actions = policy_runner.env.get_states(train_cfg.policy.actions)

    # * return the values to the parent for assertion
    return_queue.put((actions.cpu().clone(), default_actions.cpu().clone()))
    return_queue.put(policy_runner.it)
    return_queue.put(policy_runner.log_dir)


class TestDefaultIntegration():
    def test_default_integration_settings(self, args):
        # * simulation settings for the test
        args.task = 'mini_cheetah'
        args.max_iterations = 1
        args.save_interval = 1
        args.num_envs = 16
        args.headless = True
        args.disable_wandb = True

        # * create a queue to return the values for assertion
        manager = Manager()
        return_queue = manager.Queue()

        # * spin up a child process to run the simulation iteration
        test_proc = Process(
            target=_run_integration_test, args=(return_queue, args))
        test_proc.start()
        test_proc.join()

        # * extract the values to test from the child's return queue
        actions, default_actions = return_queue.get()

        it = return_queue.get()
        log_dir = return_queue.get()

        model_0_path = os.path.join(log_dir, 'model_0.pt')
        model_1_path = os.path.join(log_dir, 'model_1.pt')

        # * assert the returned values are as expected
        assert torch.equal(actions, default_actions) is False, \
            'Actions were not updated from default'
        assert it == 1, 'Iteration update incorrect'
        assert os.path.exists(model_0_path), 'model_0.pt was not saved'
        assert os.path.exists(model_1_path), 'model_1.pt was not saved'

        # * kill the child process now that we're done with it
        test_proc.kill()

tests/integration_tests/test_default_integration.py
mit-biomimetics/ORCAgym
class TestEnvironment:
    def test_all_rewards_have_right_shape(self, env_list):

        def generate_message(env, reward_name):
            name = env.__class__.__name__
            message = f"Wrong shape for {reward_name} in {name}"
            return message

        for env in env_list:
            for item in dir(env):
                if "_reward_" in item:
                    reward_name = item.replace("_reward_", "")
                    assert len(env._eval_reward(reward_name).shape) == 1, \
                        generate_message(env, reward_name)

    def test_extras(self, env_list):

        def generate_message(env):
            message = f"Environment {env.__class__.__name__} has extras: "
            message += f"{', '.join(list(env.extras.keys()))}"
            return message

        for env in env_list:
            assert len(env.extras) == 0, generate_message(env)

tests/unit_tests/test_environments.py
mit-biomimetics/ORCAgym
from gym.utils.logging_and_saving import wandb_singleton


class TestTaskRegistry:
    def test_naming_sync(self, task_registry, env_cfg, train_cfg, args):
        env_cfg, train_cfg = task_registry.create_cfgs(args)
        wandb_helper = wandb_singleton.WandbSingleton()
        task_registry.set_log_dir_name(train_cfg)
        wandb_helper.setup_wandb(env_cfg=env_cfg,
                                 train_cfg=train_cfg,
                                 args=args)
        assert wandb_helper.experiment_name == train_cfg.log_dir.split('/')[-1]

tests/unit_tests/test_task_registry.py
mit-biomimetics/ORCAgym