def _get_states()

in polymetis/python/polysim/envs/bullet_locomotor.py [0:0]

• 20 lines of code
• 4 McCabe index (conditional complexity)

    def _get_states(self):
        self.robot_state["robot_timestamp"] += self.dt
        self.robot_state["dt"] = self.dt

        state = self.robot.calc_state()
        # concatenating the joint pos and base state for controllers like ilqr
        # joint_pos should probably be renamed as generalized coordinates or something
        base_pos = self.get_current_base_pos()
        self.robot_state["joint_pos"] = np.hstack(
            [state["j_pos"], base_pos, state["base_ori_euler"]]
        ).tolist()
        base_vel = self.get_current_base_vel()
        self.robot_state["joint_vel"] = np.hstack(
            [state["j_vel"], base_vel, state["base_ang_vel"]]
        ).tolist()
        self.robot_state["joint_acc"] = state["j_eff"]  # TODO: fill in this field

        self.robot_state["base_pos"] = self.get_current_base_pos().tolist()
        self.robot_state["base_vel"] = self.get_current_base_vel().tolist()

        self.robot_state["base_ori_euler"] = self.get_current_base_ori_euler().tolist()
        self.robot_state["base_ori_quat"] = self.get_current_base_ori_quat().tolist()
        self.robot_state["base_ang_vel"] = state["base_ang_vel"].tolist()

        self.robot_state["torque_cmd"] = state["j_eff"]  # TODO: fill in this field
        return self.robot_state