diff --git a/sim/scripts/drop_urdf.py b/sim/scripts/drop_urdf.py index aa05e87a..20eb8c04 100644 --- a/sim/scripts/drop_urdf.py +++ b/sim/scripts/drop_urdf.py @@ -14,6 +14,7 @@ from sim.env import stompy_urdf_path from sim.logging import configure_logging +from sim.stompy.joints import Stompy logger = logging.getLogger(__name__) @@ -138,6 +139,9 @@ def load_gym() -> GymParams: def run_gym(gym: GymParams) -> None: + flag = False + joints = Stompy.legs.all_joints() + while not gym.gym.query_viewer_has_closed(gym.viewer): gym.gym.simulate(gym.sim) gym.gym.fetch_results(gym.sim, True) @@ -145,6 +149,12 @@ def run_gym(gym: GymParams) -> None: gym.gym.draw_viewer(gym.viewer, gym.sim, True) gym.gym.sync_frame_time(gym.sim) + # Every second, set the target effort for each joint to the reverse. + effort = 1.0 if flag else -1.0 + flag = not flag + for joint in joints: + gym.gym.set_dof_force_target(gym.sim, joint, effort) + gym.gym.destroy_viewer(gym.viewer) gym.gym.destroy_sim(gym.sim) @@ -156,5 +166,5 @@ def main() -> None: if __name__ == "__main__": - # python -m sim.stompy + # python -m sim.scripts.drop_urdf main() diff --git a/sim/stompy/joints.py b/sim/stompy/joints.py index 6795a262..c5b6edc5 100644 --- a/sim/stompy/joints.py +++ b/sim/stompy/joints.py @@ -19,6 +19,22 @@ def children(cls) -> List["Union[Node, str]"]: if isinstance(attr, (Node, str)) ] + @classmethod + def joints(cls) -> List[str]: + return [ + attr + for attr in (getattr(cls, attr) for attr in dir(cls) if not attr.startswith("__")) + if isinstance(attr, str) + ] + + @classmethod + def all_joints(cls) -> List[str]: + leaves = cls.joints() + for child in cls.children(): + if isinstance(child, Node): + leaves.extend(child.all_joints()) + return leaves + def __str__(self) -> str: parts = [str(child) for child in self.children()] parts_str = textwrap.indent("\n" + "\n".join(parts), " ")