diff --git a/.github/poetry/cpu/poetry.lock b/.github/poetry/cpu/poetry.lock index ba820f34..52fce3da 100644 --- a/.github/poetry/cpu/poetry.lock +++ b/.github/poetry/cpu/poetry.lock @@ -2892,6 +2892,25 @@ numpy = ">=1.14,<2" docs = ["matplotlib", "numpydoc (==1.1.*)", "sphinx", "sphinx-book-theme", "sphinx-remove-toctrees"] test = ["pytest", "pytest-cov"] +[[package]] +name = "sim-pusht" +version = "0.1.0" +description = "PushT environment for LeRobot" +optional = true +python-versions = "<4.0,>=3.10" +files = [ + {file = "sim_pusht-0.1.0-py3-none-any.whl", hash = "sha256:1348dcab5ea8460eff2dc97b7d62dd40f2a382df92bfdc69ff5c0224900690b0"}, + {file = "sim_pusht-0.1.0.tar.gz", hash = "sha256:d8f6a2207fd781348156206728329aa6338e9785cfc07679c5c48889b34d9b14"}, +] + +[package.dependencies] +gymnasium = ">=0.29.1,<0.30.0" +opencv-python = ">=4.9.0.80,<5.0.0.0" +pygame = ">=2.5.2,<3.0.0" +pymunk = ">=6.6.0,<7.0.0" +scikit-image = ">=0.22.0,<0.23.0" +shapely = ">=2.0.3,<3.0.0" + [[package]] name = "six" version = "1.16.0" @@ -3327,7 +3346,10 @@ files = [ docs = ["furo", "jaraco.packaging (>=9.3)", "jaraco.tidelift (>=1.4)", "rst.linker (>=1.9)", "sphinx (>=3.5)", "sphinx-lint"] testing = ["big-O", "jaraco.functools", "jaraco.itertools", "more-itertools", "pytest (>=6)", "pytest-checkdocs (>=2.4)", "pytest-cov", "pytest-enabler (>=2.2)", "pytest-ignore-flaky", "pytest-mypy", "pytest-ruff (>=0.2.1)"] +[extras] +pusht = ["sim-pusht"] + [metadata] lock-version = "2.0" python-versions = "^3.10" -content-hash = "8800bb8b24312d17b765cd2ce2799f49436171dd5fbf1bec3b07f853cfa9befd" +content-hash = "1111059e284866eabaf47b85fbb7b3785b10317feb6439f7616da3a485a17520" diff --git a/.github/poetry/cpu/pyproject.toml b/.github/poetry/cpu/pyproject.toml index e84b93c9..32d197a5 100644 --- a/.github/poetry/cpu/pyproject.toml +++ b/.github/poetry/cpu/pyproject.toml @@ -1,19 +1,23 @@ [tool.poetry] name = "lerobot" version = "0.1.0" -description = "Le robot is learning" +description = "🤗 LeRobot: State-of-the-art Machine Learning for Real-World Robotics in Pytorch" authors = [ "Rémi Cadène ", +] +maintainers = [ + "Alexander Soare ", "Simon Alibert ", ] repository = "https://github.com/Cadene/lerobot" readme = "README.md" -license = "MIT" +license = "Apache-2.0" +keywords = ["robotics, deep, reinforcement, learning, pytorch"] classifiers=[ "Development Status :: 3 - Alpha", "Intended Audience :: Developers", "Topic :: Software Development :: Build Tools", - "License :: OSI Approved :: MIT License", + "License :: OSI Approved :: Apache Software License", "Programming Language :: Python :: 3.10", ] packages = [{include = "lerobot"}] @@ -53,11 +57,19 @@ huggingface-hub = "^0.21.4" gymnasium-robotics = "^1.2.4" gymnasium = "^0.29.1" cmake = "^3.29.0.1" +sim-pusht = { version = "^0.1.0", optional = true} + + +[tool.poetry.extras] +pusht = ["sim-pusht"] [tool.poetry.group.dev.dependencies] pre-commit = "^3.6.2" debugpy = "^1.8.1" + + +[tool.poetry.group.test.dependencies] pytest = "^8.1.0" pytest-cov = "^5.0.0" diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 478be771..9c2c4f24 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -87,7 +87,7 @@ jobs: TMP: ~/tmp run: | mkdir ~/tmp - poetry install --no-interaction --no-root + poetry install --no-interaction --no-root --without dev --extras pusht - name: Save cached venv if: | diff --git a/lerobot/common/datasets/pusht.py b/lerobot/common/datasets/pusht.py index 624fb140..4ad5bd95 100644 --- a/lerobot/common/datasets/pusht.py +++ b/lerobot/common/datasets/pusht.py @@ -15,8 +15,8 @@ from torchrl.data.replay_buffers.writers import Writer from lerobot.common.datasets.abstract import AbstractDataset from lerobot.common.datasets.utils import download_and_extract_zip -from lerobot.common.envs.pusht.pusht_env import pymunk_to_shapely from lerobot.common.policies.diffusion.replay_buffer import ReplayBuffer as DiffusionPolicyReplayBuffer +from pusht.pusht_env import pymunk_to_shapely # as define in env SUCCESS_THRESHOLD = 0.95 # 95% coverage, diff --git a/lerobot/common/envs/pusht/env.py b/lerobot/common/envs/pusht/env.py index 5f7fb2c3..96811ca5 100644 --- a/lerobot/common/envs/pusht/env.py +++ b/lerobot/common/envs/pusht/env.py @@ -59,7 +59,7 @@ class PushtEnv(AbstractEnv): if not self.from_pixels: raise NotImplementedError("Use PushTEnv, instead of PushTImageEnv") - from lerobot.common.envs.pusht.pusht_image_env import PushTImageEnv + from pusht.pusht_image_env import PushTImageEnv self._env = PushTImageEnv(render_size=self.image_size) diff --git a/lerobot/common/envs/pusht/pusht_env.py b/lerobot/common/envs/pusht/pusht_env.py deleted file mode 100644 index 6ef70aec..00000000 --- a/lerobot/common/envs/pusht/pusht_env.py +++ /dev/null @@ -1,378 +0,0 @@ -import collections - -import cv2 -import gymnasium as gym -import numpy as np -import pygame -import pymunk -import pymunk.pygame_util -import shapely.geometry as sg -import skimage.transform as st -from gymnasium import spaces -from pymunk.vec2d import Vec2d - -from lerobot.common.envs.pusht.pymunk_override import DrawOptions - - -def pymunk_to_shapely(body, shapes): - geoms = [] - for shape in shapes: - if isinstance(shape, pymunk.shapes.Poly): - verts = [body.local_to_world(v) for v in shape.get_vertices()] - verts += [verts[0]] - geoms.append(sg.Polygon(verts)) - else: - raise RuntimeError(f"Unsupported shape type {type(shape)}") - geom = sg.MultiPolygon(geoms) - return geom - - -class PushTEnv(gym.Env): - metadata = {"render.modes": ["human", "rgb_array"], "video.frames_per_second": 10} - reward_range = (0.0, 1.0) - - def __init__( - self, - legacy=True, # compatibility with original - block_cog=None, - damping=None, - render_action=True, - render_size=96, - reset_to_state=None, - ): - self._seed = None - self.seed() - self.window_size = ws = 512 # The size of the PyGame window - self.render_size = render_size - self.sim_hz = 100 - # Local controller params. - self.k_p, self.k_v = 100, 20 # PD control.z - self.control_hz = self.metadata["video.frames_per_second"] - # legcay set_state for data compatibility - self.legacy = legacy - - # agent_pos, block_pos, block_angle - self.observation_space = spaces.Box( - low=np.array([0, 0, 0, 0, 0], dtype=np.float64), - high=np.array([ws, ws, ws, ws, np.pi * 2], dtype=np.float64), - shape=(5,), - dtype=np.float64, - ) - - # positional goal for agent - self.action_space = spaces.Box( - low=np.array([0, 0], dtype=np.float64), - high=np.array([ws, ws], dtype=np.float64), - shape=(2,), - dtype=np.float64, - ) - - self.block_cog = block_cog - self.damping = damping - self.render_action = render_action - - """ - If human-rendering is used, `self.window` will be a reference - to the window that we draw to. `self.clock` will be a clock that is used - to ensure that the environment is rendered at the correct framerate in - human-mode. They will remain `None` until human-mode is used for the - first time. - """ - self.window = None - self.clock = None - self.screen = None - - self.space = None - self.teleop = None - self.render_buffer = None - self.latest_action = None - self.reset_to_state = reset_to_state - - def reset(self): - seed = self._seed - self._setup() - if self.block_cog is not None: - self.block.center_of_gravity = self.block_cog - if self.damping is not None: - self.space.damping = self.damping - - # use legacy RandomState for compatibility - state = self.reset_to_state - if state is None: - rs = np.random.RandomState(seed=seed) - state = np.array( - [ - rs.randint(50, 450), - rs.randint(50, 450), - rs.randint(100, 400), - rs.randint(100, 400), - rs.randn() * 2 * np.pi - np.pi, - ] - ) - self._set_state(state) - - observation = self._get_obs() - return observation - - def step(self, action): - dt = 1.0 / self.sim_hz - self.n_contact_points = 0 - n_steps = self.sim_hz // self.control_hz - if action is not None: - self.latest_action = action - for _ in range(n_steps): - # Step PD control. - # self.agent.velocity = self.k_p * (act - self.agent.position) # P control works too. - acceleration = self.k_p * (action - self.agent.position) + self.k_v * ( - Vec2d(0, 0) - self.agent.velocity - ) - self.agent.velocity += acceleration * dt - - # Step physics. - self.space.step(dt) - - # compute reward - goal_body = self._get_goal_pose_body(self.goal_pose) - goal_geom = pymunk_to_shapely(goal_body, self.block.shapes) - block_geom = pymunk_to_shapely(self.block, self.block.shapes) - - intersection_area = goal_geom.intersection(block_geom).area - goal_area = goal_geom.area - coverage = intersection_area / goal_area - reward = np.clip(coverage / self.success_threshold, 0, 1) - done = coverage > self.success_threshold - - observation = self._get_obs() - info = self._get_info() - - return observation, reward, done, info - - def render(self, mode): - return self._render_frame(mode) - - def teleop_agent(self): - TeleopAgent = collections.namedtuple("TeleopAgent", ["act"]) - - def act(obs): - act = None - mouse_position = pymunk.pygame_util.from_pygame(Vec2d(*pygame.mouse.get_pos()), self.screen) - if self.teleop or (mouse_position - self.agent.position).length < 30: - self.teleop = True - act = mouse_position - return act - - return TeleopAgent(act) - - def _get_obs(self): - obs = np.array( - tuple(self.agent.position) + tuple(self.block.position) + (self.block.angle % (2 * np.pi),) - ) - return obs - - def _get_goal_pose_body(self, pose): - mass = 1 - inertia = pymunk.moment_for_box(mass, (50, 100)) - body = pymunk.Body(mass, inertia) - # preserving the legacy assignment order for compatibility - # the order here doesn't matter somehow, maybe because CoM is aligned with body origin - body.position = pose[:2].tolist() - body.angle = pose[2] - return body - - def _get_info(self): - n_steps = self.sim_hz // self.control_hz - n_contact_points_per_step = int(np.ceil(self.n_contact_points / n_steps)) - info = { - "pos_agent": np.array(self.agent.position), - "vel_agent": np.array(self.agent.velocity), - "block_pose": np.array(list(self.block.position) + [self.block.angle]), - "goal_pose": self.goal_pose, - "n_contacts": n_contact_points_per_step, - } - return info - - def _render_frame(self, mode): - if self.window is None and mode == "human": - pygame.init() - pygame.display.init() - self.window = pygame.display.set_mode((self.window_size, self.window_size)) - if self.clock is None and mode == "human": - self.clock = pygame.time.Clock() - - canvas = pygame.Surface((self.window_size, self.window_size)) - canvas.fill((255, 255, 255)) - self.screen = canvas - - draw_options = DrawOptions(canvas) - - # Draw goal pose. - goal_body = self._get_goal_pose_body(self.goal_pose) - for shape in self.block.shapes: - goal_points = [ - pymunk.pygame_util.to_pygame(goal_body.local_to_world(v), draw_options.surface) - for v in shape.get_vertices() - ] - goal_points += [goal_points[0]] - pygame.draw.polygon(canvas, self.goal_color, goal_points) - - # Draw agent and block. - self.space.debug_draw(draw_options) - - if mode == "human": - # The following line copies our drawings from `canvas` to the visible window - self.window.blit(canvas, canvas.get_rect()) - pygame.event.pump() - pygame.display.update() - - # the clock is already ticked during in step for "human" - - img = np.transpose(np.array(pygame.surfarray.pixels3d(canvas)), axes=(1, 0, 2)) - img = cv2.resize(img, (self.render_size, self.render_size)) - if self.render_action and self.latest_action is not None: - action = np.array(self.latest_action) - coord = (action / 512 * 96).astype(np.int32) - marker_size = int(8 / 96 * self.render_size) - thickness = int(1 / 96 * self.render_size) - cv2.drawMarker( - img, - coord, - color=(255, 0, 0), - markerType=cv2.MARKER_CROSS, - markerSize=marker_size, - thickness=thickness, - ) - return img - - def close(self): - if self.window is not None: - pygame.display.quit() - pygame.quit() - - def seed(self, seed=None): - if seed is None: - seed = np.random.randint(0, 25536) - self._seed = seed - self.np_random = np.random.default_rng(seed) - - def _handle_collision(self, arbiter, space, data): - self.n_contact_points += len(arbiter.contact_point_set.points) - - def _set_state(self, state): - if isinstance(state, np.ndarray): - state = state.tolist() - pos_agent = state[:2] - pos_block = state[2:4] - rot_block = state[4] - self.agent.position = pos_agent - # setting angle rotates with respect to center of mass - # therefore will modify the geometric position - # if not the same as CoM - # therefore should be modified first. - if self.legacy: - # for compatibility with legacy data - self.block.position = pos_block - self.block.angle = rot_block - else: - self.block.angle = rot_block - self.block.position = pos_block - - # Run physics to take effect - self.space.step(1.0 / self.sim_hz) - - def _set_state_local(self, state_local): - agent_pos_local = state_local[:2] - block_pose_local = state_local[2:] - tf_img_obj = st.AffineTransform(translation=self.goal_pose[:2], rotation=self.goal_pose[2]) - tf_obj_new = st.AffineTransform(translation=block_pose_local[:2], rotation=block_pose_local[2]) - tf_img_new = st.AffineTransform(matrix=tf_img_obj.params @ tf_obj_new.params) - agent_pos_new = tf_img_new(agent_pos_local) - new_state = np.array(list(agent_pos_new[0]) + list(tf_img_new.translation) + [tf_img_new.rotation]) - self._set_state(new_state) - return new_state - - def _setup(self): - self.space = pymunk.Space() - self.space.gravity = 0, 0 - self.space.damping = 0 - self.teleop = False - self.render_buffer = [] - - # Add walls. - walls = [ - self._add_segment((5, 506), (5, 5), 2), - self._add_segment((5, 5), (506, 5), 2), - self._add_segment((506, 5), (506, 506), 2), - self._add_segment((5, 506), (506, 506), 2), - ] - self.space.add(*walls) - - # Add agent, block, and goal zone. - self.agent = self.add_circle((256, 400), 15) - self.block = self.add_tee((256, 300), 0) - self.goal_color = pygame.Color("LightGreen") - self.goal_pose = np.array([256, 256, np.pi / 4]) # x, y, theta (in radians) - - # Add collision handling - self.collision_handeler = self.space.add_collision_handler(0, 0) - self.collision_handeler.post_solve = self._handle_collision - self.n_contact_points = 0 - - self.max_score = 50 * 100 - self.success_threshold = 0.95 # 95% coverage. - - def _add_segment(self, a, b, radius): - shape = pymunk.Segment(self.space.static_body, a, b, radius) - shape.color = pygame.Color("LightGray") # https://htmlcolorcodes.com/color-names - return shape - - def add_circle(self, position, radius): - body = pymunk.Body(body_type=pymunk.Body.KINEMATIC) - body.position = position - body.friction = 1 - shape = pymunk.Circle(body, radius) - shape.color = pygame.Color("RoyalBlue") - self.space.add(body, shape) - return body - - def add_box(self, position, height, width): - mass = 1 - inertia = pymunk.moment_for_box(mass, (height, width)) - body = pymunk.Body(mass, inertia) - body.position = position - shape = pymunk.Poly.create_box(body, (height, width)) - shape.color = pygame.Color("LightSlateGray") - self.space.add(body, shape) - return body - - def add_tee(self, position, angle, scale=30, color="LightSlateGray", mask=None): - if mask is None: - mask = pymunk.ShapeFilter.ALL_MASKS() - mass = 1 - length = 4 - vertices1 = [ - (-length * scale / 2, scale), - (length * scale / 2, scale), - (length * scale / 2, 0), - (-length * scale / 2, 0), - ] - inertia1 = pymunk.moment_for_poly(mass, vertices=vertices1) - vertices2 = [ - (-scale / 2, scale), - (-scale / 2, length * scale), - (scale / 2, length * scale), - (scale / 2, scale), - ] - inertia2 = pymunk.moment_for_poly(mass, vertices=vertices1) - body = pymunk.Body(mass, inertia1 + inertia2) - shape1 = pymunk.Poly(body, vertices1) - shape2 = pymunk.Poly(body, vertices2) - shape1.color = pygame.Color(color) - shape2.color = pygame.Color(color) - shape1.filter = pymunk.ShapeFilter(mask=mask) - shape2.filter = pymunk.ShapeFilter(mask=mask) - body.center_of_gravity = (shape1.center_of_gravity + shape2.center_of_gravity) / 2 - body.position = position - body.angle = angle - body.friction = 1 - self.space.add(body, shape1, shape2) - return body diff --git a/lerobot/common/envs/pusht/pusht_image_env.py b/lerobot/common/envs/pusht/pusht_image_env.py deleted file mode 100644 index 6547835a..00000000 --- a/lerobot/common/envs/pusht/pusht_image_env.py +++ /dev/null @@ -1,41 +0,0 @@ -import numpy as np -from gymnasium import spaces - -from lerobot.common.envs.pusht.pusht_env import PushTEnv - - -class PushTImageEnv(PushTEnv): - metadata = {"render.modes": ["rgb_array"], "video.frames_per_second": 10} - - # Note: legacy defaults to True for compatibility with original - def __init__(self, legacy=True, block_cog=None, damping=None, render_size=96): - super().__init__( - legacy=legacy, block_cog=block_cog, damping=damping, render_size=render_size, render_action=False - ) - ws = self.window_size - self.observation_space = spaces.Dict( - { - "image": spaces.Box(low=0, high=1, shape=(3, render_size, render_size), dtype=np.float32), - "agent_pos": spaces.Box(low=0, high=ws, shape=(2,), dtype=np.float32), - } - ) - self.render_cache = None - - def _get_obs(self): - img = super()._render_frame(mode="rgb_array") - - agent_pos = np.array(self.agent.position) - img_obs = np.moveaxis(img, -1, 0) - obs = {"image": img_obs, "agent_pos": agent_pos} - - self.render_cache = img - - return obs - - def render(self, mode): - assert mode == "rgb_array" - - if self.render_cache is None: - self._get_obs() - - return self.render_cache diff --git a/lerobot/common/envs/pusht/pymunk_override.py b/lerobot/common/envs/pusht/pymunk_override.py deleted file mode 100644 index 7ad76237..00000000 --- a/lerobot/common/envs/pusht/pymunk_override.py +++ /dev/null @@ -1,244 +0,0 @@ -# ---------------------------------------------------------------------------- -# pymunk -# Copyright (c) 2007-2016 Victor Blomqvist -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in -# all copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. -# ---------------------------------------------------------------------------- - -"""This submodule contains helper functions to help with quick prototyping -using pymunk together with pygame. - -Intended to help with debugging and prototyping, not for actual production use -in a full application. The methods contained in this module is opinionated -about your coordinate system and not in any way optimized. -""" - -__docformat__ = "reStructuredText" - -__all__ = [ - "DrawOptions", - "get_mouse_pos", - "to_pygame", - "from_pygame", - # "lighten", - "positive_y_is_up", -] - -from typing import Sequence, Tuple - -import numpy as np -import pygame -import pymunk -from pymunk.space_debug_draw_options import SpaceDebugColor -from pymunk.vec2d import Vec2d - -positive_y_is_up: bool = False -"""Make increasing values of y point upwards. - -When True:: - - y - ^ - | . (3, 3) - | - | . (2, 2) - | - +------ > x - -When False:: - - +------ > x - | - | . (2, 2) - | - | . (3, 3) - v - y - -""" - - -class DrawOptions(pymunk.SpaceDebugDrawOptions): - def __init__(self, surface: pygame.Surface) -> None: - """Draw a pymunk.Space on a pygame.Surface object. - - Typical usage:: - - >>> import pymunk - >>> surface = pygame.Surface((10,10)) - >>> space = pymunk.Space() - >>> options = pymunk.pygame_util.DrawOptions(surface) - >>> space.debug_draw(options) - - You can control the color of a shape by setting shape.color to the color - you want it drawn in:: - - >>> c = pymunk.Circle(None, 10) - >>> c.color = pygame.Color("pink") - - See pygame_util.demo.py for a full example - - Since pygame uses a coordinate system where y points down (in contrast - to many other cases), you either have to make the physics simulation - with Pymunk also behave in that way, or flip everything when you draw. - - The easiest is probably to just make the simulation behave the same - way as Pygame does. In that way all coordinates used are in the same - orientation and easy to reason about:: - - >>> space = pymunk.Space() - >>> space.gravity = (0, -1000) - >>> body = pymunk.Body() - >>> body.position = (0, 0) # will be positioned in the top left corner - >>> space.debug_draw(options) - - To flip the drawing its possible to set the module property - :py:data:`positive_y_is_up` to True. Then the pygame drawing will flip - the simulation upside down before drawing:: - - >>> positive_y_is_up = True - >>> body = pymunk.Body() - >>> body.position = (0, 0) - >>> # Body will be position in bottom left corner - - :Parameters: - surface : pygame.Surface - Surface that the objects will be drawn on - """ - self.surface = surface - super().__init__() - - def draw_circle( - self, - pos: Vec2d, - angle: float, - radius: float, - outline_color: SpaceDebugColor, - fill_color: SpaceDebugColor, - ) -> None: - p = to_pygame(pos, self.surface) - - pygame.draw.circle(self.surface, fill_color.as_int(), p, round(radius), 0) - pygame.draw.circle(self.surface, light_color(fill_color).as_int(), p, round(radius - 4), 0) - - # circle_edge = pos + Vec2d(radius, 0).rotated(angle) - # p2 = to_pygame(circle_edge, self.surface) - # line_r = 2 if radius > 20 else 1 - # pygame.draw.lines(self.surface, outline_color.as_int(), False, [p, p2], line_r) - - def draw_segment(self, a: Vec2d, b: Vec2d, color: SpaceDebugColor) -> None: - p1 = to_pygame(a, self.surface) - p2 = to_pygame(b, self.surface) - - pygame.draw.aalines(self.surface, color.as_int(), False, [p1, p2]) - - def draw_fat_segment( - self, - a: Tuple[float, float], - b: Tuple[float, float], - radius: float, - outline_color: SpaceDebugColor, - fill_color: SpaceDebugColor, - ) -> None: - p1 = to_pygame(a, self.surface) - p2 = to_pygame(b, self.surface) - - r = round(max(1, radius * 2)) - pygame.draw.lines(self.surface, fill_color.as_int(), False, [p1, p2], r) - if r > 2: - orthog = [abs(p2[1] - p1[1]), abs(p2[0] - p1[0])] - if orthog[0] == 0 and orthog[1] == 0: - return - scale = radius / (orthog[0] * orthog[0] + orthog[1] * orthog[1]) ** 0.5 - orthog[0] = round(orthog[0] * scale) - orthog[1] = round(orthog[1] * scale) - points = [ - (p1[0] - orthog[0], p1[1] - orthog[1]), - (p1[0] + orthog[0], p1[1] + orthog[1]), - (p2[0] + orthog[0], p2[1] + orthog[1]), - (p2[0] - orthog[0], p2[1] - orthog[1]), - ] - pygame.draw.polygon(self.surface, fill_color.as_int(), points) - pygame.draw.circle( - self.surface, - fill_color.as_int(), - (round(p1[0]), round(p1[1])), - round(radius), - ) - pygame.draw.circle( - self.surface, - fill_color.as_int(), - (round(p2[0]), round(p2[1])), - round(radius), - ) - - def draw_polygon( - self, - verts: Sequence[Tuple[float, float]], - radius: float, - outline_color: SpaceDebugColor, - fill_color: SpaceDebugColor, - ) -> None: - ps = [to_pygame(v, self.surface) for v in verts] - ps += [ps[0]] - - radius = 2 - pygame.draw.polygon(self.surface, light_color(fill_color).as_int(), ps) - - if radius > 0: - for i in range(len(verts)): - a = verts[i] - b = verts[(i + 1) % len(verts)] - self.draw_fat_segment(a, b, radius, fill_color, fill_color) - - def draw_dot(self, size: float, pos: Tuple[float, float], color: SpaceDebugColor) -> None: - p = to_pygame(pos, self.surface) - pygame.draw.circle(self.surface, color.as_int(), p, round(size), 0) - - -def get_mouse_pos(surface: pygame.Surface) -> Tuple[int, int]: - """Get position of the mouse pointer in pymunk coordinates.""" - p = pygame.mouse.get_pos() - return from_pygame(p, surface) - - -def to_pygame(p: Tuple[float, float], surface: pygame.Surface) -> Tuple[int, int]: - """Convenience method to convert pymunk coordinates to pygame surface - local coordinates. - - Note that in case positive_y_is_up is False, this function won't actually do - anything except converting the point to integers. - """ - if positive_y_is_up: - return round(p[0]), surface.get_height() - round(p[1]) - else: - return round(p[0]), round(p[1]) - - -def from_pygame(p: Tuple[float, float], surface: pygame.Surface) -> Tuple[int, int]: - """Convenience method to convert pygame surface local coordinates to - pymunk coordinates - """ - return to_pygame(p, surface) - - -def light_color(color: SpaceDebugColor): - color = np.minimum(1.2 * np.float32([color.r, color.g, color.b, color.a]), np.float32([255])) - color = SpaceDebugColor(r=color[0], g=color[1], b=color[2], a=color[3]) - return color diff --git a/poetry.lock b/poetry.lock index 72397001..b05cdc96 100644 --- a/poetry.lock +++ b/poetry.lock @@ -3096,6 +3096,25 @@ numpy = ">=1.14,<2" docs = ["matplotlib", "numpydoc (==1.1.*)", "sphinx", "sphinx-book-theme", "sphinx-remove-toctrees"] test = ["pytest", "pytest-cov"] +[[package]] +name = "sim-pusht" +version = "0.1.0" +description = "PushT environment for LeRobot" +optional = true +python-versions = "<4.0,>=3.10" +files = [ + {file = "sim_pusht-0.1.0-py3-none-any.whl", hash = "sha256:1348dcab5ea8460eff2dc97b7d62dd40f2a382df92bfdc69ff5c0224900690b0"}, + {file = "sim_pusht-0.1.0.tar.gz", hash = "sha256:d8f6a2207fd781348156206728329aa6338e9785cfc07679c5c48889b34d9b14"}, +] + +[package.dependencies] +gymnasium = ">=0.29.1,<0.30.0" +opencv-python = ">=4.9.0.80,<5.0.0.0" +pygame = ">=2.5.2,<3.0.0" +pymunk = ">=6.6.0,<7.0.0" +scikit-image = ">=0.22.0,<0.23.0" +shapely = ">=2.0.3,<3.0.0" + [[package]] name = "six" version = "1.16.0" @@ -3586,7 +3605,10 @@ files = [ docs = ["furo", "jaraco.packaging (>=9.3)", "jaraco.tidelift (>=1.4)", "rst.linker (>=1.9)", "sphinx (>=3.5)", "sphinx-lint"] testing = ["big-O", "jaraco.functools", "jaraco.itertools", "more-itertools", "pytest (>=6)", "pytest-checkdocs (>=2.4)", "pytest-cov", "pytest-enabler (>=2.2)", "pytest-ignore-flaky", "pytest-mypy", "pytest-ruff (>=0.2.1)"] +[extras] +pusht = ["sim-pusht"] + [metadata] lock-version = "2.0" python-versions = "^3.10" -content-hash = "174c7d42f8039eedd2c447a4e6cae5169782cbd94346b5606572a0010194ca05" +content-hash = "91dcb5cc851e0a348bfa36801e6cd260391a851623f0949575a5a19ca014d75c" diff --git a/pyproject.toml b/pyproject.toml index 972c1b61..cfda20c7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,19 +1,23 @@ [tool.poetry] name = "lerobot" version = "0.1.0" -description = "Le robot is learning" +description = "🤗 LeRobot: State-of-the-art Machine Learning for Real-World Robotics in Pytorch" authors = [ "Rémi Cadène ", +] +maintainers = [ + "Alexander Soare ", "Simon Alibert ", ] repository = "https://github.com/Cadene/lerobot" readme = "README.md" -license = "MIT" +license = "Apache-2.0" +keywords = ["robotics, deep, reinforcement, learning, pytorch"] classifiers=[ "Development Status :: 3 - Alpha", "Intended Audience :: Developers", "Topic :: Software Development :: Build Tools", - "License :: OSI Approved :: MIT License", + "License :: OSI Approved :: Apache Software License", "Programming Language :: Python :: 3.10", ] packages = [{include = "lerobot"}] @@ -52,11 +56,19 @@ robomimic = "0.2.0" gymnasium-robotics = "^1.2.4" gymnasium = "^0.29.1" cmake = "^3.29.0.1" +sim-pusht = { version = "^0.1.0", optional = true} + + +[tool.poetry.extras] +pusht = ["sim-pusht"] [tool.poetry.group.dev.dependencies] pre-commit = "^3.6.2" debugpy = "^1.8.1" + + +[tool.poetry.group.test.dependencies] pytest = "^8.1.0" pytest-cov = "^5.0.0" @@ -93,6 +105,9 @@ exclude = [ select = ["E4", "E7", "E9", "F", "I", "N", "B", "C4", "SIM"] ignore-init-module-imports = true +[tool.ruff.lint.isort] +known-first-party = ["pusht"] + [tool.poetry-dynamic-versioning] enable = true