-
Notifications
You must be signed in to change notification settings - Fork 1
/
solo8v2vanilla.py
135 lines (104 loc) · 4.26 KB
/
solo8v2vanilla.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
from dataclasses import dataclass
from typing import Any, Dict, List, Tuple
import numpy as np
import pkg_resources
import pybullet as p
import pybullet_data as pbd
import random
import time
import gym
from gym import error, spaces
from gym_solo.core.configs import Solo8BaseConfig
from gym_solo.core import obs
from gym_solo import solo_types
@dataclass
class Solo8VanillaConfig(Solo8BaseConfig):
urdf_path: str = 'assets/solo8v2/solo.urdf'
class Solo8VanillaEnv(gym.Env):
"""An unmodified solo8 gym environment.
Note that the model corresponds to the solo8v2.
"""
def __init__(self, use_gui: bool = False, realtime: bool = False,
config=None, **kwargs) -> None:
"""Create a solo8 env"""
self._realtime = realtime
self._config = config
self.obs_factory = obs.ObservationFactory()
self._client = p.connect(p.GUI if use_gui else p.DIRECT)
p.setAdditionalSearchPath(pbd.getDataPath())
p.setGravity(*self._config.gravity)
p.setPhysicsEngineParameter(fixedTimeStep=self._config.dt, numSubSteps=1)
self.plane = p.loadURDF('plane.urdf')
self.robot, joint_cnt = self._load_robot()
self._zero_gains = np.zeros(joint_cnt)
self.action_space = spaces.Box(-self._config.motor_torque_limit,
self._config.motor_torque_limit,
shape=(joint_cnt,))
self.reset()
def step(self, action: List[float]) -> Tuple[solo_types.obs, float, bool,
Dict[Any, Any]]:
"""The agent takes a step in the environment.
Args:
action (List[float]): The torques applied to the motors in N•m. Note
len(action) == the # of actuator
Returns:
Tuple[solo_types.obs, float, bool, Dict[Any, Any]]: A tuple of the next
observation, the reward for that step, whether or not the episode
terminates, and an info dict for misc diagnostic details.
"""
p.setJointMotorControlArray(self.robot,
np.arange(self.action_space.shape[0]),
p.TORQUE_CONTROL, forces=action,
positionGains=self._zero_gains,
velocityGains=self._zero_gains)
p.stepSimulation()
if self._realtime:
time.sleep(self._config.dt)
# TODO: Fix rewards
obs_values, obs_labels = self.obs_factory.get_obs()
return obs_values, 0.0, False, {'labels': obs_labels}
def reset(self) -> solo_types.obs:
"""Reset the state of the environment and returns an initial observation.
Returns:
solo_types.obs: The initial observation of the space.
"""
p.removeBody(self.robot)
self.robot, _ = self._load_robot()
# Let gravity do it's thing and reset the environment
for i in range(1000):
self.step(self._zero_gains)
obs_values, _ = self.obs_factory.get_obs()
return obs_values
@property
def observation_space(self):
# TODO: Write tests for this function
return self.obs_factory.get_observation_space()
def _load_robot(self) -> Tuple[int, int]:
"""Load the robot from URDF and reset the dynamics.
Returns:
Tuple[int, int]: the id of the robot object and the number of joints.
"""
robot_id = p.loadURDF(
self._config.urdf, self._config.robot_start_pos,
p.getQuaternionFromEuler(self._config.robot_start_orientation_euler),
flags=p.URDF_USE_INERTIA_FROM_FILE, useFixedBase=False)
joint_cnt = p.getNumJoints(robot_id)
p.setJointMotorControlArray(robot_id, np.arange(joint_cnt),
p.VELOCITY_CONTROL, forces=np.zeros(joint_cnt))
for joint in range(joint_cnt):
p.changeDynamics(robot_id, joint,
linearDamping=self._config.linear_damping,
angularDamping=self._config.angular_damping,
restitution=self._config.restitution,
lateralFriction=self._config.lateral_friction)
return robot_id, joint_cnt
def _close(self) -> None:
"""Soft shutdown the environment. """
p.disconnect()
def _seed(self, seed: int) -> None:
"""Set the seeds for random and numpy
Args:
seed (int): The seed to set
"""
np.random.seed(seed)
random.seed(seed)