def _state_body()

in env_humanoid_base.py [0:0]

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

    def _state_body(self, 
                    agent, 
                    T_ref=None, 
                    include_com=True, 
                    include_p=True, 
                    include_Q=True, 
                    include_v=True, 
                    include_w=True, 
                    return_stacked=True):
        if T_ref is None: 
            T_ref = agent.get_facing_transform(self.get_ground_height())

        R_ref, p_ref = conversions.T2Rp(T_ref)
        R_ref_inv = R_ref.transpose()

        link_states = []
        link_states.append(agent.get_root_state())
        ps, Qs, vs, ws = agent.get_link_states()
        for j in agent._joint_indices:
            link_states.append((ps[j], Qs[j], vs[j], ws[j]))

        state = []
        for i, s in enumerate(link_states):
            p, Q, v, w = s[0], s[1], s[2], s[3]
            if include_p:
                p_rel = np.dot(R_ref_inv, p - p_ref)
                state.append(p_rel) # relative position w.r.t. the reference frame
            if include_Q:
                Q_rel = conversions.R2Q(np.dot(R_ref_inv, conversions.Q2R(Q)))
                Q_rel = quaternion.Q_op(Q_rel, op=["normalize", "halfspace"])
                state.append(Q_rel) # relative rotation w.r.t. the reference frame
            if include_v:
                v_rel = np.dot(R_ref_inv, v)
                state.append(v_rel) # relative linear vel w.r.t. the reference frame
            if include_w:
                w_rel = np.dot(R_ref_inv, w)
                state.append(w_rel) # relative angular vel w.r.t. the reference frame
            if include_com:
                if i==0:
                    p_com = agent._link_masses[i] * p
                    v_com = agent._link_masses[i] * v
                else:
                    p_com += agent._link_masses[i] * p
                    v_com += agent._link_masses[i] * v

        if include_com:
            p_com /= agent._link_total_mass
            v_com /= agent._link_total_mass
            state.append(np.dot(R_ref_inv, p_com - p_ref))
            state.append(np.dot(R_ref_inv, v_com))
        
        if return_stacked:
            return np.hstack(state)
        else:
            return state