diff --git a/habitat/sims/pyrobot/__init__.py b/habitat/sims/pyrobot/__init__.py index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..91bfd32df7548b22c7dc8f02d07e0d16c4c84515 100644 --- a/habitat/sims/pyrobot/__init__.py +++ b/habitat/sims/pyrobot/__init__.py @@ -0,0 +1,2 @@ +# TODO(akadian): come up with better module structure +from habitat.sims.pyrobot.pyrobot import PyRobot \ No newline at end of file diff --git a/habitat/sims/pyrobot/pyrobot.py b/habitat/sims/pyrobot/pyrobot.py index 26018ec0e1980caef3b526e8f030e94f40722536..6d48cf57d3e2dff562b73e377589424049d1f190 100644 --- a/habitat/sims/pyrobot/pyrobot.py +++ b/habitat/sims/pyrobot/pyrobot.py @@ -4,18 +4,34 @@ # This source code is licensed under the MIT license found in the # LICENSE file in the root directory of this source tree. - +from habitat.core.registry import registry from habitat.core.simulator import ( Config, SensorSuite, Simulator, Space ) +import pyrobot +import numpy as np + + +# TODO(akadian): make MAX_DEPTH a parameter of the class +MAX_DEPTH = 5.0 +@registry.register_simulator(name="PyRobot-v0") class PyRobot(Simulator): def __init__(self, config: Config) -> None: - raise NotImplementedError + config_pyrobot = { + "base_controller": "proportional" + } + # self._forward_step_change = [0.25, 0, 0] + + self._angle_step = (10 / 180) * np.pi + # self._left_step_change = [0, 0, self._angle_step] + # self._right_step_change = [0, 0, -self._angle_step] + + self._robot = pyrobot.Robot("locobot", base_config=config_pyrobot) @property def sensor_suite(self) -> SensorSuite: @@ -25,8 +41,76 @@ class PyRobot(Simulator): def action_space(self) -> Space: return self._action_space + def reset(self): + self._robot.camera.reset() + rgb_observation = self.rgb_observation() + depth_observation = self.depth_observation() + + return { + "rgb": rgb_observation, + "depth": depth_observation, + } + + def _forward(self): + target_position = [0.25, 0, 0] + # TODO(akadian): see if there is a command to check if the step was successful + self._robot.base.go_to_relative( + target_position, + smooth=False, + close_loop=True + ) + + def _left(self): + target_position = [0, 0, self._angle_step] + self._robot.base.go_to_relative( + target_position, + smooth=False, + close_loop=True + ) + + def _right(self): + target_position = [0, 0, -self._angle_step] + self._robot.base.go_to_relative( + target_position, + smooth=False, + close_loop=True + ) + def step(self, action): - raise NotImplementedError + if action == 0: + # forward + self._forward() + elif action == 1: + # left + self._left() + elif action == 2: + # right + self._right() + else: + raise ValueError + + rgb_observation = self.rgb_observation() + depth_observation = self.depth_observation() + + return { + "rgb": rgb_observation, + "depth": depth_observation, + } + + # TODO(akadian): expose a method to just get the observations, + # this will be useful for the case if the state of robot + # changes due to external factors. + + def rgb_observation(self): + rgb = self._robot.camera.get_rgb() + return rgb + + def depth_observation(self): + depth = self._robot.camera.get_depth() + depth = depth / 1000 # convert from mm to m + depth[depth > MAX_DEPTH] = MAX_DEPTH + depth = depth / MAX_DEPTH # scale to [0, 1] + return depth def close(self): raise NotImplementedError diff --git a/habitat/sims/registration.py b/habitat/sims/registration.py index d756906d51fc0d68f3f5d8f083f418cce9f631e9..ae7744c85304585a795defecb2c02ebae770237e 100644 --- a/habitat/sims/registration.py +++ b/habitat/sims/registration.py @@ -8,6 +8,7 @@ from habitat.core.logging import logger from habitat.core.registry import registry from habitat.core.simulator import Simulator from habitat.sims.habitat_simulator import _try_register_habitat_sim +from habitat.sims.pyrobot.pyrobot import PyRobot def make_sim(id_sim, **kwargs): diff --git a/simple_pyrobot.py b/simple_pyrobot.py new file mode 100644 index 0000000000000000000000000000000000000000..ed3a35e8cabf9ec45c75a4814c933e436c0c04f5 --- /dev/null +++ b/simple_pyrobot.py @@ -0,0 +1,19 @@ +from habitat.sims import make_sim +from PIL import Image + + +def main(): + reality = make_sim( + id_sim="PyRobot-v0", config=None + ) + print("Reality created") + + observations = reality.reset() + observation = reality.step(0) + + print("RGB:", observations["rgb"].shape) + print("Depth:", observations["depth"].shape) + + +if __name__ == "__main__": + main()