Locomotion/dynamics.py at master · rbapat/Locomotion · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
from isaacgym import gymapi, gymutil, gymtorch
from stable_baselines3.common.vec_env import VecEnv
from scipy.spatial.transform import Rotation as R
import torchgeometry as tgm
import numpy as np
import torch
import torch.nn as nn
import time
import gym

from isaacgym.torch_utils import *

class AgentDynamics:
    def __init__(self, ctx):
        self.ctx = ctx
        self.gym, self.sim = ctx.gym, ctx.sim

        self.root_state_desc = self.gym.acquire_actor_root_state_tensor(self.sim)
        self._root_state = gymtorch.wrap_tensor(self.root_state_desc)

        self.dof_state_desc = self.gym.acquire_dof_state_tensor(self.sim)
        self._dof_state = gymtorch.wrap_tensor(self.dof_state_desc)

        self._rb_states = gymtorch.wrap_tensor(self.gym.acquire_rigid_body_state_tensor(self.sim))
        self._net_forces = gymtorch.wrap_tensor(self.gym.acquire_net_contact_force_tensor(self.sim))
        self._joint_forces = gymtorch.wrap_tensor(self.gym.acquire_dof_force_tensor(self.sim))

        root_state_view = self._root_state.view(self.ctx.num_envs, 13)
        self.root_position = root_state_view[:, :3]
        self.root_rotation =  root_state_view[:, 3:7]
        self.root_lin_vel = root_state_view[:, 7:10]
        self.root_ang_vel = root_state_view[:, 10:]

        self.feet_idx = [5, 9, 13, 17]
        self.body_idx = [0, 1, 2, 3, 6, 7, 10, 11, 14, 15]


        dof_state_view = self._dof_state.view(self.ctx.num_envs, self.ctx.num_dof, 2)
        self.dof_pos = dof_state_view[:, :, 0]
        self.dof_vel = dof_state_view[:, :, 1]

        rb_view = self._rb_states.view(self.ctx.num_envs, self.ctx.num_rb, 13)
        self.rb_pos = rb_view[:, :, :3]
        self.rb_rot = rb_view[:, :, 3:7]
        self.rb_lin_vel = rb_view[:, :, 7:10]
        self.rb_ang_vel = rb_view[:, :, 10:]

        self.net_forces = self._net_forces.view(self.ctx.num_envs, self.ctx.num_rb, 3)
        self.joint_forces = self._joint_forces.view(self.ctx.num_envs, self.ctx.num_dof)

    def update(self):
        self.gym.refresh_actor_root_state_tensor(self.sim)
        self.gym.refresh_dof_state_tensor(self.sim)
        self.gym.refresh_rigid_body_state_tensor(self.sim)
        self.gym.refresh_net_contact_force_tensor(self.sim)
        self.gym.refresh_dof_force_tensor(self.sim)

    def get_joint_forces(self):
        return self.joint_forces

    def get_position(self):
        return self.root_position

    def get_rotation(self):
        return self.root_rotation

    def get_linear_velocity(self):
        #return self.root_lin_vel
        return quat_rotate_inverse(self.root_rotation, self.root_lin_vel)

    def get_angular_velocity(self):
        #return self.root_ang_vel
        return quat_rotate_inverse(self.root_rotation, self.root_ang_vel)

    def get_dof_position(self):
        return self.dof_pos

    def get_dof_velocity(self):
        return self.dof_vel

    def get_feet_position(self, relative = False):
        feet_pos = self.rb_pos[:, self.feet_idx, :] # (num_envs, 4, 3)
        if relative:
            for i in range(3):
                feet_pos[:, i, :] -= self.root_position

        return feet_pos

    def get_feet_velocity(self):
        feet_pos = self.rb_lin_vel[:, self.feet_idx, :] # (num_envs, 4, 3)

        for i in range(4):
            feet_pos[:, i, :] = quat_rotate_inverse(self.root_rotation, feet_pos[:, i, :])

        return feet_pos

    def get_collisions(self, idx_to_check, collision_thresh):
        contact_forces = torch.norm(self.net_forces[:, idx_to_check, :], dim = 2) # (num_envs, len(idx_to_check))
        collision_idx = contact_forces > collision_thresh # (num_envs, len(idx_to_check))
        is_collision = torch.any(collision_idx, dim = 1) # (num_envs)
        return is_collision, collision_idx

    def get_feet_collsions(self, collision_thresh = 1.0):
        return self.get_collisions(self.feet_idx, collision_thresh)[1]

    def get_body_collisions(self, collision_thresh = 1.0):
        return self.get_collisions(self.body_idx, collision_thresh)[0]