def step()

in habitat/tasks/rearrange/rearrange_sim.py [0:0]

• 48 lines of code
• 12 McCabe index (conditional complexity)

    def step(self, action: Union[str, int]) -> Observations:
        rom = self.get_rigid_object_manager()

        self._update_markers()

        if self._should_render_debug:
            self._try_acquire_context()
            for obj_handle, _ in self.ep_info["targets"].items():
                self.set_object_bb_draw(
                    False, rom.get_object_by_handle(obj_handle).object_id
                )

        add_back_viz_objs = {}
        for name, viz_id in self.viz_ids.items():
            if viz_id is None:
                continue
            rom = self.get_rigid_object_manager()
            viz_obj = rom.get_object_by_id(viz_id)
            before_pos = viz_obj.translation
            rom.remove_object_by_id(viz_id)
            add_back_viz_objs[name] = before_pos
        self.viz_ids = defaultdict(lambda: None)
        self.grasp_mgr.update()

        if self._concur_render:
            self._prev_sim_obs = self.start_async_render()

            for _ in range(self.ac_freq_ratio):
                self.internal_step(-1)
            # self.internal_step(0.008 * self.ac_freq_ratio)

            self._prev_sim_obs = self.get_sensor_observations_async_finish()
            obs = self._sensor_suite.get_observations(self._prev_sim_obs)
        else:
            for _ in range(self.ac_freq_ratio):
                self.internal_step(-1)
            # self.internal_step(0.008 * self.ac_freq_ratio)
            self._prev_sim_obs = self.get_sensor_observations()
            obs = self._sensor_suite.get_observations(self._prev_sim_obs)

        # TODO: Make debug cameras more flexible
        if "robot_third_rgb" in obs:
            self._should_render_debug = True
            self._try_acquire_context()
            for k, pos in add_back_viz_objs.items():
                self.viz_ids[k] = self.visualize_position(pos, self.viz_ids[k])

            # Also render debug information
            self._create_obj_viz(self.ep_info)

            # Always draw the target
            for obj_handle, _ in self.ep_info["targets"].items():
                self.set_object_bb_draw(
                    True, rom.get_object_by_handle(obj_handle).object_id
                )

            debug_obs = self.get_sensor_observations()
            obs["robot_third_rgb"] = debug_obs["robot_third_rgb"][:, :, :3]

        if self.habitat_config.HABITAT_SIM_V0.get(
            "ENABLE_GFX_REPLAY_SAVE", False
        ):
            self.gfx_replay_manager.save_keyframe()

        return obs