Skip to content

Commit

Permalink
Merge pull request #64 from maxspahn/ft-resizing-and-rectangle-plots
Browse files Browse the repository at this point in the history
Adds feature to resize and adds plotting of rectangle.
  • Loading branch information
maxspahn authored Mar 22, 2024
2 parents 000c45e + 66fb0c7 commit 74506b5
Show file tree
Hide file tree
Showing 3 changed files with 58 additions and 33 deletions.
75 changes: 52 additions & 23 deletions planarenvs/planar_common/planar_env.py
Original file line number Diff line number Diff line change
Expand Up @@ -120,7 +120,7 @@ def check_box(

class PlanarEnv(core.Env):

SCREEN_DIM = 500
SCREEN_DIM = 30

def __init__(self, render: bool = False, dt=0.01):
self._viewer = None
Expand All @@ -138,6 +138,7 @@ def __init__(self, render: bool = False, dt=0.01):
self._obsts = []
self._goals = []
self._sensors = []
self._outside_limits = False
self.observation_space = None
self._n = None

Expand Down Expand Up @@ -230,12 +231,12 @@ def _get_ob(self):
observation['joint_state'],
self.observation_space['joint_state'],
)
self._outside_limits = True
warnings.warn(str(err))
return observation

@abstractmethod
def _terminal(self):
pass
return self._outside_limits

@abstractmethod
def continuous_dynamics(self, x, t):
Expand All @@ -262,14 +263,14 @@ def render_line(self, point_a, point_b, angle=0, color=(0, 0, 0)):
tf_matrix = np.array(((c, -s), (s, c)))
point_a = np.dot(tf_matrix, point_a)
point_b = np.dot(tf_matrix, point_b)
point_a[0] += self._offset
point_a[1] += self._offset
point_b[0] += self._offset
point_b[1] += self._offset
point_a[0] *= self._scale
point_a[1] *= self._scale
point_b[0] *= self._scale
point_b[1] *= self._scale
point_a[0] *= self.SCREEN_DIM
point_a[1] *= self.SCREEN_DIM
point_b[0] *= self.SCREEN_DIM
point_b[1] *= self.SCREEN_DIM
point_a[0] += self._offsets[0]
point_a[1] += self._offsets[1]
point_b[0] += self._offsets[0]
point_b[1] += self._offsets[1]
pygame.draw.line(
self.surf,
start_pos=point_a,
Expand All @@ -279,19 +280,29 @@ def render_line(self, point_a, point_b, angle=0, color=(0, 0, 0)):

def render_polygone(self, coordinates, color=(204, 204, 0)):
for coordinate in coordinates:
coordinate[0] = (coordinate[0] + self._offset) * self._scale
coordinate[1] = (coordinate[1] + self._offset) * self._scale
coordinate[0] = coordinate[0] * self.SCREEN_DIM + self._offsets[0]
coordinate[1] = coordinate[1] * self.SCREEN_DIM + self._offsets[1]
gfxdraw.filled_polygon(self.surf, coordinates, color)

def render_rectangle(self, position, size, color=(204, 204, 0)):
corners = [
[position[0] - size[0] / 2, position[1] - size[1] / 2],
[position[0] + size[0] / 2, position[1] - size[1] / 2],
[position[0] + size[0] / 2, position[1] + size[1] / 2],
[position[0] - size[0] / 2, position[1] + size[1] / 2],
]
self.render_polygone(corners, color=color)



def render_point(self, point_a, color=(204, 204, 0), radius=0.1):
pos_x = point_a[0] * self._scale + self._offset * self._scale
pos_y = point_a[1] * self._scale + self._offset * self._scale
pos_x = point_a[0] * self.SCREEN_DIM + self._offsets[0]
pos_y = point_a[1] * self.SCREEN_DIM + self._offsets[1]
gfxdraw.filled_circle(
self.surf,
int(pos_x),
int(pos_y),
int(radius * self._scale),
int(radius * self.SCREEN_DIM),
color,
)
def render_sub_goal(self, sub_goal):
Expand Down Expand Up @@ -320,11 +331,18 @@ def render_sub_goal(self, sub_goal):

def render_scene(self):
for obst in self._obsts:
self.render_point(
obst.position(t=self.t()),
radius=obst.radius(),
color=(0, 0,0)
)
if obst.type() == 'box':
self.render_rectangle(
obst.position(t=self.t()),
obst.size(),
color=(0, 0, 0),
)
if obst.type() == 'sphere':
self.render_point(
obst.position(t=self.t()),
radius=obst.radius(),
color=(0, 0,0)
)
for goal in self._goals:
if isinstance(goal, GoalComposition):
for sub_goal in goal.sub_goals():
Expand All @@ -334,12 +352,23 @@ def render_scene(self):

def render(self):
if self.screen is None:
scale_x = self._limits["pos"]["high"][0] - self._limits["pos"]["low"][0]
scale_y = self._limits["pos"]["high"][1] - self._limits["pos"]["low"][1]
#self._scale = self.SCREEN_DIM / (
# self._limits["pos"]["high"][0] - self._limits["pos"]["low"][0]
#)
#self._offset = self.SCREEN_DIM / (2 * self._scale)
self._scales = [self.SCREEN_DIM * 0.5 * scale_x, self.SCREEN_DIM * 0.5 * scale_y]
self._offsets = [
(0 - self._limits['pos']['low'][0])/scale_x * self.SCREEN_DIM * scale_x,
(0 - self._limits['pos']['low'][1])/scale_y * self.SCREEN_DIM * scale_y,
]
pygame.init()
pygame.display.init()
self.screen = pygame.display.set_mode(
(self.SCREEN_DIM, self.SCREEN_DIM)
(2 * self._scales[0], 2 * self._scales[1])
)
self.surf = pygame.Surface((self.SCREEN_DIM, self.SCREEN_DIM))
self.surf = pygame.Surface((2 * self._scales[0], 2 * self._scales[1]))
self.surf.fill((255, 255, 255))
self.render_specific()
self.render_scene()
Expand Down
14 changes: 5 additions & 9 deletions planarenvs/point_robot/envs/point_robot_env.py
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,7 @@ def __init__(self, n=2, dt=0.01, render=False):
},
}
self._lim_up_pos = self._limits["pos"]["high"]
self._lim_low_pos = self._limits["pos"]["low"]
self._lim_up_vel = self._limits["vel"]["high"]
self._lim_up_acc = self._limits["acc"]["high"]
self._lim_up_for = self._limits["for"]["high"]
Expand All @@ -52,11 +53,12 @@ def reset_limits(self, **kwargs):
)

