This repository has been archived by the owner on Oct 8, 2022. It is now read-only.
-
Notifications
You must be signed in to change notification settings - Fork 1
/
scene.py
139 lines (103 loc) · 4.07 KB
/
scene.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
136
137
138
139
from pybullet_utils import bullet_client
import pybullet_data
import pybullet
import dataclasses
from typing import Sequence, Optional, Union
from .color import Color
@dataclasses.dataclass(frozen=True)
class DebugBody:
body_id: int
def remove_from_scene(self, scene):
# TODO: annotate scene with Scene type
scene.client.removeBody(self.body_id)
@dataclasses.dataclass(frozen=True)
class DebugSphere(DebugBody):
pass
@dataclasses.dataclass(frozen=True)
class DebugItem:
item_id: int
def remove_from_scene(self, scene):
# TODO: annotate scene with Scene type
scene.client.removeUserDebugItem(self.item_id)
@dataclasses.dataclass(frozen=True)
class DebugLine(DebugItem):
pass
@dataclasses.dataclass(frozen=True)
class DebugText(DebugItem):
pass
@dataclasses.dataclass
class Scene:
timestep: float
frame_skip: int
gravity: float = 9.8
client: bullet_client.BulletClient = None
dt: float = dataclasses.field(init=False)
ts: float = dataclasses.field(init=False)
def __post_init__(self):
self.dt = self.timestep * self.frame_skip
self.connect(self.client)
self.episode_restart()
def connect(self, client):
if client:
self.client = client
else:
self.client = bullet_client.BulletClient()
def clean_everything(self):
self.client.resetSimulation()
self.configure_simulation()
def configure_simulation(self):
self.client.setAdditionalSearchPath(pybullet_data.getDataPath())
self.client.setGravity(0, 0, -self.gravity)
self.client.setPhysicsEngineParameter(fixedTimeStep=self.dt, numSubSteps=self.frame_skip)
def load_plane(self):
self.plane_id = self.client.loadURDF("plane.urdf")
def episode_restart(self):
self.clean_everything()
self.ts = 0
self.load_plane()
def step(self):
self.client.stepSimulation()
self.ts += self.dt
def draw_line(self, from_pos: Sequence[float], to_pos: Sequence[float], color: Optional[Color] = None, width: Optional[float] = None, replace: Optional[DebugItem] = None) -> DebugLine:
args = {
'lineFromXYZ': from_pos,
'lineToXYZ': to_pos
}
if color is not None:
args['lineColorRGB'] = color.as_rgb()
if replace is not None:
args['replaceItemUniqueId'] = replace.item_id
if width is not None:
args['lineWidth'] = width
item_id = self.client.addUserDebugLine(**args)
return DebugLine(item_id)
def draw_text(self, text: str, pos: Sequence[float], orientation: Optional[Sequence[float]] = None, color: Optional[Color] = None, size: Optional[float] = None, replace: Optional[DebugItem] = None) -> DebugText:
args = {
'text': text,
'textPosition': pos,
}
if color is not None:
args['textColorRGB'] = color.as_rgb()
if orientation is not None:
args['textOrientation'] = orientation
if replace is not None:
args['replaceItemUniqueId'] = replace.item_id
if size is not None:
args['textSize'] = size
item_id = self.client.addUserDebugText(**args)
return DebugText(item_id)
def draw_sphere(self, pos: Sequence[float], radius: Optional[float] = None, color: Optional[Color] = None) -> DebugSphere:
args = {
'shapeType': pybullet.GEOM_SPHERE
}
if color is not None:
args['rgbaColor'] = color.as_rgba()
if radius is not None:
args['radius'] = radius
visual_id = self.client.createVisualShape(**args)
body_id = self.client.createMultiBody(baseVisualShapeIndex=visual_id, basePosition=pos)
return DebugSphere(body_id)
def move_debug_body(self, body: DebugBody, position: Sequence[float], orientation: Sequence[float]):
self.client.resetBasePositionAndOrientation(body.body_id, position, orientation)
def remove_debug_object(self, o: Union[DebugItem, DebugBody]):
o.remove_from_scene(self)