self._lim_up_pos = self._limits["pos"]["high"]
self._lim_low_pos = self._limits["pos"]["low"]
self._lim_up_vel = self._limits["vel"]["high"]
self._lim_up_acc = self._limits["acc"]["high"]
self._lim_up_for = self._limits["for"]["high"]
self.observation_space.spaces["joint_state"]["position"] = spaces.Box(
low=-self._lim_up_pos, high=self._lim_up_pos, dtype=np.float64
low=self._lim_low_pos, high=self._lim_up_pos, dtype=np.float64
)
self.observation_space.spaces["joint_state"]["velocity"] = spaces.Box(
low=-self._lim_up_vel, high=self._lim_up_vel, dtype=np.float64
Expand All @@ -70,8 +72,6 @@ def _reward(self):
reward = -1.0 if not self._terminal() else 0.0
return reward

def _terminal(self):
return False

@abstractmethod
def continuous_dynamics(self, x, t):
Expand All @@ -81,10 +81,6 @@ def render_specific(self):

if self._state is None:
return None
self._scale = self.SCREEN_DIM / (
self._limits["pos"]["high"][0] - self._limits["pos"]["low"][0]
)
self._offset = self.SCREEN_DIM / (2 * self._scale)

x = self._state["joint_state"]["position"][0:2]

Expand All @@ -93,7 +89,7 @@ def render_specific(self):
[self._limits["pos"]["high"][0], 0],
)
self.render_line(
[0, self._limits["pos"]["low"][0]],
[0, self._limits["pos"]["high"][0]],
[0, self._limits["pos"]["low"][1]],
[0, self._limits["pos"]["high"][1]],
)
self.render_point(x)
2 changes: 1 addition & 1 deletion pyproject.toml
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
[tool.poetry]
name = "planarenvs"
version = "1.4.0"
version = "1.4.1"
description = "Lightweight open-ai gym environments for planar kinematic chains."
authors = ["Max Spahn <m.spahn@tudelft.nl>"]

Expand Down

0 comments on commit 74506b5

Please sign in to comment.